From c2f7883a5ce3433cfa4daf44e465310191416b0e Mon Sep 17 00:00:00 2001
From: Marc 'risson' Schmitt <marc.schmitt@risson.space>
Date: Thu, 12 Jun 2025 12:12:31 +0200
Subject: [PATCH] start work on moving to a package

Signed-off-by: Marc 'risson' Schmitt <marc.schmitt@risson.space>
---
 authentik/root/settings.py                    | 34 +++++++++++++--
 .../django_dramatiq_postgres/apps.py          | 42 +++++++++++++++++++
 .../django_dramatiq_postgres/conf.py          | 26 ++++++++++++
 3 files changed, 99 insertions(+), 3 deletions(-)
 create mode 100644 packages/django-dramatiq-postgres/django_dramatiq_postgres/apps.py
 create mode 100644 packages/django-dramatiq-postgres/django_dramatiq_postgres/conf.py

diff --git a/authentik/root/settings.py b/authentik/root/settings.py
index 479311ab63..7889e9032b 100644
--- a/authentik/root/settings.py
+++ b/authentik/root/settings.py
@@ -344,6 +344,37 @@ USE_TZ = True
 LOCALE_PATHS = ["./locale"]
 
 
+# Tests
+
+TEST = False
+TEST_RUNNER = "authentik.root.test_runner.PytestTestRunner"
+
+
+# DramatiQ
+
+DRAMATIQ = {
+    "middlewares": (
+        # TODO: fixme
+        # ("dramatiq.middleware.prometheus.Prometheus", {}),
+        ("dramatiq.middleware.age_limit.AgeLimit", {}),
+        (
+            "dramatiq.middleware.time_limit.TimeLimit",
+            {
+                # 5 minutes task timeout by default for all tasks
+                "time_limit": 600 * 1000,
+            },
+        ),
+        ("dramatiq.middleware.shutdown.ShutdownNotifications", {}),
+        ("dramatiq.middleware.callbacks.Callbacks", {}),
+        ("dramatiq.middleware.pipelines.Pipelines", {}),
+        ("dramatiq.middleware.retries.Retries", {"max_retries": 20 if not TEST else 0}),
+        # TODO: results
+        ("authentik.tasks.middleware.FullyQualifiedActorName", {}),
+        ("authentik.tasks.middleware.CurrentTask", {}),
+    )
+}
+
+
 # Sentry integration
 
 env = get_env()
@@ -406,9 +437,6 @@ else:
     MEDIA_ROOT = STORAGES["default"]["OPTIONS"]["location"]
     MEDIA_URL = STORAGES["default"]["OPTIONS"]["base_url"]
 
-TEST = False
-TEST_RUNNER = "authentik.root.test_runner.PytestTestRunner"
-
 structlog_configure()
 LOGGING = get_logger_config()
 
diff --git a/packages/django-dramatiq-postgres/django_dramatiq_postgres/apps.py b/packages/django-dramatiq-postgres/django_dramatiq_postgres/apps.py
new file mode 100644
index 0000000000..677658ecd0
--- /dev/null
+++ b/packages/django-dramatiq-postgres/django_dramatiq_postgres/apps.py
@@ -0,0 +1,42 @@
+import dramatiq
+from django.apps import AppConfig
+from django.core.exceptions import ImproperlyConfigured
+from django.utils.module_loading import import_string
+
+from django_dramatiq_postgres.conf import Conf
+
+
+class DjangoDramatiqPostgres(AppConfig):
+    name = "django_dramatiq_postgres"
+    verbose_name = "Django DramatiQ postgres"
+
+    def ready(self):
+        old_broker = dramatiq.get_broker()
+
+        if len(old_broker.actors) != 0:
+            raise ImproperlyConfigured(
+                "Actors were previously registered. "
+                "Make sure your actors are not imported too early."
+            )
+
+        encoder: dramatiq.encoder.Encoder = import_string(Conf.encoder_class)()
+        dramatiq.set_encoder(encoder)
+
+        broker_args = Conf.broker_args
+        broker_kwargs = {
+            **Conf.broker_kwargs,
+            "middleware": [],
+        }
+        broker: dramatiq.broker.Broker = import_string(Conf.broker_class)(
+            *broker_args, **broker_kwargs
+        )
+
+        for middleware_class, middleware_kwargs in Conf.middlewares.items():
+            middleware: dramatiq.middleware.middleware.Middleware = import_string(middleware_class)(
+                **middleware_kwargs
+            )
+            broker.add_middleware(middleware)
+
+        dramatiq.set_broker(broker)
+
+        return super().ready()
diff --git a/packages/django-dramatiq-postgres/django_dramatiq_postgres/conf.py b/packages/django-dramatiq-postgres/django_dramatiq_postgres/conf.py
new file mode 100644
index 0000000000..64a26c8803
--- /dev/null
+++ b/packages/django-dramatiq-postgres/django_dramatiq_postgres/conf.py
@@ -0,0 +1,26 @@
+from django.conf import settings
+
+
+class Conf:
+    try:
+        conf = settings.DRAMATIQ.copy()
+    except AttributeError:
+        conf = {}
+
+    encoder_class = conf.get("encoder_class", "dramatiq.encoder.PickleEncoder")
+
+    broker_class = conf.get("broker_class", "django_dramatiq_postgres.broker.PostgresBroker")
+    broker_args = conf.get("broker_args", ())
+    broker_kwargs = conf.get("broker_kwargs", {})
+
+    middlewares = conf.get(
+        "middlewares",
+        (
+            ("dramatiq.middleware.age_limit.AgeLimit", {}),
+            ("dramatiq.middleware.time_limit.TimeLimit", {}),
+            ("dramatiq.middleware.shutdown.ShutdownNotifications", {}),
+            ("dramatiq.middleware.callbacks.Callbacks", {}),
+            ("dramatiq.middleware.pipelines.Pipelines", {}),
+            ("dramatiq.middleware.retries.Retries", {}),
+        ),
+    )