diff --git a/backend/notification_v2/internal_api_views.py b/backend/notification_v2/internal_api_views.py index be2bbe57e4..04d77b81cb 100644 --- a/backend/notification_v2/internal_api_views.py +++ b/backend/notification_v2/internal_api_views.py @@ -14,6 +14,7 @@ from datetime import timedelta from typing import Any, cast +from account_v2.models import Organization from api_v2.models import APIDeployment from django.conf import settings from django.db import transaction @@ -37,6 +38,10 @@ webhook_url_hash, ) from notification_v2.models import Notification, NotificationBuffer +from notification_v2.notification_dispatch import ( + PermanentDispatchError, + dispatch_webhook_notification, +) logger = logging.getLogger(__name__) @@ -513,6 +518,38 @@ def _reclaim_stale_sending() -> int: return int(reclaimed) +def _org_identifier(org_pk: int) -> str | None: + """Resolve the string ``Organization.organization_id`` from the buffer's org pk. + + ``resolve_transport`` keys its Flipt decision on the org's string identifier, + but the buffer stores/uses the Organization pk. One indexed pk lookup per + dispatch group (post-commit) — negligible relative to the downstream webhook + dispatch. + + Data-anomaly guard: ``NotificationBuffer.organization`` is + ``on_delete=CASCADE``, so a live buffer row with a missing org is unreachable + in normal operation. A ``None`` here therefore signals a dangling FK — we log + it (the only org-traceable breadcrumb; resolve_transport's own warning is keyed + on the random dispatch uuid) and fail closed to Celery in resolve_transport. + """ + org_string_id = ( + Organization.objects.filter(pk=org_pk) + .values_list("organization_id", flat=True) + .first() + ) + if org_string_id is None: + # Sentry-routed (logger.error): a live buffer row with no org is a data + # anomaly (dangling FK / corruption) that shouldn't happen under the + # CASCADE constraint, not routine noise. Routing still fails closed to + # Celery in resolve_transport. + logger.error( + "metric=notification_org_identifier_missing_total org_pk=%s " + "(dangling FK; notification routing falls back to Celery)", + org_pk, + ) + return org_string_id + + def _send_clubbed( *, url: str, @@ -521,7 +558,7 @@ def _send_clubbed( platform: str, max_retries: int, buffer_ids: list[str], - org_id: Any, + org_id: int, ) -> None: """Send the clubbed Celery task after the DB transition has committed. @@ -540,8 +577,12 @@ def _send_clubbed( ``buffer_row_ids`` + ``organization_id`` to the worker so it can mark them. """ try: - celery_app.send_task( - "send_webhook_notification", + # Flag-gated transport (UN-3753): PG queue when pg_queue_enabled for this + # org, else Celery (byte-identical to the prior send_task). resolve_transport + # keys on the org STRING id, but the buffer/worker contract below uses the + # org pk — hence _org_identifier(org_id) for routing, org_id in kwargs. + dispatch_webhook_notification( + celery_app=celery_app, args=[url, body, headers, settings.NOTIFICATION_TIMEOUT], kwargs={ "max_retries": max_retries, @@ -557,6 +598,7 @@ def _send_clubbed( "organization_id": org_id, }, queue="notifications", + org_string_id=_org_identifier(org_id), ) logger.info( "metric=notification_batch_dispatched_total platform=%s result=success " @@ -566,7 +608,37 @@ def _send_clubbed( webhook_url_hash(url), len(buffer_ids), ) + except PermanentDispatchError: + # PG-ONLY permanent failure. From this path the only reachable cause is + # payload JSON-serialization (e.g. a NaN/Infinity float that jsonb rejects + # at insert): this call passes no reply_key/callback and a default in-range + # priority, so enqueue_task's other permanent checks can't fire here (the + # full set lives on PermanentDispatchError's docstring). It fails + # identically every flush tick — dead-letter now (distinct metric) instead + # of reverting to PENDING and re-rendering + re-dispatching + emitting a + # broker_failure traceback until the attempt cap. The Celery path never + # raises this, so the flag-off flow is unchanged: a Celery send_task failure + # is an ordinary Exception handled by the transient branch below, exactly as + # before UN-3753. Guard on SENDING so a row the worker already resolved + # isn't clobbered. + logger.exception( + "metric=notification_batch_dispatched_total platform=%s " + "result=dispatch_error org_id=%s webhook_url_hash=%s rows=%d", + platform, + org_id, + webhook_url_hash(url), + len(buffer_ids), + ) + NotificationBuffer.objects.filter( + id__in=buffer_ids, + status=BufferStatus.SENDING.value, + ).update(status=BufferStatus.DEAD_LETTER.value) except Exception: + # TRANSIENT transport/broker failure — revert to PENDING (outside the + # committed txn) so the next flush tick retries; refund the SENDING-claim + # attempt since nothing was queued or sent. Guard on SENDING so a row the + # worker already marked terminal (broker raised post-delivery) isn't + # resurrected into a duplicate. logger.exception( "metric=notification_batch_dispatched_total platform=%s " "result=broker_failure org_id=%s webhook_url_hash=%s rows=%d", @@ -575,10 +647,6 @@ def _send_clubbed( webhook_url_hash(url), len(buffer_ids), ) - # Revert to PENDING (outside the committed txn) so a transient broker - # outage retries next tick; refund the SENDING-claim attempt since nothing - # was queued or sent. Guard on SENDING so a row the worker already marked - # terminal (broker raised post-delivery) isn't resurrected into a duplicate. NotificationBuffer.objects.filter( id__in=buffer_ids, status=BufferStatus.SENDING.value, @@ -589,7 +657,7 @@ def _send_clubbed( ) -def _penalize_render_failure(buffer_ids: list[str], org_id: Any, platform: str) -> None: +def _penalize_render_failure(buffer_ids: list[str], org_id: int, platform: str) -> None: """Charge a dispatch attempt to a group whose payloads failed to render. The SENDING-claim increment never runs on a render failure, so count it here @@ -609,7 +677,7 @@ def _penalize_render_failure(buffer_ids: list[str], org_id: Any, platform: str) def _dispatch_group( - org_id: Any, + org_id: int, webhook_url: str, auth_sig: str, platform: str, diff --git a/backend/notification_v2/notification_dispatch.py b/backend/notification_v2/notification_dispatch.py new file mode 100644 index 0000000000..a7331a1618 --- /dev/null +++ b/backend/notification_v2/notification_dispatch.py @@ -0,0 +1,121 @@ +"""Transport-routed dispatch for buffered webhook notifications (UN-3753). + +Routes the ``send_webhook_notification`` task through the same +:func:`resolve_transport` flag as the execution path: the PG queue when +``pg_queue_enabled`` for this org, else Celery. **Fail-closed** — with the gate +off (the production default) it resolves to Celery, behaving exactly like the +prior unconditional ``celery_app.send_task`` (zero regression). On PG the task +is drained by the PG notification consumer on the ``notifications`` queue. + +The same ``args``/``kwargs``/``queue`` are forwarded on both paths; on PG they're +JSON-normalized by ``enqueue_task`` (UUIDs/datetimes → str), so the consumer sees +the same payload the Celery worker would. +""" + +from __future__ import annotations + +import logging +import uuid +from typing import Any + +from pg_queue.producer import enqueue_task +from workflow_manager.workflow_v2.transport import resolve_transport + +from unstract.core.data_models import is_pg_transport + +logger = logging.getLogger(__name__) + +# The fired task name — mirrors the Celery task registered by the notification +# worker; kept as a local constant so the backend doesn't import the workers pkg. +WEBHOOK_NOTIFICATION_TASK = "send_webhook_notification" + + +class PermanentDispatchError(Exception): + """A dispatch failure that would fail identically on every retry. + + Raised ONLY on the PG path, when ``enqueue_task`` rejects the message for a + permanent reason (priority range / reply_key+callback exclusivity validation, + or a payload that can't be JSON-serialized). The Celery path never raises it, + so a caller can dead-letter on this exception without altering the flag-off + (Celery) error flow — a Celery ``send_task`` failure stays an ordinary + ``Exception`` the caller's transient handler owns, exactly as before. + """ + + +def dispatch_webhook_notification( + *, + celery_app: Any, + args: list[Any], + kwargs: dict[str, Any], + queue: str, + org_string_id: str | None, +) -> str: + """Dispatch ``send_webhook_notification`` on the resolved transport. + + ``args``/``kwargs``/``queue`` are forwarded unchanged on both paths, so the + flag-off (Celery) path is byte-identical to the legacy ``send_task`` call. + + Args: + celery_app: Injected Celery app (the backend's ``celery_service.app``); + passed in rather than imported so this seam stays trivially testable. + args: Positional task args, forwarded verbatim. + kwargs: Keyword task args, forwarded verbatim. For the buffered path this + carries ``organization_id`` = the buffer's org **pk** (the worker's + buffer-mark contract) — deliberately a DIFFERENT identifier from the + ``org_string_id`` param below (the two must not be conflated). + queue: Target queue name, forwarded verbatim. + org_string_id: The org's **string** identifier + (``Organization.organization_id``), used solely for the Flipt + transport decision — NOT the org pk carried in ``kwargs``. ``None`` (or + empty) fails closed to Celery. + + Returns: + A task id string — the Celery ``AsyncResult`` id on the Celery path, or + the minted PG task id on the PG path. Returned for symmetry / any future + caller; the sole current caller (fire-and-forget) discards it. + """ + # A buffered notification is a single fire-and-forget task with no natural + # sticky entity, so mint a fresh id to drive Flipt's percentage bucketing and + # to serve as the PG task id. + dispatch_id = str(uuid.uuid4()) + # resolve_transport already normalizes falsy input to Celery, so pass the id + # straight through (no `or None` needed). + transport = resolve_transport( + execution_id=dispatch_id, + organization_id=org_string_id, + ) + # Use the shared is_pg_transport() — the single source for "what counts as PG + # transport" — rather than opening a second comparison site. + if is_pg_transport(transport): + try: + enqueue_task( + task_name=WEBHOOK_NOTIFICATION_TASK, + queue=queue, + args=args, + kwargs=kwargs, + # enqueue_task types org_id as str; None→"" here satisfies that + # type only, not runtime — enqueue_task itself re-coerces + # ``org_id or ""`` at insert (producer.py), so this has no runtime + # effect. + org_id=org_string_id or "", + task_id=dispatch_id, + ) + except (ValueError, TypeError) as exc: + # PG-only permanent failure (enqueue_task validation / JSON encode): + # re-raise as PermanentDispatchError so the caller dead-letters it. + # A transient PG error (DB down) is NOT wrapped — it propagates as an + # ordinary Exception into the caller's retry (revert-to-PENDING) path. + raise PermanentDispatchError(str(exc)) from exc + logger.info( + "Webhook notification enqueued on PG '%s' queue (task_id=%s)", + queue, + dispatch_id, + ) + return dispatch_id + result = celery_app.send_task( + WEBHOOK_NOTIFICATION_TASK, + args=args, + kwargs=kwargs, + queue=queue, + ) + return result.id diff --git a/backend/notification_v2/tests/__init__.py b/backend/notification_v2/tests/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/backend/notification_v2/tests/test_notification_dispatch.py b/backend/notification_v2/tests/test_notification_dispatch.py new file mode 100644 index 0000000000..e53514d5de --- /dev/null +++ b/backend/notification_v2/tests/test_notification_dispatch.py @@ -0,0 +1,136 @@ +"""Unit tests for the buffered-webhook transport routing (UN-3753). + +``resolve_transport`` + ``enqueue_task`` are patched on the module, so no Flipt / +DB is needed — these pin the routing contract: PG when the flag resolves PG, +Celery otherwise (fail-closed), with byte-identical args/kwargs/queue on both +paths. +""" + +from __future__ import annotations + +from unittest.mock import MagicMock, patch + +import pytest + +import notification_v2.notification_dispatch as nd + +_ARGS = ["https://hook.test", {"text": "hi"}, {"Content-Type": "application/json"}, 30] +_KWARGS = { + "max_retries": 3, + "retry_delay": 10, + "platform": "SLACK", + "raise_on_final_failure": True, + "buffer_row_ids": ["b1", "b2"], + "organization_id": 7, # the org pk (worker's buffer-mark contract) +} +_QUEUE = "notifications" + + +def _dispatch(celery_app, org_string_id="org_x"): + return nd.dispatch_webhook_notification( + celery_app=celery_app, + args=_ARGS, + kwargs=_KWARGS, + queue=_QUEUE, + org_string_id=org_string_id, + ) + + +class TestDispatchWebhookNotification: + def test_routes_to_pg_when_flag_resolves_pg(self): + celery = MagicMock() + with ( + patch.object(nd, "resolve_transport", return_value="pg_queue"), + patch.object(nd, "enqueue_task", return_value=42) as enqueue, + ): + task_id = _dispatch(celery) + enqueue.assert_called_once() + kwargs = enqueue.call_args.kwargs + assert kwargs["task_name"] == "send_webhook_notification" + assert kwargs["queue"] == _QUEUE + assert kwargs["args"] == _ARGS + assert kwargs["kwargs"] == _KWARGS + assert kwargs["org_id"] == "org_x" + # The minted PG task id is returned and threaded into the enqueue row. + assert kwargs["task_id"] == task_id + celery.send_task.assert_not_called() + + def test_routes_to_celery_when_flag_resolves_celery(self): + celery = MagicMock() + with ( + patch.object(nd, "resolve_transport", return_value="celery"), + patch.object(nd, "enqueue_task") as enqueue, + ): + result = _dispatch(celery) + celery.send_task.assert_called_once_with( + "send_webhook_notification", args=_ARGS, kwargs=_KWARGS, queue=_QUEUE + ) + enqueue.assert_not_called() + assert result is celery.send_task.return_value.id + + def test_pg_enqueue_failure_propagates_with_no_celery_fallback(self): + # No silent Celery fallback on a TRANSIENT PG failure — it propagates raw + # (not wrapped), so the caller (_send_clubbed) reverts rows to PENDING. + celery = MagicMock() + with ( + patch.object(nd, "resolve_transport", return_value="pg_queue"), + patch.object(nd, "enqueue_task", side_effect=RuntimeError("pg down")), + ): + with pytest.raises(RuntimeError, match="pg down"): + _dispatch(celery) + celery.send_task.assert_not_called() + + def test_pg_permanent_enqueue_error_is_wrapped(self): + # A PERMANENT enqueue error (ValueError/TypeError: validation / JSON encode) + # is re-raised as PermanentDispatchError so the caller dead-letters it — + # raised ONLY on the PG path, so the Celery flow is never affected. + celery = MagicMock() + for exc in (ValueError("priority out of range"), TypeError("not serializable")): + with ( + patch.object(nd, "resolve_transport", return_value="pg_queue"), + patch.object(nd, "enqueue_task", side_effect=exc), + ): + with pytest.raises(nd.PermanentDispatchError): + _dispatch(celery) + celery.send_task.assert_not_called() + + def test_none_org_fails_closed_to_celery(self): + # A missing org string (org deleted) must not route to PG — resolve_transport + # fails closed, and we pass organization_id=None straight through to it. + celery = MagicMock() + with ( + patch.object(nd, "resolve_transport", return_value="celery") as resolve, + patch.object(nd, "enqueue_task") as enqueue, + ): + _dispatch(celery, org_string_id=None) + assert resolve.call_args.kwargs["organization_id"] is None + celery.send_task.assert_called_once() + enqueue.assert_not_called() + + def test_resolve_transport_buckets_by_minted_dispatch_id(self): + # Fire-and-forget: entity_id is a freshly minted uuid (str), and it equals + # the PG task_id so the row and the Flipt bucket agree. + celery = MagicMock() + with ( + patch.object(nd, "resolve_transport", return_value="pg_queue") as resolve, + patch.object(nd, "enqueue_task", return_value=1) as enqueue, + ): + task_id = _dispatch(celery) + assert resolve.call_args.kwargs["execution_id"] == task_id + assert enqueue.call_args.kwargs["task_id"] == task_id + + def test_args_and_kwargs_identical_on_both_paths(self): + # The consumer must behave the same regardless of transport. + celery = MagicMock() + with patch.object(nd, "resolve_transport", return_value="celery"): + _dispatch(celery) + celery_call = celery.send_task.call_args.kwargs + celery2 = MagicMock() + with ( + patch.object(nd, "resolve_transport", return_value="pg_queue"), + patch.object(nd, "enqueue_task", return_value=1) as enqueue, + ): + _dispatch(celery2) + assert enqueue.call_args.kwargs["args"] == celery_call["args"] + assert enqueue.call_args.kwargs["kwargs"] == celery_call["kwargs"] + assert enqueue.call_args.kwargs["queue"] == celery_call["queue"] diff --git a/backend/notification_v2/tests/test_send_clubbed.py b/backend/notification_v2/tests/test_send_clubbed.py new file mode 100644 index 0000000000..4ebdc6a33c --- /dev/null +++ b/backend/notification_v2/tests/test_send_clubbed.py @@ -0,0 +1,117 @@ +"""Call-site tests for ``_send_clubbed`` / ``_org_identifier`` (UN-3753). + +The seam's own routing is covered by ``test_notification_dispatch``; these lock +the two highest-risk regressions at the buffer-flush call site: + +1. the two-org-identifier contract — the org **string** id routes the flag, while + the org **pk** stays in the worker kwargs (swapping them passes every seam test + yet mis-routes every org to Celery and strands the mark endpoint); +2. failure recovery — a transient/broker error refunds and reverts SENDING rows to + PENDING, while a permanent PG enqueue error (surfaced as ``PermanentDispatchError``, + raised only on the PG path) dead-letters instead of retrying forever. A Celery + send_task failure is an ordinary ``Exception`` → the PENDING path, so flag-off is + byte-identical. + +Mock-based (no broker/DB): patch the module's collaborators. +""" + +from __future__ import annotations + +from unittest.mock import patch + +from notification_v2 import internal_api_views as views +from notification_v2.enums import BufferStatus + + +def _send(**overrides): + kw = { + "url": "https://hook.test", + "body": {"text": "x"}, + "headers": {"Content-Type": "application/json"}, + "platform": "SLACK", + "max_retries": 3, + "buffer_ids": ["b1"], + "org_id": 7, + } + kw.update(overrides) + views._send_clubbed(**kw) + + +class TestSendClubbedOrgContract: + def test_routes_string_id_and_keeps_pk_in_kwargs(self): + with ( + patch.object(views, "_org_identifier", return_value="org-str-id") as ident, + patch.object(views, "dispatch_webhook_notification") as seam, + ): + _send(org_id=7) + ident.assert_called_once_with(7) + call = seam.call_args.kwargs + # Routing uses the STRING id; the worker buffer-mark contract keeps the PK. + assert call["org_string_id"] == "org-str-id" + assert call["kwargs"]["organization_id"] == 7 + assert call["queue"] == "notifications" + + +class TestSendClubbedFailureRecovery: + def test_transient_failure_reverts_sending_rows_to_pending(self): + with ( + patch.object(views, "_org_identifier", return_value="o"), + patch.object( + views, + "dispatch_webhook_notification", + side_effect=RuntimeError("broker down"), + ), + patch.object(views, "NotificationBuffer") as buf, + ): + _send(buffer_ids=["b1", "b2"]) + # Guarded on SENDING, reverted to PENDING (attempt refunded). + fkw = buf.objects.filter.call_args.kwargs + assert fkw["status"] == BufferStatus.SENDING.value + assert set(fkw["id__in"]) == {"b1", "b2"} + ukw = buf.objects.filter.return_value.update.call_args.kwargs + assert ukw["status"] == BufferStatus.PENDING.value + assert ukw["dispatched_at"] is None + + def test_permanent_pg_error_dead_letters(self): + # The PG path surfaces a permanent enqueue failure as PermanentDispatchError; + # only that dead-letters. (A raw ValueError from the Celery path can't occur, + # and would fall through to the transient PENDING branch, unchanged.) + with ( + patch.object(views, "_org_identifier", return_value="o"), + patch.object( + views, + "dispatch_webhook_notification", + side_effect=views.PermanentDispatchError("priority out of range"), + ), + patch.object(views, "NotificationBuffer") as buf, + ): + _send() + # Permanent error → terminal DEAD_LETTER, no PENDING revert / refund. + # The dead-letter update is guarded to only touch rows still SENDING + # (same clobber-guard the transient-revert path asserts). + fkw = buf.objects.filter.call_args.kwargs + assert fkw["status"] == BufferStatus.SENDING.value + ukw = buf.objects.filter.return_value.update.call_args.kwargs + assert ukw == {"status": BufferStatus.DEAD_LETTER.value} + + +class TestOrgIdentifier: + def test_returns_string_id(self): + with patch.object(views, "Organization") as org: + chain = org.objects.filter.return_value.values_list.return_value + chain.first.return_value = "org-uuid" + assert views._org_identifier(7) == "org-uuid" + org.objects.filter.assert_called_once_with(pk=7) + + def test_missing_org_returns_none_and_logs_error(self): + with ( + patch.object(views, "Organization") as org, + patch.object(views.logger, "error") as err, + ): + chain = org.objects.filter.return_value.values_list.return_value + chain.first.return_value = None + assert views._org_identifier(7) is None + # A dangling FK is a data anomaly: logged at error (Sentry-routed), + # org-traceable, rather than swallowed. + err.assert_called_once() + assert "org_pk" in err.call_args.args[0] diff --git a/backend/pg_queue/producer.py b/backend/pg_queue/producer.py index 3c4585fc4d..028fddfb6d 100644 --- a/backend/pg_queue/producer.py +++ b/backend/pg_queue/producer.py @@ -50,8 +50,15 @@ def _json_safe(value: Any) -> Any: ``PgQueueMessage.message`` is a plain ``JSONField`` (no Django encoder), and the worker consumer already receives string ids on the existing PG dispatch path, so coercing here keeps both transports consistent. + + ``allow_nan=False`` rejects ``NaN``/``Infinity`` here with a ``ValueError`` + rather than letting Python's default lenient encoder emit the non-standard + ``NaN``/``Infinity`` tokens: Postgres ``jsonb`` rejects those at insert with a + ``django.db.DataError`` (a permanent failure the notification dispatcher's + ``(ValueError, TypeError)`` seam would otherwise miss, looping the row). Fail + at the intended seam instead. """ - return json.loads(json.dumps(value, default=str)) + return json.loads(json.dumps(value, default=str, allow_nan=False)) def enqueue_task( @@ -97,33 +104,37 @@ def enqueue_task( f"[{FAIRNESS_MIN_PRIORITY}, {FAIRNESS_MAX_PRIORITY}]: {priority!r}" ) pg_queue = queue or DEFAULT_GENERAL_QUEUE - message: TaskPayload = { - "task_name": task_name, - "args": _json_safe(list(args) if args is not None else []), - "kwargs": _json_safe(dict(kwargs) if kwargs is not None else {}), - "queue": pg_queue, - # Coerce like args/kwargs/on_success/on_error: FairnessPayload may carry a - # UUID/enum/datetime that a plain JSONField insert can't serialise, which - # would raise at enqueue and drop the task on the PG path only. - "fairness": _json_safe(fairness) if fairness is not None else fairness, - } - # Each optional key is set only when present — keeps fire-and-forget rows - # byte-identical to before these fields existed. - if reply_key is not None: - message["reply_key"] = reply_key - # Continuation specs carry a nested callback ``kwargs`` dict that may hold a - # UUID/datetime — coerce like ``args``/``kwargs`` above, else the JSONField - # insert raises at dispatch time (caller-visible). - if on_success is not None: - message["on_success"] = _json_safe(on_success) - if on_error is not None: - message["on_error"] = _json_safe(on_error) - if task_id is not None: - message["task_id"] = task_id # Mirror the worker _enqueue_pg path: log the failure with breadcrumbs before # it propagates, so a DB/constraint/serialization error isn't mislabeled by - # the caller's broad handler. + # the caller's broad handler. Message construction (the _json_safe coercions) + # lives inside the try because serialization failures surface there — e.g. a + # NaN/Infinity float raises ValueError under allow_nan=False — and on the + # orchestrator path (no silent fallback) would otherwise drop the task with no + # task/queue/org breadcrumb for on-call. try: + message: TaskPayload = { + "task_name": task_name, + "args": _json_safe(list(args) if args is not None else []), + "kwargs": _json_safe(dict(kwargs) if kwargs is not None else {}), + "queue": pg_queue, + # Coerce like args/kwargs/on_success/on_error: FairnessPayload may carry + # a UUID/enum/datetime that a plain JSONField insert can't serialise, + # which would raise at enqueue and drop the task on the PG path only. + "fairness": _json_safe(fairness) if fairness is not None else fairness, + } + # Each optional key is set only when present — keeps fire-and-forget rows + # byte-identical to before these fields existed. + if reply_key is not None: + message["reply_key"] = reply_key + # Continuation specs carry a nested callback ``kwargs`` dict that may hold a + # UUID/datetime — coerce like ``args``/``kwargs`` above, else the JSONField + # insert raises at dispatch time (caller-visible). + if on_success is not None: + message["on_success"] = _json_safe(on_success) + if on_error is not None: + message["on_error"] = _json_safe(on_error) + if task_id is not None: + message["task_id"] = task_id row = PgQueueMessage.objects.create( queue_name=pg_queue, message=message, diff --git a/backend/pg_queue/tests/test_producer.py b/backend/pg_queue/tests/test_producer.py index 40df8a79fe..37cdb89b49 100644 --- a/backend/pg_queue/tests/test_producer.py +++ b/backend/pg_queue/tests/test_producer.py @@ -106,6 +106,20 @@ def test_json_safe_coerces_datetime(self): when = model.objects.create.call_args.kwargs["message"]["kwargs"]["when"] assert isinstance(when, str) and "2026-06-18" in when + def test_json_safe_rejects_nan_with_value_error(self): + # NaN would slip past the default lenient encoder and only fail at the + # jsonb insert (DataError); allow_nan=False surfaces it as a ValueError at + # the enqueue seam so the notification dispatcher can dead-letter it. + # The coercion runs inside the try, so the failure is logged with the + # task/queue/org breadcrumb (not dropped silently) and never reaches the + # DB insert. + nan_kwargs = {"score": float("nan")} + with patch(_MODEL) as model, patch.object(producer, "logger") as log: + with pytest.raises(ValueError): + producer.enqueue_task(task_name="t", queue="celery", kwargs=nan_kwargs) + model.objects.create.assert_not_called() + log.exception.assert_called_once() + def test_enqueue_failure_logs_and_propagates(self): with patch(_MODEL) as model: model.objects.create.side_effect = RuntimeError("db down") diff --git a/workers/shared/enums/worker_enums_base.py b/workers/shared/enums/worker_enums_base.py index 9c4dfbdb46..97e4799768 100644 --- a/workers/shared/enums/worker_enums_base.py +++ b/workers/shared/enums/worker_enums_base.py @@ -58,13 +58,22 @@ def to_import_path(self) -> str: if self.is_pluggable(): return f"pluggable_worker.{self.value}.tasks" - # Map to actual directory structure + return f"{self.to_directory()}.tasks" + + def to_directory(self) -> str: + """Return the on-disk directory name for this (non-pluggable) worker. + + Single source of truth for the naming-convention mapping: enum values use + underscores (Python module names), but a few on-disk dirs use hyphens + (e.g. ``api-deployment``). ``to_import_path`` builds on this, and the + file-path task loader in ``worker.py`` reads it directly rather than + slicing the import path. + """ directory_mapping = { - "api_deployment": "api-deployment", + WorkerType.API_DEPLOYMENT: "api-deployment", # All others use same name for directory and module } - directory = directory_mapping.get(self.value, self.value) - return f"{directory}.tasks" + return directory_mapping.get(self, self.value) def is_pluggable(self) -> bool: """Check if this worker type is a pluggable worker. diff --git a/workers/tests/test_worker_enums_directory.py b/workers/tests/test_worker_enums_directory.py new file mode 100644 index 0000000000..931fc45f1e --- /dev/null +++ b/workers/tests/test_worker_enums_directory.py @@ -0,0 +1,30 @@ +"""WorkerType.to_directory() — the single source for the on-disk dir mapping (UN-3798). + +worker.py's file-path task loader reads to_directory() directly instead of slicing +to_import_path(); these pin the hyphen/underscore mapping and that to_import_path +is built on top of it, so the two can't drift. +""" + +from celery import Celery +from shared.enums.worker_enums_base import WorkerType + +# The workers autouse conftest fixture finalizes celery's default_app around every +# test; this pure-enum test builds no Celery app of its own, so establish one. +Celery("test-worker-enums").set_default() + + +def test_to_directory_maps_underscored_value_to_hyphenated_dir(): + # The one worker whose on-disk dir differs from its enum value. + assert WorkerType.API_DEPLOYMENT.value == "api_deployment" + assert WorkerType.API_DEPLOYMENT.to_directory() == "api-deployment" + + +def test_to_directory_passthrough_when_dir_equals_value(): + assert WorkerType.GENERAL.to_directory() == "general" + + +def test_to_import_path_is_built_on_to_directory(): + # to_import_path must derive from to_directory (not a separate mapping), so the + # directory naming lives in exactly one place. + wt = WorkerType.API_DEPLOYMENT + assert wt.to_import_path() == f"{wt.to_directory()}.tasks" diff --git a/workers/worker.py b/workers/worker.py index 3822e81a16..295f5c9b00 100755 --- a/workers/worker.py +++ b/workers/worker.py @@ -5,6 +5,7 @@ It uses WorkerBuilder to ensure proper configuration including chord retry settings. """ +import importlib.util import logging import os import sys @@ -440,40 +441,83 @@ def on_task_postrun(sender=None, task_id=None, **kwargs): initialize_worker_infrastructure() logger.info("✅ Worker infrastructure initialized successfully") -# Import tasks from the worker-specific directory -# Determine worker path dynamically based on worker type -base_dir = os.path.dirname(os.path.abspath(__file__)) -if worker_type.is_pluggable(): - # Pluggable workers live inside workers/pluggable_worker/{worker_name} - worker_directory = os.path.join("pluggable_worker", worker_type.value) - worker_path = os.path.join(base_dir, worker_directory) -else: - # Enum values use underscores (Python module names); a few on-disk dirs - # still use hyphens (e.g. api-deployment). Derive the directory from the - # authoritative import-path map on WorkerType instead of a blind replace. - worker_directory = worker_type.to_import_path().rsplit(".", 1)[0] + +def load_worker_tasks(worker_type: WorkerType) -> None: + """Register the worker type's Celery tasks. + + For pluggable workers, ``build_celery_app()`` above already imported the + plugin package via ``WorkerBuilder._verify_pluggable_worker_exists`` + (``importlib.import_module("pluggable_worker..worker")`` — a proper + PACKAGE import that RUNS the plugin's own registration code, e.g. + ``from . import tasks`` with relative imports such as + ``from .clients import ...``). That registration's ``@shared_task`` bindings + complete when the app FINALIZES, which can be *after* this point — so the + tasks are not necessarily bound to the app yet here. The generic file-path + load below is SKIPPED for pluggable workers not because they are already + bound, but because loading ``tasks.py`` under a bare ``"tasks"`` spec (no + parent package) would break the plugin's relative imports ("attempted + relative import with no known parent package"). + + Non-pluggable (top-level) workers use absolute imports; their ``tasks.py`` is + loaded by file path with the worker directory on ``sys.path``. Their tasks + bind eagerly here — there is no later finalize step to rescue them — so a + missing directory / ``tasks.py`` is a broken deploy and hard-fails rather + than booting a no-op worker. + """ + if worker_type.is_pluggable(): + logger.info( + f"✅ Pluggable worker {worker_type.value} tasks registered via " + "WorkerBuilder (package import); skipping file-path task load" + ) + return + + base_dir = os.path.dirname(os.path.abspath(__file__)) + worker_directory = worker_type.to_directory() worker_path = os.path.join(base_dir, worker_directory) + if not os.path.exists(worker_path): + # Non-pluggable only (pluggable returned above): tasks bind eagerly here + # with no later finalize step, so a missing worker dir is terminal. + raise RuntimeError(f"Worker directory not found: {worker_path}") -# Add worker directory to path for task imports -if os.path.exists(worker_path): sys.path.append(worker_path) logger.info(f"✅ Added {worker_directory} to Python path for task imports") - # Import tasks module to register tasks tasks_file = os.path.join(worker_path, "tasks.py") - if os.path.exists(tasks_file): - logger.info(f"📋 Loading tasks from: {tasks_file}") - # Import the tasks module to register tasks with the app - import importlib.util - - spec = importlib.util.spec_from_file_location("tasks", tasks_file) - tasks_module = importlib.util.module_from_spec(spec) - spec.loader.exec_module(tasks_module) - logger.info(f"✅ Tasks loaded successfully from {worker_directory}") + if not os.path.exists(tasks_file): + raise RuntimeError(f"No tasks.py found for worker at: {tasks_file}") + + logger.info(f"📋 Loading tasks from: {tasks_file}") + spec = importlib.util.spec_from_file_location("tasks", tasks_file) + tasks_module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(tasks_module) + logger.info(f"✅ Tasks loaded successfully from {worker_directory}") + + +load_worker_tasks(worker_type) + +# A worker that boots with no registered tasks starts but silently processes +# nothing (Celery does not error on an empty registry). Surface that misconfig, +# but split by worker kind: pluggable tasks bind on app finalize (@shared_task / +# connect_on_app_finalize), which can be after this point, so a raise here would +# false-positive on a correctly-configured pluggable worker — only WARN. A +# non-pluggable worker has no later binding step (load_worker_tasks bound its +# tasks eagerly above), so an empty registry here is terminal — RAISE. +if not any(not name.startswith("celery.") for name in app.tasks): + _empty_registry_msg = ( + f"No non-celery tasks registered for worker '{worker_type.value}' " + f"(pluggable={worker_type.is_pluggable()})." + ) + if worker_type.is_pluggable(): + logger.warning( + f"⚠️ {_empty_registry_msg} If this persists past app finalize the " + "worker will start but process nothing — check the plugin task " + "registration." + ) else: - logger.warning(f"⚠️ No tasks.py found at: {tasks_file}") -else: - logger.error(f"❌ Worker directory not found: {worker_path}") + raise RuntimeError( + f"{_empty_registry_msg} The worker would start but process nothing — " + "check the worker task registration." + ) # Log successful configuration logger.info(f"✅ Successfully loaded {worker_type} worker using WorkerBuilder")