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

10
packages/django-dramatiq-postgres/django_dramatiq_postgres/apps.py
View File

@ -19,20 +19,20 @@ class DjangoDramatiqPostgres(AppConfig):
"Make sure your actors are not imported too early."
)
encoder: dramatiq.encoder.Encoder = import_string(Conf.encoder_class)()
encoder: dramatiq.encoder.Encoder = import_string(Conf().encoder_class)()
dramatiq.set_encoder(encoder)
broker_args = Conf.broker_args
broker_args = Conf().broker_args
broker_kwargs = {
**Conf.broker_kwargs,
**Conf().broker_kwargs,
"middleware": [],
}
broker: dramatiq.broker.Broker = import_string(Conf.broker_class)(
broker: dramatiq.broker.Broker = import_string(Conf().broker_class)(
*broker_args,
**broker_kwargs,
)
for middleware_class, middleware_kwargs in Conf.middlewares:
for middleware_class, middleware_kwargs in Conf().middlewares:
middleware: dramatiq.middleware.middleware.Middleware = import_string(middleware_class)(
**middleware_kwargs,
)

51
packages/django-dramatiq-postgres/django_dramatiq_postgres/broker.py
View File

@ -7,6 +7,7 @@ from random import randint
from typing import Any
import tenacity
from django.core.exceptions import ImproperlyConfigured
from django.db import (
DEFAULT_DB_ALIAS,
DatabaseError,
@ -24,14 +25,7 @@ from dramatiq.common import compute_backoff, current_millis, dq_name, xq_name
from dramatiq.errors import ConnectionError, QueueJoinTimeout
from dramatiq.message import Message
from dramatiq.middleware import (
AgeLimit,
Callbacks,
Middleware,
Pipelines,
Prometheus,
Retries,
ShutdownNotifications,
TimeLimit,
)
from pglock.core import _cast_lock_id
from psycopg import Notify, sql
@ -39,7 +33,6 @@ from psycopg.errors import AdminShutdown
from structlog.stdlib import get_logger
from django_dramatiq_postgres.conf import Conf
from django_dramatiq_postgres.middleware import DbConnectionMiddleware
from django_dramatiq_postgres.models import CHANNEL_PREFIX, ChannelIdentifier, TaskBase, TaskState
LOGGER = get_logger()
@ -75,20 +68,10 @@ class PostgresBroker(Broker):
self.db_alias = db_alias
self.middleware = []
if middleware is None:
for m in (
DbConnectionMiddleware,
Prometheus,
AgeLimit,
TimeLimit,
ShutdownNotifications,
Callbacks,
Pipelines,
Retries,
):
self.add_middleware(m())
for m in middleware or []:
self.add_middleware(m)
if middleware:
raise ImproperlyConfigured(
"Middlewares should be set in django settings, not passed directly to the broker."
)
@property
def connection(self) -> DatabaseWrapper:
@ -100,7 +83,7 @@ class PostgresBroker(Broker):
@cached_property
def model(self) -> type[TaskBase]:
return import_string(Conf.task_class)
return import_string(Conf().task_class)
@property
def query_set(self) -> QuerySet:
@ -267,7 +250,6 @@ class _PostgresConsumer(Consumer):
@raise_connection_error
def ack(self, message: Message):
self.unlock_queue.put_nowait(message)
self.query_set.filter(
message_id=message.message_id,
queue_name=message.queue_name,
@ -276,11 +258,11 @@ class _PostgresConsumer(Consumer):
state=TaskState.DONE,
message=message.encode(),
)
self.unlock_queue.put_nowait(message.message_id)
self.in_processing.remove(message.message_id)
@raise_connection_error
def nack(self, message: Message):
self.unlock_queue.put_nowait(message)
self.query_set.filter(
message_id=message.message_id,
queue_name=message.queue_name,
@ -290,18 +272,18 @@ class _PostgresConsumer(Consumer):
state=TaskState.REJECTED,
message=message.encode(),
)
self.unlock_queue.put_nowait(message.message_id)
self.in_processing.remove(message.message_id)
@raise_connection_error
def requeue(self, messages: Iterable[Message]):
for message in messages:
self.unlock_queue.put_nowait(message)
self.query_set.filter(
message_id__in=[message.message_id for message in messages],
).update(
state=TaskState.QUEUED,
)
for message in messages:
self.unlock_queue.put_nowait(message.message_id)
self.in_processing.remove(message.message_id)
self._purge_locks()
@ -328,9 +310,9 @@ class _PostgresConsumer(Consumer):
)
self.notifies += notifies
def _get_message_lock_id(self, message: Message) -> int:
def _get_message_lock_id(self, message_id: str) -> int:
return _cast_lock_id(
f"{channel_name(self.queue_name, ChannelIdentifier.LOCK)}.{message.message_id}"
f"{channel_name(self.queue_name, ChannelIdentifier.LOCK)}.{message_id}"
)
def _consume_one(self, message: Message) -> bool:
@ -345,7 +327,7 @@ class _PostgresConsumer(Consumer):
)
.extra(
where=["pg_try_advisory_lock(%s)"],
params=[self._get_message_lock_id(message)],
params=[self._get_message_lock_id(message.message_id)],
)
.update(
state=TaskState.CONSUMED,
@ -391,6 +373,7 @@ class _PostgresConsumer(Consumer):
notify = self.notifies.pop(0)
task = self.query_set.get(message_id=notify.payload)
message = Message.decode(task.message)
message.task = task
if self._consume_one(message):
self.in_processing.add(message.message_id)
return MessageProxy(message)
@ -404,17 +387,18 @@ class _PostgresConsumer(Consumer):
def _purge_locks(self):
while True:
try:
message = self.unlock_queue.get(block=False)
message_id = self.unlock_queue.get(block=False)
except Empty:
return
self.logger.debug(f"Unlocking {message.message_id}@{message.queue_name}")
self.logger.debug(f"Unlocking message {message_id}")
with self.connection.cursor() as cursor:
cursor.execute(
"SELECT pg_advisory_unlock(%s)", (self._get_message_lock_id(message),)
"SELECT pg_advisory_unlock(%s)", (self._get_message_lock_id(message_id),)
)
self.unlock_queue.task_done()
def _auto_purge(self):
# TODO: allow configuring this
# Automatically purge messages on average every 100k iteration.
# Dramatiq defaults to 1s, so this means one purge every 28 hours.
if randint(0, 100_000): # nosec
@ -422,6 +406,7 @@ class _PostgresConsumer(Consumer):
self.logger.debug("Running garbage collector")
count = self.query_set.filter(
state__in=(TaskState.DONE, TaskState.REJECTED),
# TODO: allow configuring this
mtime__lte=timezone.now() - timezone.timedelta(days=30),
result_expiry__lte=timezone.now(),
).delete()

70
packages/django-dramatiq-postgres/django_dramatiq_postgres/conf.py
View File

@ -1,33 +1,57 @@
from typing import Any
from django.conf import settings
from django.core.exceptions import ImproperlyConfigured
class Conf:
try:
conf = settings.DRAMATIQ.copy()
except AttributeError:
conf = {}
def __init__(self):
try:
self.conf = settings.DRAMATIQ
except AttributeError as exc:
raise ImproperlyConfigured("Setting DRAMATIQ not set.") from exc
if "task_class" not in self.conf:
raise ImproperlyConfigured("DRAMATIQ.task_class not defined")
encoder_class = conf.get("encoder_class", "dramatiq.encoder.PickleEncoder")
@property
def encoder_class(self) -> str:
return self.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", {})
@property
def broker_class(self) -> str:
return self.conf.get("broker_class", "django_dramatiq_postgres.broker.PostgresBroker")
middlewares = conf.get(
"middlewares",
(
("django_dramatiq_postgres.middleware.DbConnectionMiddleware", {}),
("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", {}),
),
)
@property
def broker_args(self) -> tuple[Any]:
return self.conf.get("broker_args", ())
channel_prefix = conf.get("channel_prefix", "dramatiq.tasks")
@property
def broker_kwargs(self) -> dict[str, Any]:
return self.conf.get("broker_kwargs", {})
task_class = conf.get("task_class", None)
@property
def middlewares(self) -> tuple[tuple[str, dict[str, Any]]]:
return self.conf.get(
"middlewares",
(
("django_dramatiq_postgres.middleware.DbConnectionMiddleware", {}),
("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", {}),
),
)
test = conf.get("test", False)
@property
def channel_prefix(self) -> str:
return self.conf.get("channel_prefix", "dramatiq.tasks")
@property
def task_class(self) -> str:
return self.conf["task_class"]
@property
def test(self) -> bool:
return self.conf.get("test", False)

2
packages/django-dramatiq-postgres/django_dramatiq_postgres/middleware.py
View File

@ -9,7 +9,7 @@ from django_dramatiq_postgres.conf import Conf
class DbConnectionMiddleware(Middleware):
def _close_old_connections(self, *args, **kwargs):
if Conf.test:
if Conf().test:
return
close_old_connections()

2
packages/django-dramatiq-postgres/django_dramatiq_postgres/models.py
View File

@ -8,7 +8,7 @@ from django.utils.translation import gettext_lazy as _
from django_dramatiq_postgres.conf import Conf
CHANNEL_PREFIX = f"{Conf.channel_prefix}.tasks"
CHANNEL_PREFIX = f"{Conf().channel_prefix}.tasks"
class ChannelIdentifier(StrEnum):