habsi/authentik
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

move all broker stuff to package, schedule is still todo

Browse Source 
Signed-off-by: Marc 'risson' Schmitt <marc.schmitt@risson.space>
This commit is contained in:
Marc 'risson' Schmitt
2025-06-18 17:01:05 +02:00
parent 8c7818a252
commit 16fd9cab67
9 changed files with 113 additions and 187 deletions
Download Patch File Download Diff File Expand all files Collapse all files

69
authentik/tasks/apps.py
View File

@ -1,17 +1,3 @@
import dramatiq
from dramatiq.broker import Broker, get_broker
from dramatiq.encoder import PickleEncoder
from dramatiq.middleware import (
AgeLimit,
Callbacks,
Pipelines,
# Prometheus,
Retries,
ShutdownNotifications,
TimeLimit,
)
from dramatiq.results.middleware import Results
from authentik.blueprints.apps import ManagedAppConfig
@ -21,47 +7,14 @@ class AuthentikTasksConfig(ManagedAppConfig):
verbose_name = "authentik Tasks"
default = True
def _set_dramatiq_middlewares(self, broker: Broker, max_retries: int = 20) -> None:
from authentik.tasks.middleware import CurrentTask, FullyQualifiedActorName
from authentik.tasks.results import PostgresBackend
# TODO: fixme
# broker.add_middleware(Prometheus())
broker.add_middleware(AgeLimit())
# Task timeout, 5 minutes by default for all tasks
broker.add_middleware(TimeLimit(time_limit=600 * 1000))
broker.add_middleware(ShutdownNotifications())
broker.add_middleware(Callbacks())
broker.add_middleware(Pipelines())
broker.add_middleware(Retries(max_retries=max_retries))
broker.add_middleware(Results(backend=PostgresBackend(), store_results=True))
broker.add_middleware(FullyQualifiedActorName())
broker.add_middleware(CurrentTask())
def ready(self) -> None:
from authentik.tasks.broker import PostgresBroker
old_broker = dramatiq.get_broker()
if len(old_broker.actors) != 0:
raise RuntimeError("Mis-registered actors")
dramatiq.set_encoder(PickleEncoder())
broker = PostgresBroker(middleware=[])
self._set_dramatiq_middlewares(broker)
dramatiq.set_broker(broker)
return super().ready()
def use_test_broker(self) -> None:
from authentik.tasks.test import TestBroker
old_broker = get_broker()
broker = TestBroker(middleware=[])
self._set_dramatiq_middlewares(broker, max_retries=0)
dramatiq.set_broker(broker)
for actor_name in old_broker.get_declared_actors():
actor = old_broker.get_actor(actor_name)
actor.broker = broker
actor.broker.declare_actor(actor)
# def use_test_broker(self) -> None:
# from authentik.tasks.test import TestBroker
#
# old_broker = get_broker()
# broker = TestBroker(middleware=[])
# self._set_dramatiq_middlewares(broker, max_retries=0)
# dramatiq.set_broker(broker)
# for actor_name in old_broker.get_declared_actors():
# actor = old_broker.get_actor(actor_name)
# actor.broker = broker
# actor.broker.declare_actor(actor)

62
authentik/tasks/broker.py
View File

@ -1,73 +1,15 @@
from typing import Any
from django_dramatiq_postgres.middleware import DbConnectionMiddleware
from django.db import (
DEFAULT_DB_ALIAS,
)
from dramatiq.broker import Broker
from django_dramatiq_postgres.broker import PostgresBroker
from dramatiq.message import Message
from dramatiq.middleware import (
AgeLimit,
Callbacks,
Middleware,
Pipelines,
Prometheus,
Retries,
ShutdownNotifications,
TimeLimit,
)
from structlog.stdlib import get_logger
from django_dramatiq_postgres.broker import PostgresBroker as PostgresBrokerBase
from authentik.tasks.models import Task
from authentik.tenants.models import Tenant
from authentik.tenants.utils import get_current_tenant
LOGGER = get_logger()
class TenantMiddleware(Middleware):
def before_process_message(self, broker: Broker, message: Message):
Task.objects.select_related("tenant").get(message_id=message.message_id).tenant.activate()
def after_process_message(self, *args, **kwargs):
Tenant.deactivate()
class PostgresBroker(PostgresBrokerBase):
def __init__(
self,
*args,
middleware: list[Middleware] | None = None,
db_alias: str = DEFAULT_DB_ALIAS,
**kwargs,
):
super().__init__(*args, middleware=[], **kwargs)
self.logger = get_logger().bind()
self.queues = set()
self.actor_options = {
"rel_obj",
}
self.db_alias = db_alias
self.middleware = []
self.add_middleware(DbConnectionMiddleware())
self.add_middleware(TenantMiddleware())
if middleware is None:
for m in (
Prometheus,
AgeLimit,
TimeLimit,
ShutdownNotifications,
Callbacks,
Pipelines,
Retries,
):
self.add_middleware(m())
for m in middleware or []:
self.add_middleware(m)
class Broker(PostgresBroker):
def model_defaults(self, message: Message) -> dict[str, Any]:
rel_obj = message.options.get("rel_obj")
if rel_obj:

15
authentik/tasks/middleware.py
View File

@ -8,10 +8,17 @@ from dramatiq.middleware import Middleware
from structlog.stdlib import get_logger
from authentik.tasks.models import Task
from authentik.tenants.models import Tenant
LOGGER = get_logger()
class RelObjMiddleware(Middleware):
@property
def actor_options(self):
return {"rel_obj"}
class FullyQualifiedActorName(Middleware):
def before_declare_actor(self, broker: Broker, actor: Actor):
actor.actor_name = f"{actor.fn.__module__}.{actor.fn.__name__}"
@ -52,3 +59,11 @@ class CurrentTask(Middleware):
else:
tasks[-1].save()
self._TASK.set(tasks[:-1])
class TenantMiddleware(Middleware):
def before_process_message(self, broker: Broker, message: Message):
Task.objects.select_related("tenant").get(message_id=message.message_id).tenant.activate()
def after_process_message(self, *args, **kwargs):
Tenant.deactivate()