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

wip

Browse Source 
Signed-off-by: Marc 'risson' Schmitt <marc.schmitt@risson.space>
This commit is contained in:
Marc 'risson' Schmitt
2025-03-09 13:38:15 +01:00
parent c4b988c632
commit 6662611347
4 changed files with 115 additions and 27 deletions
Download Patch File Download Diff File Expand all files Collapse all files

24
authentik/tasks/broker.py
View File

@ -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")

84
authentik/tasks/migrations/0001_initial.py Normal file
View File

@ -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",
),
),
),
]

30
authentik/tasks/models.py
View File

@ -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;

4
authentik/tasks/results.py
View File

@ -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)