From 3c8232b9a5c945dd6beb4e2fb04b1e86e07fd5a1 Mon Sep 17 00:00:00 2001 From: Jens Langhammer Date: Thu, 18 Apr 2024 01:15:27 +0200 Subject: [PATCH] add database aware scheduler with tenant support Signed-off-by: Jens Langhammer --- authentik/enterprise/reporting/__init__.py | 0 authentik/enterprise/reporting/apps.py | 12 ++++++++ authentik/enterprise/reporting/models.py | 18 +++++++++++ authentik/enterprise/reporting/signals.py | 35 ++++++++++++++++++++++ authentik/enterprise/reporting/tasks.py | 10 +++++++ authentik/enterprise/settings.py | 1 + authentik/root/settings.py | 1 + authentik/tenants/scheduler.py | 14 +++++---- 8 files changed, 86 insertions(+), 5 deletions(-) create mode 100644 authentik/enterprise/reporting/__init__.py create mode 100644 authentik/enterprise/reporting/apps.py create mode 100644 authentik/enterprise/reporting/models.py create mode 100644 authentik/enterprise/reporting/signals.py create mode 100644 authentik/enterprise/reporting/tasks.py diff --git a/authentik/enterprise/reporting/__init__.py b/authentik/enterprise/reporting/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/authentik/enterprise/reporting/apps.py b/authentik/enterprise/reporting/apps.py new file mode 100644 index 0000000000..c212e5d204 --- /dev/null +++ b/authentik/enterprise/reporting/apps.py @@ -0,0 +1,12 @@ +"""Reporting app config""" + +from authentik.enterprise.apps import EnterpriseConfig + + +class AuthentikEnterpriseReporting(EnterpriseConfig): + """authentik enterprise reporting app config""" + + name = "authentik.enterprise.reporting" + label = "authentik_reporting" + verbose_name = "authentik Enterprise.Reporting" + default = True diff --git a/authentik/enterprise/reporting/models.py b/authentik/enterprise/reporting/models.py new file mode 100644 index 0000000000..0c0c0a91f6 --- /dev/null +++ b/authentik/enterprise/reporting/models.py @@ -0,0 +1,18 @@ +from uuid import uuid4 + +from django.db import models + + +class Report(models.Model): + + report_uuid = models.UUIDField(primary_key=True, editable=False, default=uuid4) + + name = models.TextField() + + schedule = models.TextField() + + def __str__(self) -> str: + return self.name + + def do_the_thing(self): + pass diff --git a/authentik/enterprise/reporting/signals.py b/authentik/enterprise/reporting/signals.py new file mode 100644 index 0000000000..c93ac81c9a --- /dev/null +++ b/authentik/enterprise/reporting/signals.py @@ -0,0 +1,35 @@ +from json import dumps + +from celery.schedules import crontab +from django.db.models.signals import post_save, pre_delete +from django.dispatch import receiver +from django_celery_beat.models import CrontabSchedule, PeriodicTask + +from authentik.enterprise.reporting.models import Report + + +@receiver(post_save, sender=Report) +def report_post_save(sender, instance: Report, **_): + schedule = CrontabSchedule.from_schedule(crontab()) + schedule.save() + PeriodicTask.objects.update_or_create( + name=str(instance.pk), + defaults={ + "crontab": schedule, + "task": "authentik.enterprise.reporting.tasks.process_report", + "queue": "authentik_reporting", + "description": f"Report {instance.name}", + "kwargs": dumps( + { + "report_uuid": str(instance.pk), + } + ), + }, + ) + + +@receiver(pre_delete, sender=Report) +def report_pre_delete(sender, instance: Report, **_): + PeriodicTask.objects.filter(name=str(instance.pk)).delete() + # Cleanup schedules without any tasks + CrontabSchedule.objects.filter(periodictask__isnull=True).delete() diff --git a/authentik/enterprise/reporting/tasks.py b/authentik/enterprise/reporting/tasks.py new file mode 100644 index 0000000000..63cfb8fd6e --- /dev/null +++ b/authentik/enterprise/reporting/tasks.py @@ -0,0 +1,10 @@ +from authentik.enterprise.reporting.models import Report +from authentik.root.celery import CELERY_APP + + +@CELERY_APP.task() +def process_report(report_uuid: str): + report = Report.objects.filter(pk=report_uuid).first() + if not report: + return + report.do_the_thing() diff --git a/authentik/enterprise/settings.py b/authentik/enterprise/settings.py index 7f735eb312..fc55da66e4 100644 --- a/authentik/enterprise/settings.py +++ b/authentik/enterprise/settings.py @@ -17,6 +17,7 @@ TENANT_APPS = [ "authentik.enterprise.providers.google_workspace", "authentik.enterprise.providers.microsoft_entra", "authentik.enterprise.providers.ssf", + "authentik.enterprise.reporting", "authentik.enterprise.stages.authenticator_endpoint_gdtc", "authentik.enterprise.stages.source", ] diff --git a/authentik/root/settings.py b/authentik/root/settings.py index 24f56e7b39..43979bf212 100644 --- a/authentik/root/settings.py +++ b/authentik/root/settings.py @@ -125,6 +125,7 @@ TENANT_APPS = [ "authentik.brands", "authentik.blueprints", "guardian", + "django_celery_beat", ] TENANT_MODEL = "authentik_tenants.Tenant" diff --git a/authentik/tenants/scheduler.py b/authentik/tenants/scheduler.py index 753831ae0b..3ae81c149d 100644 --- a/authentik/tenants/scheduler.py +++ b/authentik/tenants/scheduler.py @@ -1,14 +1,18 @@ """Tenant-aware Celery beat scheduler""" -from tenant_schemas_celery.scheduler import ( - TenantAwarePersistentScheduler as BaseTenantAwarePersistentScheduler, -) -from tenant_schemas_celery.scheduler import TenantAwareScheduleEntry +from django_celery_beat.schedulers import DatabaseScheduler, ModelEntry +from tenant_schemas_celery.scheduler import TenantAwareScheduleEntry, TenantAwareSchedulerMixin -class TenantAwarePersistentScheduler(BaseTenantAwarePersistentScheduler): +class SchedulerEntry(ModelEntry, TenantAwareScheduleEntry): + pass + + +class TenantAwarePersistentScheduler(TenantAwareSchedulerMixin, DatabaseScheduler): """Tenant-aware Celery beat scheduler""" + Entry = SchedulerEntry + @classmethod def get_queryset(cls): return super().get_queryset().filter(ready=True)