UN-3445 [GATED-FEAT] Complete the PostgreSQL queue transport: scheduler, metrics, log streaming and strand recovery

backend/backend/internal_base_urls.py

@@ -191,6 +191,13 @@ def test_middleware_debug(request):
         include("workflow_manager.workflow_execution_internal_urls"),
         name="workflow_execution_internal",
     ),
+    # Dashboard-metrics periodics (UN-3796) — called by the thin worker tasks the PG
+    # scheduler fires, replacing Beat + workerMetrics.
+    path(
+        "v1/dashboard-metrics/",
+        include("dashboard_metrics.internal_urls"),
+        name="dashboard_metrics_internal",
+    ),
     # Workflow management and pipeline APIs
     path(
         "v1/workflow-manager/",

backend/dashboard_metrics/internal_urls.py

@@ -0,0 +1,29 @@
+"""Internal API URLs for the dashboard-metrics periodics (UN-3796).
+
+Called by the thin worker tasks that the PG scheduler fires, replacing Beat +
+``workerMetrics``. Mirrors the shape of ``execution_log_internal_urls``.
+"""
+
+from django.urls import path
+
+from . import internal_views
+
+app_name = "dashboard_metrics_internal"
+
+urlpatterns = [
+    path(
+        "aggregate/",
+        internal_views.AggregateMetricsAPIView.as_view(),
+        name="aggregate_metrics",
+    ),
+    path(
+        "cleanup/hourly/",
+        internal_views.CleanupHourlyMetricsAPIView.as_view(),
+        name="cleanup_hourly_metrics",
+    ),
+    path(
+        "cleanup/daily/",
+        internal_views.CleanupDailyMetricsAPIView.as_view(),
+        name="cleanup_daily_metrics",
+    ),
+]

backend/dashboard_metrics/internal_views.py

@@ -0,0 +1,115 @@
+"""Internal API for running the dashboard-metrics periodics (UN-3796).
+
+Beat + ``workerMetrics`` are the last two Celery deployments blocking the
+"scale every Celery deployment to zero" gate. The PG scheduler replaces Beat, but the
+three ``dashboard_metrics.*`` tasks are **Django** — ORM, cache, a Redis lock — and the
+PG consumer bootstraps its tasks from ``workers/`` with no Django, so it cannot run them
+directly.
+
+These endpoints are the execution half, following the pattern already used by
+``ProcessLogHistoryAPIView`` / ``process_notification_buffer``: a thin worker-side task
+POSTs here, and the backend runs the real function. The task bodies are plain functions
+that happen to carry ``@shared_task``, so they are called **verbatim** — the aggregation
+logic, its windows and its Redis lock are reused unchanged, not reimplemented.
+
+Auth is entirely ``InternalAPIAuthMiddleware``, which fires on any ``/internal/`` path.
+
+**Org context is deliberately absent.** The Celery path runs with no organization set —
+which is why ``tasks.py`` uses ``_base_manager`` throughout — and the middleware
+populates ``StateStore`` from ``X-Organization-ID``. The callers must not send that
+header, and each view clears the slot defensively: ``StateStore`` is a thread-local and
+gunicorn reuses threads, so a value left behind by an earlier request on the same thread
+would silently scope these global aggregations to one tenant.
+"""
+
+import contextlib
+import logging
+from typing import Any
+
+from rest_framework import status
+from rest_framework.request import Request
+from rest_framework.response import Response
+from rest_framework.views import APIView
+from utils.constants import Account
+from utils.local_context import StateStore
+
+from dashboard_metrics.tasks import (
+    aggregate_metrics_from_sources,
+    cleanup_daily_metrics,
+    cleanup_hourly_metrics,
+)
+
+logger = logging.getLogger(__name__)
+
+# Mirrors the defaults the Beat rows carry in their kwargs
+# (dashboard_metrics/migrations/0002_setup_periodic_tasks.py), so a caller that omits
+# them gets the same retention the Celery path applies.
+DEFAULT_HOURLY_RETENTION_DAYS = 30
+DEFAULT_DAILY_RETENTION_DAYS = 365
+
+
+def _clear_org_context() -> None:
+    """Drop any organization left in this thread's StateStore. See the module docstring.
+
+    ``suppress`` because ``StateStore.clear`` raises when the slot was never set, which
+    is the normal case and not an error.
+    """
+    with contextlib.suppress(Exception):
+        StateStore.clear(Account.ORGANIZATION_ID)
+
+
+def _int_arg(request: Request, key: str, default: int) -> int:
+    """Read an optional positive integer from the request body."""
+    raw = request.data.get(key, default) if isinstance(request.data, dict) else default
+    try:
+        value = int(raw)
+    except (TypeError, ValueError) as exc:
+        raise ValueError(f"{key} must be an integer, got {raw!r}") from exc
+    if value < 1:
+        raise ValueError(f"{key} must be >= 1, got {value}")
+    return value
+
+
+class _MetricsTaskAPIView(APIView):
+    """Shared plumbing: clear org context, run, translate errors."""
+
+    def _run(self, fn, *args: Any, **kwargs: Any) -> Response:
+        _clear_org_context()
+        try:
+            return Response(fn(*args, **kwargs))
+        except ValueError as exc:  # bad request body
+            return Response({"error": str(exc)}, status=status.HTTP_400_BAD_REQUEST)
+        except Exception as exc:
+            logger.error("dashboard-metrics internal call failed: %s", exc, exc_info=True)
+            return Response(
+                {"error": str(exc)}, status=status.HTTP_500_INTERNAL_SERVER_ERROR
+            )
+
+
+class AggregateMetricsAPIView(_MetricsTaskAPIView):
+    """Run the metrics aggregation.
+
+    Calls the Celery task body verbatim, Redis lock included — this endpoint exists
+    only because the PG consumer has no Django, not to change what the job does.
+    """
+
+    def post(self, request: Request) -> Response:
+        return self._run(aggregate_metrics_from_sources)
+
+
+class CleanupHourlyMetricsAPIView(_MetricsTaskAPIView):
+    def post(self, request: Request) -> Response:
+        try:
+            days = _int_arg(request, "retention_days", DEFAULT_HOURLY_RETENTION_DAYS)
+        except ValueError as exc:
+            return Response({"error": str(exc)}, status=status.HTTP_400_BAD_REQUEST)
+        return self._run(cleanup_hourly_metrics, retention_days=days)
+
+
+class CleanupDailyMetricsAPIView(_MetricsTaskAPIView):
+    def post(self, request: Request) -> Response:
+        try:
+            days = _int_arg(request, "retention_days", DEFAULT_DAILY_RETENTION_DAYS)
+        except ValueError as exc:
+            return Response({"error": str(exc)}, status=status.HTTP_400_BAD_REQUEST)
+        return self._run(cleanup_daily_metrics, retention_days=days)

backend/dashboard_metrics/migrations/0004_pg_periodic_tasks.py

@@ -0,0 +1,97 @@
+"""Declare the dashboard-metrics periodics for the PG scheduler (UN-3796).
+
+The PG twin of ``0002_setup_periodic_tasks``, which declares the same three schedules
+for Celery Beat. **Deliberately in this app rather than in ``pg_queue``**: the failure
+mode that matters is the two declarations drifting apart, and a reviewer editing a
+schedule sees both only if they sit side by side.
+
+Why a data migration here, when the pipeline mirror is a management command: these are
+three fixed rows known at build time, so a migration is the right tool — tiny, idempotent
+(``update_or_create``), and it reaches every environment including on-prem with no
+operator step. Pipeline schedules are bulk and per-environment, so they stay a chunked
+command that runs outside the migrate transaction.
+
+Rows land **inert**: ``pg_owned=False`` and ``next_run_at=NULL``. Nothing fires from this
+migration — the PG scheduler skips rows it does not own, and a NULL ``next_run_at`` means
+"record a baseline next tick" rather than "overdue, fire now", so enabling the flag never
+produces a burst of catch-up runs.
+
+``task_kwargs`` is stored **decoded**: Beat keeps ``kwargs`` as a JSON *string*
+(``'{"retention_days": 30}'``) while ``PgPeriodicTask.task_kwargs`` is a JSONField, so the
+dispatcher can build a payload without re-parsing per tick.
+"""
+
+from django.db import migrations
+
+# Single source for both directions, and importable by the drift test. Mirrors the Beat
+# declarations in 0002_setup_periodic_tasks one-for-one — same names (the mirror key is
+# PeriodicTask.name), same queue, same kwargs, and cron strings equivalent to the
+# IntervalSchedule/CrontabSchedule rows there.
+PG_PERIODIC_TASKS = [
+    {
+        "name": "dashboard_metrics_aggregate_from_sources",
+        "task_name": "dashboard_metrics.aggregate_from_sources",
+        "queue": "dashboard_metric_events",
+        "task_args": [],
+        "task_kwargs": {},
+        # Beat: IntervalSchedule(every=15, period="minutes")
+        "cron_string": "*/15 * * * *",
+    },
+    {
+        "name": "dashboard_metrics_cleanup_hourly",
+        "task_name": "dashboard_metrics.cleanup_hourly_data",
+        "queue": "dashboard_metric_events",
+        "task_args": [],
+        "task_kwargs": {"retention_days": 30},
+        # Beat: CrontabSchedule(minute=0, hour=2, every day) UTC
+        "cron_string": "0 2 * * *",
+    },
+    {
+        "name": "dashboard_metrics_cleanup_daily",
+        "task_name": "dashboard_metrics.cleanup_daily_data",
+        "queue": "dashboard_metric_events",
+        "task_args": [],
+        "task_kwargs": {"retention_days": 365},
+        # Beat: CrontabSchedule(minute=0, hour=3, day_of_week=0 → Sunday) UTC
+        "cron_string": "0 3 * * 0",
+    },
+]
+
+
+def create_pg_periodic_tasks(apps, schema_editor):
+    PgPeriodicTask = apps.get_model("pg_queue", "PgPeriodicTask")
+    for spec in PG_PERIODIC_TASKS:
+        PgPeriodicTask.objects.update_or_create(
+            name=spec["name"],
+            defaults={
+                "task_name": spec["task_name"],
+                "queue": spec["queue"],
+                "task_args": spec["task_args"],
+                "task_kwargs": spec["task_kwargs"],
+                "cron_string": spec["cron_string"],
+                "org_id": "",
+                "enabled": True,
+                # Inert until the rollout flag decides otherwise; never fired by
+                # applying this migration.
+                "pg_owned": False,
+            },
+        )
+
+
+def remove_pg_periodic_tasks(apps, schema_editor):
+    PgPeriodicTask = apps.get_model("pg_queue", "PgPeriodicTask")
+    PgPeriodicTask.objects.filter(
+        name__in=[spec["name"] for spec in PG_PERIODIC_TASKS]
+    ).delete()
+
+
+class Migration(migrations.Migration):
+    dependencies = [
+        ("dashboard_metrics", "0003_alter_eventmetricsdaily_organization_and_more"),
+        # The table this seeds.
+        ("pg_queue", "0003_pgperiodictask"),
+    ]
+
+    operations = [
+        migrations.RunPython(create_pg_periodic_tasks, remove_pg_periodic_tasks),
+    ]

backend/dashboard_metrics/tests/test_pg_periodic_task_declarations.py

@@ -0,0 +1,113 @@
+"""Drift guard between the Beat and PG declarations of the metrics periodics (UN-3796).
+
+Two migrations declare the same three schedules — ``0002_setup_periodic_tasks`` for Celery
+Beat and ``0004_pg_periodic_tasks`` for the PG scheduler. They are separate rows in
+separate tables, so nothing stops someone editing one and forgetting the other. That is
+the whole failure mode this file exists for: a schedule changed on Beat but not on PG means
+the task silently runs on a different cadence the moment the flag flips.
+
+DB-free — both migration modules are imported and their declared specs compared directly,
+so this runs in the unit tier rather than needing a migrated database.
+"""
+
+from __future__ import annotations
+
+import importlib
+import json
+
+import pytest
+
+_BEAT_MIGRATION = "dashboard_metrics.migrations.0002_setup_periodic_tasks"
+_PG_MIGRATION = "dashboard_metrics.migrations.0004_pg_periodic_tasks"
+
+# Cron equivalent of each Beat schedule, asserted against what the Beat migration builds.
+# Written out rather than derived: deriving it from the same code under test would make
+# the comparison vacuous.
+_EXPECTED_CRON = {
+    "dashboard_metrics_aggregate_from_sources": "*/15 * * * *",
+    "dashboard_metrics_cleanup_hourly": "0 2 * * *",
+    "dashboard_metrics_cleanup_daily": "0 3 * * 0",
+}
+
+
+@pytest.fixture(scope="module")
+def pg_specs() -> dict[str, dict]:
+    mod = importlib.import_module(_PG_MIGRATION)
+    return {spec["name"]: spec for spec in mod.PG_PERIODIC_TASKS}
+
+
+class _FakeQuerySet:
+    """Captures update_or_create calls from the Beat migration without a database."""
+
+    def __init__(self, sink: dict):
+        self._sink = sink
+
+    def get_or_create(self, **kwargs):
+        # Schedule rows (Interval/Crontab) — return the kwargs so the PeriodicTask
+        # call can be inspected for which schedule it was given.
+        return kwargs, True
+
+    def update_or_create(self, name=None, defaults=None, **_kw):
+        self._sink[name] = defaults or {}
+        return defaults, True
+
+    def filter(self, *_a, **_k):
+        return self
+
+    def delete(self):
+        return (0, {})
+
+
+@pytest.fixture(scope="module")
+def beat_specs() -> dict[str, dict]:
+    """Run the Beat migration's forward function against fakes and capture what it declares."""
+    mod = importlib.import_module(_BEAT_MIGRATION)
+    captured: dict[str, dict] = {}
+
+    class _Apps:
+        def get_model(self, _app, model):
+            if model == "PeriodicTask":
+                return type("PT", (), {"objects": _FakeQuerySet(captured)})
+            return type("S", (), {"objects": _FakeQuerySet({})})
+
+    mod.create_periodic_tasks(_Apps(), None)
+    return captured
+
+
+class TestDeclarationsAgree:
+    def test_same_set_of_schedules(self, beat_specs, pg_specs):
+        # A schedule added to Beat but not PG stops firing the moment the flag flips;
+        # the reverse fires something Beat never knew about.
+        assert set(beat_specs) == set(pg_specs)
+
+    @pytest.mark.parametrize("name", sorted(_EXPECTED_CRON))
+    def test_task_path_and_queue_match(self, beat_specs, pg_specs, name):
+        assert pg_specs[name]["task_name"] == beat_specs[name]["task"]
+        assert pg_specs[name]["queue"] == beat_specs[name]["queue"]
+
+    @pytest.mark.parametrize("name", sorted(_EXPECTED_CRON))
+    def test_kwargs_match_once_decoded(self, beat_specs, pg_specs, name):
+        # Beat stores kwargs as a JSON *string*; PgPeriodicTask.task_kwargs is a
+        # JSONField. A mismatch here means the cleanup runs with the wrong retention.
+        beat_kwargs = json.loads(beat_specs[name].get("kwargs") or "{}")
+        assert pg_specs[name]["task_kwargs"] == beat_kwargs
+
+    @pytest.mark.parametrize("name,cron", sorted(_EXPECTED_CRON.items()))
+    def test_cron_matches_the_beat_cadence(self, pg_specs, name, cron):
+        assert pg_specs[name]["cron_string"] == cron
+
+
+class TestSeededInert:
+    """Applying the migration must not cause anything to fire."""
+
+    def test_no_spec_declares_itself_pg_owned(self, pg_specs):
+        # pg_owned is set to False in the migration's defaults, never from the spec —
+        # this pins that no spec can smuggle ownership in.
+        assert not any("pg_owned" in spec for spec in pg_specs.values())
+
+    def test_no_spec_presets_a_run_time(self, pg_specs):
+        # A non-NULL next_run_at in the past would read as "overdue" and fire a burst
+        # of catch-up runs the moment the flag is enabled.
+        for spec in pg_specs.values():
+            assert "next_run_at" not in spec
+            assert "last_run_at" not in spec