Feat/un 3445 pg queue integration

backend/backend/settings/base.py

@@ -340,6 +340,8 @@ def filter(self, record):
     # For the organization model
     "account_v2",
     "account_usage",
+    # PG Queue — extension-free bespoke queue (cross-org infra, shared schema)
+    "pg_queue",
     # Django apps should go below this line
     "django.contrib.admin",
     "django.contrib.auth",

backend/pg_queue/apps.py

@@ -0,0 +1,7 @@
+from django.apps import AppConfig
+
+
+class PgQueueConfig(AppConfig):
+    default_auto_field = "django.db.models.BigAutoField"
+    name = "pg_queue"
+    verbose_name = "PG Queue"

backend/pg_queue/executor_rpc.py

@@ -0,0 +1,113 @@
+"""Executor-RPC for the PG path — backend (Django) transport adapter.
+
+The gate + reply_key/timeout orchestration + routing live ONCE in
+``unstract.workflow_execution.executor_rpc`` (shared with the workers). This module
+is the thin Django half: a :class:`DjangoQueueTransport` that enqueues via the ORM
+(``enqueue_task``) and polls ``PgTaskResult``, plus the :func:`get_executor_dispatcher`
+factory that wires them together.
+
+Zero-regression: with the ``pg_queue_enabled`` Flipt flag off the routing dispatcher
+delegates every mode to the unchanged Celery ``ExecutionDispatcher`` and no
+``pg_task_result`` row is created.
+"""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+from django.db import close_old_connections
+
+from pg_queue.models import PgTaskResult
+from pg_queue.producer import enqueue_task
+from unstract.core.polling import poll_for_row
+from unstract.sdk1.execution.dispatcher import ExecutionDispatcher
+from unstract.workflow_execution.executor_rpc import (
+    EXECUTE_TASK,
+    ExecResultRow,
+    PgExecutionDispatcher,
+    QueueTransport,
+    RoutingExecutionDispatcher,
+    resolve_pg_transport,
+)
+
+if TYPE_CHECKING:
+    from unstract.core.data_models import ContinuationSpec
+    from unstract.sdk1.execution.context import ExecutionContext
+
+# Re-exported so existing ``from pg_queue.executor_rpc import …`` imports keep working.
+__all__ = [
+    "DjangoQueueTransport",
+    "PgExecutionDispatcher",
+    "RoutingExecutionDispatcher",
+    "get_executor_dispatcher",
+    "resolve_executor_transport",
+]
+
+
+def resolve_executor_transport(context: ExecutionContext) -> bool:
+    """True → route this executor dispatch over PG; False → Celery (default).
+
+    The single ``pg_queue_enabled`` Flipt flag (fail-closed).
+    """
+    return resolve_pg_transport(context)
+
+
+class DjangoQueueTransport(QueueTransport):
+    """:class:`QueueTransport` over the Django ORM (the backend half).
+
+    Inherits the Protocol so a type-checker verifies this implementation against the
+    seam independently of the ``PgExecutionDispatcher(...)`` construction site.
+    """
+
+    def enqueue(
+        self,
+        *,
+        queue: str,
+        context: ExecutionContext,
+        org_id: str,
+        reply_key: str | None = None,
+        on_success: ContinuationSpec | None = None,
+        on_error: ContinuationSpec | None = None,
+        task_id: str | None = None,
+    ) -> None:
+        enqueue_task(
+            task_name=EXECUTE_TASK,
+            queue=queue,
+            args=[context.to_dict()],
+            org_id=org_id,
+            reply_key=reply_key,
+            on_success=on_success,
+            on_error=on_error,
+            task_id=task_id,
+        )
+
+    def wait_for_result(self, reply_key: str, timeout: float) -> ExecResultRow | None:
+        """Poll ``pg_task_result`` until the row appears or *timeout* elapses.
+
+        Uses the shared :func:`poll_for_row` backoff skeleton, releasing the DB
+        connection between polls (``close_old_connections``) so a long-running RPC
+        does not pin a backend connection and exhaust the pool. Each poll is its own
+        autocommit query, so a row committed by the executor consumer becomes visible
+        — **dispatch must NOT be called inside an open transaction**
+        (``transaction.atomic`` / ``ATOMIC_REQUESTS`` would pin one snapshot and never
+        see the new row).
+        """
+
+        def _fetch() -> ExecResultRow | None:
+            row = PgTaskResult.objects.filter(pk=reply_key).first()
+            if row is None:
+                return None
+            return ExecResultRow(status=row.status, result=row.result, error=row.error)
+
+        return poll_for_row(_fetch, timeout, between_polls=close_old_connections)
+
+
+def get_executor_dispatcher(
+    celery_app: object | None = None,
+) -> RoutingExecutionDispatcher:
+    """Factory: the gate-routed executor dispatcher (PG when enabled, else Celery)."""
+    return RoutingExecutionDispatcher(
+        celery=ExecutionDispatcher(celery_app=celery_app),
+        pg=PgExecutionDispatcher(DjangoQueueTransport()),
+        resolve=resolve_executor_transport,
+    )

backend/pg_queue/flags.py

@@ -0,0 +1,12 @@
+"""Single source of truth for the PG-queue rollout flag key.
+
+One Flipt flag gates the whole PG-queue feature — execution
+(``workflow_v2/transport.py``), scheduler (``scheduler/ownership.py``), and
+executor (``pg_queue/executor_rpc.py``) all read this one key. Kept in a neutral
+leaf module so the three resolvers import a single constant instead of
+duplicating the literal (a grep on ``PG_QUEUE_FLAG_KEY`` finds every use), making
+"one flag" a structural guarantee. This flag is the **sole** rollout control —
+fail-closed to Celery on a blind/unreachable Flipt or any error.
+"""
+
+PG_QUEUE_FLAG_KEY = "pg_queue_enabled"

backend/pg_queue/management/commands/reconcile_pg_schedules.py

@@ -0,0 +1,148 @@
+"""Backfill the pg_periodic_schedule mirror + reconcile Beat/PG schedule
+ownership (Phase 9, ②c).
+
+Run this:
+- **once** after deploying the mirror, to backfill rows for schedules created
+  before the mirror existed (the dual-write only covers schedules touched since);
+- **after each Flipt ramp change** to ``pg_scheduler_enabled``, to apply the new
+  percentage — flipping ``pg_owned`` and the matching Beat ``PeriodicTask`` for
+  every schedule (the create/update path only reconciles the schedule it edits).
+
+It is idempotent and safe to run anytime: with the rollout off it leaves every
+schedule on Beat. Could later be driven periodically (e.g. by the orchestrator);
+kept a command here so the ramp stays an explicit, auditable ops action.
+"""
+
+import json
+from typing import Any
+
+from django.core.management.base import BaseCommand, CommandError
+from django_celery_beat.models import CrontabSchedule, PeriodicTask
+from scheduler.ownership import reconcile_ownership_for, resolve_schedule_owner
+from scheduler.tasks import mirror_periodic_schedule_upsert
+
+from pg_queue.models import PgPeriodicSchedule
+
+# Only the pipeline-trigger PeriodicTasks are scheduled pipelines (other periodic
+# tasks — metrics, audit — are not mirrored).
+_PIPELINE_TASK_PATH = "scheduler.tasks.execute_pipeline_task"
+
+
+def _cron_from_crontab(crontab: CrontabSchedule | None) -> str:
+    """Reconstruct the 5-field cron string from a CrontabSchedule row."""
+    if crontab is None:
+        return ""
+    return (
+        f"{crontab.minute} {crontab.hour} {crontab.day_of_month} "
+        f"{crontab.month_of_year} {crontab.day_of_week}"
+    )
+
+
+class Command(BaseCommand):
+    help = (
+        "Backfill pg_periodic_schedule mirrors for pre-existing schedules and "
+        "reconcile Beat/PG ownership against the current pg_scheduler_enabled "
+        "rollout. Idempotent; with the rollout off, leaves everything on Beat."
+    )
+
+    def add_arguments(self, parser: Any) -> None:
+        parser.add_argument(
+            "--dry-run",
+            action="store_true",
+            help="Report what would change without writing.",
+        )
+
+    def handle(self, *args: Any, **options: Any) -> None:
+        dry_run = options["dry_run"]
+        backfilled = self._backfill_mirrors(dry_run)
+        reconciled, pg_owned, failed = self._reconcile_all(dry_run)
+
+        prefix = "[dry-run] " if dry_run else ""
+        summary = (
+            f"{prefix}backfilled={backfilled} reconciled={reconciled} "
+            f"pg_owned={pg_owned} failed={failed}"
+        )
+        if failed:
+            # Surface failures where the operator looks (and to automation).
+            self.stderr.write(self.style.ERROR(summary))
+            raise CommandError(f"{failed} schedule(s) failed to reconcile")
+        self.stdout.write(self.style.SUCCESS(summary))
+
+    def _mirror_fields_from_args(self, pt: Any, pipeline_id: str) -> dict | None:
+        """Extract the mirror fields from PeriodicTask.args, or None (logged) for a
+        malformed/non-array row — a bad row must not abort the whole command.
+        """
+        try:
+            # json.JSONDecodeError is a ValueError subclass, so one except covers
+            # both the parse error and the non-array guard below.
+            task_args = json.loads(pt.args or "[]")
+            if not isinstance(task_args, list):
+                raise ValueError(f"expected JSON array, got {type(task_args).__name__}")
+        except ValueError as exc:
+            self.stderr.write(
+                self.style.ERROR(
+                    f"skipping pipeline {pipeline_id}: bad PeriodicTask.args ({exc})"
+                )
+            )
+            return None
+        return {
+            "workflow_id": task_args[0] if len(task_args) > 0 else None,
+            "organization_id": (task_args[1] if len(task_args) > 1 else "") or "",
+            # args[6] is the synthetic "Pipeline job-<id>" label; the real name
+            # self-heals via the dual-write on the next schedule edit.
+            "pipeline_name": task_args[6] if len(task_args) > 6 else "",
+        }
+
+    def _backfill_mirrors(self, dry_run: bool) -> int:
+        """Create a mirror row for every pipeline-trigger PeriodicTask lacking one."""
+        # Pre-fetch the already-mirrored ids in one query (avoid an EXISTS per row).
+        mirrored = {
+            str(pk)
+            for pk in PgPeriodicSchedule.objects.values_list("pipeline_id", flat=True)
+        }
+        backfilled = 0
+        for pt in PeriodicTask.objects.filter(task=_PIPELINE_TASK_PATH):
+            pipeline_id = pt.name  # = str(pipeline.pk)
+            if pipeline_id in mirrored:
+                continue
+            fields = self._mirror_fields_from_args(pt, pipeline_id)
+            if fields is None:
+                continue
+            self.stdout.write(
+                f"backfill mirror for pipeline {pipeline_id} (enabled={pt.enabled})"
+            )
+            if not dry_run:
+                mirror_periodic_schedule_upsert(
+                    pipeline_id=pipeline_id,
+                    cron_string=_cron_from_crontab(pt.crontab),
+                    enabled=pt.enabled,
+                    **fields,
+                )
+            backfilled += 1
+        return backfilled
+
+    def _reconcile_all(self, dry_run: bool) -> tuple[int, int, int]:
+        """Reconcile ownership for every mirror row against the current rollout.
+        Returns (reconciled, pg_owned, failed).
+        """
+        reconciled = pg_owned = failed = 0
+        for row in PgPeriodicSchedule.objects.all():
+            if dry_run:
+                # Preview only — read the would-be owner (no DB write) so an
+                # operator can see how many a ramp change would hand to PG.
+                reconciled += 1
+                if resolve_schedule_owner(str(row.pipeline_id), row.organization_id):
+                    pg_owned += 1
+                continue
+            # mirror.enabled tracks pipeline.active (dual-write); use it as the
+            # 'active' input so a paused schedule isn't re-enabled by reconcile.
+            result = reconcile_ownership_for(
+                str(row.pipeline_id), row.organization_id, active=row.enabled
+            )
+            if result is None:  # transaction failed (already logged)
+                failed += 1
+                continue
+            reconciled += 1
+            if result:
+                pg_owned += 1
+        return reconciled, pg_owned, failed

backend/pg_queue/migrations/0001_initial.py

@@ -0,0 +1,34 @@
+# Generated by Django 4.2.1 on 2026-06-11 14:42
+
+import django.utils.timezone
+from django.db import migrations, models
+
+
+class Migration(migrations.Migration):
+    initial = True
+
+    dependencies = []
+
+    operations = [
+        migrations.CreateModel(
+            name="PgQueueMessage",
+            fields=[
+                ("msg_id", models.BigAutoField(primary_key=True, serialize=False)),
+                ("queue_name", models.TextField()),
+                ("message", models.JSONField()),
+                ("org_id", models.TextField(blank=True, default="")),
+                ("enqueued_at", models.DateTimeField(default=django.utils.timezone.now)),
+                ("vt", models.DateTimeField(default=django.utils.timezone.now)),
+                ("read_ct", models.IntegerField(default=0)),
+            ],
+            options={
+                "db_table": "pg_queue_message",
+                "indexes": [
+                    models.Index(
+                        fields=["queue_name", "vt", "msg_id"],
+                        name="pg_queue_message_dequeue_idx",
+                    )
+                ],
+            },
+        ),
+    ]

backend/pg_queue/migrations/0002_remove_pgqueuemessage_pg_queue_message_dequeue_idx_and_more.py

@@ -0,0 +1,30 @@
+# Generated by Django 4.2.1 on 2026-06-12 11:47
+
+from django.db import migrations, models
+
+
+class Migration(migrations.Migration):
+    dependencies = [
+        ("pg_queue", "0001_initial"),
+    ]
+
+    operations = [
+        migrations.RemoveIndex(
+            model_name="pgqueuemessage",
+            name="pg_queue_message_dequeue_idx",
+        ),
+        migrations.AddField(
+            model_name="pgqueuemessage",
+            name="priority",
+            field=models.SmallIntegerField(default=5),
+        ),
+        migrations.AddIndex(
+            model_name="pgqueuemessage",
+            index=models.Index(
+                models.F("queue_name"),
+                models.OrderBy(models.F("priority"), descending=True),
+                models.F("msg_id"),
+                name="pg_queue_message_dequeue_idx",
+            ),
+        ),
+    ]

backend/pg_queue/migrations/0003_pgqueuemessage_pg_queue_message_priority_range.py

@@ -0,0 +1,25 @@
+# Generated by Django 4.2.1 on 2026-06-12 14:05
+
+from django.db import migrations, models
+
+
+class Migration(migrations.Migration):
+    dependencies = [
+        ("pg_queue", "0002_remove_pgqueuemessage_pg_queue_message_dequeue_idx_and_more"),
+    ]
+
+    operations = [
+        # `check=` is correct for the pinned Django (4.2). It's deprecated in
+        # 5.1 in favour of `condition=` and removed in 6.0 — but fresh installs
+        # always replay this under the Django we ship, so leave it as-is.
+        # When the pin is bumped to >= 6.0, squash these migrations (or do the
+        # behaviour-preserving `check=` -> `condition=` edit) as part of that
+        # upgrade so a from-scratch migrate still runs.
+        migrations.AddConstraint(
+            model_name="pgqueuemessage",
+            constraint=models.CheckConstraint(
+                check=models.Q(("priority__gte", 1), ("priority__lte", 10)),
+                name="pg_queue_message_priority_range",
+            ),
+        ),
+    ]