From ffc695f7b8d0a09246d6dca6651381be23622841 Mon Sep 17 00:00:00 2001 From: Marc 'risson' Schmitt Date: Thu, 27 Mar 2025 15:42:26 +0100 Subject: [PATCH] wip Signed-off-by: Marc 'risson' Schmitt --- authentik/admin/apps.py | 9 +++++ authentik/blueprints/apps.py | 33 +++++++++++++++++++ authentik/tasks/broker.py | 4 +-- authentik/tasks/models.py | 4 +-- authentik/tasks/schedules/api.py | 7 ++-- .../migrations/0002_remove_schedule_name.py | 17 ++++++++++ authentik/tasks/schedules/models.py | 10 +++--- authentik/tasks/schedules/scheduler.py | 11 ++++--- 8 files changed, 77 insertions(+), 18 deletions(-) create mode 100644 authentik/tasks/schedules/migrations/0002_remove_schedule_name.py diff --git a/authentik/admin/apps.py b/authentik/admin/apps.py index def7d51ede..bc9748a6f9 100644 --- a/authentik/admin/apps.py +++ b/authentik/admin/apps.py @@ -3,6 +3,7 @@ from prometheus_client import Info from authentik.blueprints.apps import ManagedAppConfig +from authentik.lib.utils.time import fqdn_rand PROM_INFO = Info("authentik_version", "Currently running authentik version") @@ -14,3 +15,11 @@ class AuthentikAdminConfig(ManagedAppConfig): label = "authentik_admin" verbose_name = "authentik Admin" default = True + + def get_tenant_schedules(self): + return [ + { + "actor_name": "authentik.admin.tasks.update_latest_version", + "crontab": f"{fqdn_rand('admin_latest_version')} * * * *", + }, + ] diff --git a/authentik/blueprints/apps.py b/authentik/blueprints/apps.py index eeef482764..30080b3484 100644 --- a/authentik/blueprints/apps.py +++ b/authentik/blueprints/apps.py @@ -1,8 +1,10 @@ """authentik Blueprints app""" +import pickle # nosec from collections.abc import Callable from importlib import import_module from inspect import ismethod +from typing import Any from django.apps import AppConfig from django.db import DatabaseError, InternalError, ProgrammingError @@ -80,6 +82,35 @@ class ManagedAppConfig(AppConfig): func._authentik_managed_reconcile = ManagedAppConfig.RECONCILE_GLOBAL_CATEGORY return func + def get_tenant_schedules(self) -> list[dict[str, Any]]: + return [] + + def get_global_schedules(self) -> list[dict[str, Any]]: + return [] + + def _reconcile_schedules(self, schedules: list[dict[str, Any]]): + from authentik.tasks.schedules.models import Schedule + + for schedule in schedules: + query = { + "uid": schedule.get("uid", schedule["actor_name"]), + } + defaults = { + **query, + "actor_name": schedule["actor_name"], + "args": pickle.dumps(schedule.get("args", ())), + "kwargs": pickle.dumps(schedule.get("kwargs", {})), + } + create_defaults = { + **defaults, + "crontab": schedule["crontab"], + } + Schedule.objects.update_or_create( + **query, + defaults=defaults, + create_defaults=create_defaults, + ) + def _reconcile_tenant(self) -> None: """reconcile ourselves for tenanted methods""" from authentik.tenants.models import Tenant @@ -92,6 +123,7 @@ class ManagedAppConfig(AppConfig): for tenant in tenants: with tenant: self._reconcile(self.RECONCILE_TENANT_CATEGORY) + self._reconcile_schedules(self.get_tenant_schedules()) def _reconcile_global(self) -> None: """ @@ -102,6 +134,7 @@ class ManagedAppConfig(AppConfig): with schema_context(get_public_schema_name()): self._reconcile(self.RECONCILE_GLOBAL_CATEGORY) + self._reconcile_schedules(self.get_global_schedules()) class AuthentikBlueprintsConfig(ManagedAppConfig): diff --git a/authentik/tasks/broker.py b/authentik/tasks/broker.py index f7a9e71495..5badf3c7d1 100644 --- a/authentik/tasks/broker.py +++ b/authentik/tasks/broker.py @@ -30,9 +30,9 @@ from structlog.stdlib import get_logger from authentik.tasks.models import CHANNEL_PREFIX, ChannelIdentifier, Task, TaskState from authentik.tasks.results import PostgresBackend +from authentik.tasks.schedules.scheduler import Scheduler from authentik.tenants.models import Tenant from authentik.tenants.utils import get_current_tenant -from authentik.tasks.schedules.scheduler import Scheduler LOGGER = get_logger() @@ -176,7 +176,7 @@ class PostgresBroker(Broker): **query, **defaults, } - self.query_set.update_or_create( + obj, created = self.query_set.update_or_create( **query, defaults=defaults, create_defaults=create_defaults, diff --git a/authentik/tasks/models.py b/authentik/tasks/models.py index 35b1c6fe2f..5539b06b20 100644 --- a/authentik/tasks/models.py +++ b/authentik/tasks/models.py @@ -107,7 +107,7 @@ class Task(SerializerModel): self.messages = list(messages) for idx, msg in enumerate(self.messages): if not isinstance(msg, LogEvent): - self.messages[idx] = LogEvent(msg, logger=self.__name__, log_level="info") + self.messages[idx] = LogEvent(msg, logger=str(self), log_level="info") self.messages = sanitize_item(self.messages) def set_error(self, exception: Exception, *messages: LogEvent | str): @@ -115,6 +115,6 @@ class Task(SerializerModel): self.status = TaskStatus.ERROR self.messages = list(messages) self.messages.extend( - [LogEvent(exception_to_string(exception), logger=self.__name__, log_level="error")] + [LogEvent(exception_to_string(exception), logger=str(self), log_level="error")] ) self.messages = sanitize_item(self.messages) diff --git a/authentik/tasks/schedules/api.py b/authentik/tasks/schedules/api.py index 1936cdc5ae..614b369226 100644 --- a/authentik/tasks/schedules/api.py +++ b/authentik/tasks/schedules/api.py @@ -14,7 +14,8 @@ class ScheduleSerializer(ModelSerializer): model = Schedule fields = [ "id", - "name", + "uid", + "actor_name", "crontab", "next_run", ] @@ -30,6 +31,6 @@ class ScheduleViewSet( serializer_class = ScheduleSerializer search_fields = ( "id", - "name", + "uid", ) - ordering = ("id",) + ordering = ("-next_run", "uid") diff --git a/authentik/tasks/schedules/migrations/0002_remove_schedule_name.py b/authentik/tasks/schedules/migrations/0002_remove_schedule_name.py new file mode 100644 index 0000000000..b3aa0aa63a --- /dev/null +++ b/authentik/tasks/schedules/migrations/0002_remove_schedule_name.py @@ -0,0 +1,17 @@ +# Generated by Django 5.0.13 on 2025-03-27 14:00 + +from django.db import migrations + + +class Migration(migrations.Migration): + + dependencies = [ + ("authentik_tasks_schedules", "0001_initial"), + ] + + operations = [ + migrations.RemoveField( + model_name="schedule", + name="name", + ), + ] diff --git a/authentik/tasks/schedules/models.py b/authentik/tasks/schedules/models.py index f5a498ea57..67b7653786 100644 --- a/authentik/tasks/schedules/models.py +++ b/authentik/tasks/schedules/models.py @@ -21,17 +21,15 @@ def validate_crontab(value): class Schedule(SerializerModel): id = models.UUIDField(primary_key=True, default=uuid4, editable=False) - uid = models.TextField(unique=True, editable=False) - - name = models.TextField(editable=False, help_text=_("Schedule display name")) + uid = models.TextField(unique=True, editable=False, help_text=_("Unique schedule identifier")) actor_name = models.TextField(editable=False, help_text=_("Dramatiq actor to call")) args = models.BinaryField(editable=False, help_text=_("Args to send to the actor")) kwargs = models.BinaryField(editable=False, help_text=_("Kwargs to send to the actor")) - crontab = models.TextField(validators=[validate_crontab]) + crontab = models.TextField(validators=[validate_crontab], help_text=_("When to schedule tasks")) - next_run = models.DateTimeField(auto_now_add=True) + next_run = models.DateTimeField(auto_now_add=True, editable=False) class Meta: verbose_name = _("Schedule") @@ -42,7 +40,7 @@ class Schedule(SerializerModel): ) def __str__(self): - return self.name + return self.uid @property def serializer(self): diff --git a/authentik/tasks/schedules/scheduler.py b/authentik/tasks/schedules/scheduler.py index a373fc27cf..01763f65a6 100644 --- a/authentik/tasks/schedules/scheduler.py +++ b/authentik/tasks/schedules/scheduler.py @@ -2,7 +2,8 @@ import pickle # nosec import pglock from django.db import router, transaction -from django.utils.timezone import now +from django.utils.timezone import now, timedelta +from dramatiq.actor import Actor from dramatiq.broker import Broker from structlog.stdlib import get_logger @@ -22,15 +23,15 @@ class Scheduler: next_run = schedule.calculate_next_run(next_run) if next_run > now(): break + # Force to calculate the one after + next_run += timedelta(minutes=2) schedule.next_run = next_run - actor = self.broker.get_actor(schedule.actor_name) + actor: Actor = self.broker.get_actor(schedule.actor_name) actor.send_with_options( args=pickle.loads(schedule.args), # nosec kwargs=pickle.loads(schedule.kwargs), # nosec - options={ - "schedule_uid": schedule.uid, - }, + schedule_uid=schedule.uid, ) schedule.save()