From 6662611347bd80c58f898dbc19fb7dfcddcf2503 Mon Sep 17 00:00:00 2001 From: Marc 'risson' Schmitt Date: Sun, 9 Mar 2025 13:38:15 +0100 Subject: [PATCH] wip Signed-off-by: Marc 'risson' Schmitt --- authentik/tasks/broker.py | 24 +++---- authentik/tasks/migrations/0001_initial.py | 84 ++++++++++++++++++++++ authentik/tasks/models.py | 30 ++++---- authentik/tasks/results.py | 4 +- 4 files changed, 115 insertions(+), 27 deletions(-) create mode 100644 authentik/tasks/migrations/0001_initial.py diff --git a/authentik/tasks/broker.py b/authentik/tasks/broker.py index 6028820c6a..2c3040cb1d 100644 --- a/authentik/tasks/broker.py +++ b/authentik/tasks/broker.py @@ -21,7 +21,7 @@ from psycopg import Notify from psycopg.errors import AdminShutdown from structlog.stdlib import get_logger -from authentik.tasks.models import Task, CHANNEL_PREFIX, ChannelIdentifier +from authentik.tasks.models import Task, CHANNEL_PREFIX, ChannelIdentifier, TaskState from authentik.tasks.results import PostgresBackend from authentik.tenants.utils import get_current_tenant @@ -125,7 +125,7 @@ class PostgresBroker(Broker): defaults = { "tenant": get_current_tenant(), "queue_name": message.queue_name, - "state": Task.State.QUEUED, + "state": TaskState.QUEUED, "message": encoded, } create_defaults = { @@ -169,7 +169,7 @@ class PostgresBroker(Broker): if ( self.query_set.filter( queue_name=queue_name, - state__in=(Task.State.QUEUED, Task.State.CONSUMED), + state__in=(TaskState.QUEUED, TaskState.CONSUMED), ) == 0 ): @@ -220,9 +220,9 @@ class _PostgresConsumer(Consumer): self.query_set.filter( message_id=message.message_id, queue_name=message.queue_name, - state=Task.State.CONSUMED, + state=TaskState.CONSUMED, ).update( - state=Task.State.DONE, + state=TaskState.DONE, message=message.encode(), ) self.in_processing.remove(message.message_id) @@ -233,9 +233,9 @@ class _PostgresConsumer(Consumer): self.query_set.filter( message_id=message.message_id, queue_name=message.queue_name, - state__ne=Task.State.REJECTED, + state__ne=TaskState.REJECTED, ).update( - state=Task.State.REJECTED, + state=TaskState.REJECTED, message=message.encode(), ) self.in_processing.remove(message.message_id) @@ -245,14 +245,14 @@ class _PostgresConsumer(Consumer): self.query_set.filter( message_id__in=[message.message_id for message in messages], ).update( - state=Task.State.QUEUED, + state=TaskState.QUEUED, ) # We don't care about locks, requeue occurs on worker stop def _fetch_pending_notifies(self) -> list[Notify]: self.logger.debug(f"Polling for lost messages in {self.queue_name}") notifies = self.query_set.filter( - state__in=(Task.State.QUEUED, Task.State.CONSUMED), queue_name=self.queue_name + state__in=(TaskState.QUEUED, TaskState.CONSUMED), queue_name=self.queue_name ) channel = channel_name(self.queue_name, ChannelIdentifier.ENQUEUE) return [Notify(pid=0, channel=channel, payload=item.message) for item in notifies] @@ -276,9 +276,9 @@ class _PostgresConsumer(Consumer): result = ( self.query_set.filter( message_id=message.message_id, - state__in=(Task.State.QUEUED, Task.State.CONSUMED), + state__in=(TaskState.QUEUED, TaskState.CONSUMED), ) - .update(state=Task.State.CONSUMED, mtime=timezone.now()) + .update(state=TaskState.CONSUMED, mtime=timezone.now()) .extra(where=["pg_try_advisory_lock(%s)"], params=[self._get_message_lock_id(message)]) ) return result == 1 @@ -354,7 +354,7 @@ class _PostgresConsumer(Consumer): return self.logger.debug("Running garbage collector") count = self.query_set.filter( - state__in=(Task.State.DONE, Task.State.REJECTED), + state__in=(TaskState.DONE, TaskState.REJECTED), mtime__lte=timezone.now() - timezone.timedelta(days=30), ).delete() self.logger.info(f"Purged {count} messages in all queues") diff --git a/authentik/tasks/migrations/0001_initial.py b/authentik/tasks/migrations/0001_initial.py new file mode 100644 index 0000000000..2d68423bf7 --- /dev/null +++ b/authentik/tasks/migrations/0001_initial.py @@ -0,0 +1,84 @@ +# Generated by Django 5.0.12 on 2025-03-09 12:37 + +import django.db.models.deletion +import django.utils.timezone +import pgtrigger.compiler +import pgtrigger.migrations +import uuid +from django.db import migrations, models + + +class Migration(migrations.Migration): + + initial = True + + dependencies = [ + ("authentik_tenants", "0004_tenant_impersonation_require_reason"), + ] + + operations = [ + migrations.CreateModel( + name="Task", + fields=[ + ( + "message_id", + models.UUIDField( + default=uuid.uuid4, editable=False, primary_key=True, serialize=False + ), + ), + ("queue_name", models.TextField(default="default", editable=False)), + ( + "state", + models.CharField( + choices=[ + ("queued", "Queued"), + ("consumed", "Consumed"), + ("rejected", "Rejected"), + ("done", "Done"), + ], + default="queued", + editable=False, + ), + ), + ("mtime", models.DateTimeField(default=django.utils.timezone.now, editable=False)), + ("message", models.JSONField(editable=False, null=True)), + ("result", models.JSONField(editable=False, null=True)), + ("result_ttl", models.DateTimeField(editable=False, null=True)), + ("description", models.TextField(blank=True)), + ("messages", models.JSONField(blank=True, editable=False, null=True)), + ( + "tenant", + models.ForeignKey( + editable=False, + on_delete=django.db.models.deletion.CASCADE, + to="authentik_tenants.tenant", + ), + ), + ], + options={ + "indexes": [ + models.Index(fields=["state", "mtime"], name="authentik_t_state_bb4a31_idx") + ], + }, + ), + migrations.RunSQL( + "ALTER TABLE authentik_tasks_queue SET WITHOUT OIDS;", migrations.RunSQL.noop + ), + pgtrigger.migrations.AddTrigger( + model_name="task", + trigger=pgtrigger.compiler.Trigger( + name="notify_enqueueing", + sql=pgtrigger.compiler.UpsertTriggerSql( + condition="WHEN (NEW.\"state\" = 'queued')", + constraint="CONSTRAINT", + func="\n PERFORM pg_notify(\n 'authentik.tasks' || NEW.queue_name || '.enqueue',\n CASE WHEN octet_length(NEW.message::text) >= 8000\n THEN jsonb_build_object('message_id', NEW.message_id)::text\n ELSE NEW.message::text\n END\n );\n RETURN NEW;\n ", + hash="37184bcf29160694f794426a0246c3c1a5e8e702", + operation="INSERT OR UPDATE", + pgid="pgtrigger_notify_enqueueing_0bc94", + table="authentik_tasks_task", + timing="DEFERRABLE INITIALLY DEFERRED", + when="AFTER", + ), + ), + ), + ] diff --git a/authentik/tasks/models.py b/authentik/tasks/models.py index b081c44ec7..d5a57e0b0d 100644 --- a/authentik/tasks/models.py +++ b/authentik/tasks/models.py @@ -16,21 +16,25 @@ class ChannelIdentifier(StrEnum): LOCK = auto() -class Task(SerializerModel): - class State(models.TextChoices): - QUEUED = "queued" - CONSUMED = "consumed" - REJECTED = "rejected" - DONE = "done" +class TaskState(models.TextChoices): + QUEUED = "queued" + CONSUMED = "consumed" + REJECTED = "rejected" + DONE = "done" + +class Task(SerializerModel): message_id = models.UUIDField(primary_key=True, default=uuid4, editable=False) - tenant = models.ForeignKey(Tenant, on_delete=models.CASCADE, editable=False) queue_name = models.TextField(default="default", editable=False) - state = models.CharField(default=State.QUEUED, choices=State.choices, editable=False) + + tenant = models.ForeignKey(Tenant, on_delete=models.CASCADE, editable=False) + state = models.CharField(default=TaskState.QUEUED, choices=TaskState.choices, editable=False) mtime = models.DateTimeField(default=timezone.now, editable=False) - message = models.JSONField(blank=True, null=True, editable=False) - result = models.JSONField(blank=True, null=True, editable=False) - result_ttl = models.DateTimeField(blank=True, null=True, editable=False) + message = models.JSONField(null=True, editable=False) + + result = models.JSONField(null=True, editable=False) + result_ttl = models.DateTimeField(null=True, editable=False) + description = models.TextField(blank=True) messages = models.JSONField(blank=True, null=True, editable=False) @@ -41,14 +45,14 @@ class Task(SerializerModel): name="notify_enqueueing", operation=pgtrigger.Insert | pgtrigger.Update, when=pgtrigger.After, - condition=pgtrigger.Q(new__state="queued"), + condition=pgtrigger.Q(new__state=TaskState.QUEUED), timing=pgtrigger.Deferred, func=f""" PERFORM pg_notify( '{CHANNEL_PREFIX}' || NEW.queue_name || '.{ChannelIdentifier.ENQUEUE.value}', CASE WHEN octet_length(NEW.message::text) >= 8000 THEN jsonb_build_object('message_id', NEW.message_id)::text - ELSE message::text + ELSE NEW.message::text END ); RETURN NEW; diff --git a/authentik/tasks/results.py b/authentik/tasks/results.py index 774f5e0969..363f15467f 100644 --- a/authentik/tasks/results.py +++ b/authentik/tasks/results.py @@ -4,7 +4,7 @@ from django.utils import timezone from dramatiq.message import Message, get_encoder from dramatiq.results.backend import Missing, MResult, Result, ResultBackend -from authentik.tasks.models import Task +from authentik.tasks.models import Task, TaskState class PostgresBackend(ResultBackend): @@ -42,6 +42,6 @@ class PostgresBackend(ResultBackend): **query, **defaults, "queue_name": "__RQ__", - "state": Queue.State.DONE, + "state": TaskState.DONE, } self.query_set.update_or_create(**query, defaults=defaults, create_defaults=create_defaults)