Github/authentik
1
0
Fork 0
mirror of https://github.com/goauthentik/authentik.git synced 2026-08-30 18:51:39 -07:00
Code Issues Packages Projects Releases Wiki Activity

packages/django-dramatiq-postgres/broker: use positive state filter for pending messages (#24074)

Browse Source 
Co-authored-by: Justin Schuhmann <jschuhmann@Justins-MBP.home>
This commit is contained in:
Justin Schuhmann
2026-07-17 12:47:47 -04:00
committed by GitHub
parent aedf7230d2
commit ae6882c95c
2 changed files with 31 additions and 14 deletions
Download Patch File Download Diff File Expand all files Collapse all files

23
authentik/tasks/tests/test_broker.py
View File

@@ -1,7 +1,7 @@
from unittest.mock import MagicMock
from django.test import SimpleTestCase
from django_dramatiq_postgres.broker import PostgresBroker
from django_dramatiq_postgres.broker import CONSUMABLE_TASK_STATES, PostgresBroker
from django_dramatiq_postgres.models import TaskState
from dramatiq.broker import MessageProxy
from dramatiq.message import Message
@@ -93,3 +93,24 @@ class TestPostgresConsumer(SimpleTestCase):
{first_message.message_id, second_message.message_id},
)
self.assertNotIn(first_message.message_id, consumer.in_processing)
def test_fetch_pending_messages_filters_consumable_states(self):
consumer = self._consumer()
consumer.queue_name = "default"
consumer.timeout = 30
query_set = consumer.query_set
queue_query = query_set.exclude.return_value.filter.return_value
pending_query = queue_query.filter.return_value
values_list = pending_query.exclude.return_value.order_by.return_value.values_list
values_list.return_value = ["00000000-0000-0000-0000-000000000001"]
pending = consumer._fetch_pending_messages()
query_set.exclude.assert_called_once_with(message_id__in=consumer.in_processing)
query_set.exclude.return_value.filter.assert_called_once_with(
queue_name=consumer.queue_name,
)
queue_query.filter.assert_called_once_with(
state__in=CONSUMABLE_TASK_STATES,
)
self.assertEqual(pending, {"00000000-0000-0000-0000-000000000001"})

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

@@ -48,6 +48,12 @@ DATABASE_ERRORS = (
OperationalError,
)
CONSUMABLE_TASK_STATES: set[TaskState] = set(TaskState) - {
TaskState.DONE,
TaskState.REJECTED,
TaskState.WAITING_FOR_DEPENDENCIES,
}
def channel_name(queue_name: str, identifier: ChannelIdentifier) -> str:
return f"{CHANNEL_PREFIX}.{queue_name}.{identifier.value}"
@@ -375,13 +381,7 @@ class _PostgresConsumer(Consumer):
pending = set(
self.query_set.exclude(message_id__in=self.in_processing)
.filter(queue_name=self.queue_name)
.exclude(
state__in=(
TaskState.DONE,
TaskState.REJECTED,
TaskState.WAITING_FOR_DEPENDENCIES,
)
)
.filter(state__in=CONSUMABLE_TASK_STATES)
.exclude(eta__gte=timezone.now() + timedelta(seconds=self.timeout))
.order_by(F("eta").asc(nulls_first=True))
.values_list("message_id", flat=True)
@@ -417,7 +417,7 @@ class _PostgresConsumer(Consumer):
WHERE
{table}.{message_id} = %(message_id)s
AND
{table}.{state} != ALL(%(excluded_states)s)
{table}.{state} = ANY(%(consumable_states)s)
AND
({table}.{eta} < %(maximum_eta)s OR {table}.{eta} IS NULL)
AND
@@ -433,11 +433,7 @@ class _PostgresConsumer(Consumer):
"state": TaskState.CONSUMED.value,
"mtime": timezone.now(),
"message_id": message_id,
"excluded_states": [
TaskState.DONE.value,
TaskState.REJECTED.value,
TaskState.WAITING_FOR_DEPENDENCIES.value,
],
"consumable_states": [state.value for state in CONSUMABLE_TASK_STATES],
"maximum_eta": timezone.now() + timedelta(seconds=self.timeout),
"lock_id": self._get_message_lock_id(message_id),
},