Files
authentik/packages/django-dramatiq-postgres/django_dramatiq_postgres/middleware.py
Marc 'risson' Schmitt 16fd9cab67 move all broker stuff to package, schedule is still todo
Signed-off-by: Marc 'risson' Schmitt <marc.schmitt@risson.space>
2025-06-18 17:01:05 +02:00

25 lines
695 B
Python

from django.db import (
close_old_connections,
connections,
)
from dramatiq.middleware.middleware import Middleware
from django_dramatiq_postgres.conf import Conf
class DbConnectionMiddleware(Middleware):
def _close_old_connections(self, *args, **kwargs):
if Conf().test:
return
close_old_connections()
before_process_message = _close_old_connections
after_process_message = _close_old_connections
def _close_connections(self, *args, **kwargs):
connections.close_all()
before_consumer_thread_shutdown = _close_connections
before_worker_thread_shutdown = _close_connections
before_worker_shutdown = _close_connections