DataOps-platform/app/core/mcp/persistence.py

"""PostgreSQL persistence adapters for the DataOps MCP control plane."""

from __future__ import annotations

import hashlib
import json
import uuid
from datetime import UTC, datetime
from zoneinfo import ZoneInfo

from croniter import croniter
from sqlalchemy import text

from app.core.common.identifiers import new_governance_uid


def _json(value):
    return json.dumps(
        value,
        sort_keys=True,
        separators=(",", ":"),
        ensure_ascii=False,
    )


def _schedule_hash(schedule_plan):
    return hashlib.sha256(_json(schedule_plan).encode("utf-8")).hexdigest()


def _required_string(value, label, maximum=500):
    if not isinstance(value, str) or not value.strip():
        raise ValueError(f"{label} is required")
    normalized = value.strip()
    if len(normalized) > maximum:
        raise ValueError(f"{label} exceeds {maximum} characters")
    return normalized


class PostgresSchedulingPlanStore:
    def __init__(self, engine):
        self.engine = engine

    def create_candidate(self, record):
        dataflow_uid = _required_string(record.get("dataflow_uid"), "dataflow_uid", 100)
        business_domain = _required_string(
            record.get("business_domain"), "business_domain", 200
        )
        environment = _required_string(record.get("environment"), "environment", 20)
        created_by = _required_string(record.get("created_by"), "created_by", 200)
        workflow_spec = dict(record.get("workflow_spec") or {})
        schedule_plan = dict(record.get("schedule_plan") or {})
        workflow_hash = _required_string(
            record.get("workflow_hash"), "workflow_hash", 64
        )
        candidate_id = new_governance_uid()
        binding_id = new_governance_uid()
        schedule_id = new_governance_uid()
        placeholder = f"candidate_{candidate_id.replace('-', '')}"
        metadata = {
            "gateway_status": "candidate",
            "correlation_id": record.get("correlation_id"),
        }

        with self.engine.begin() as connection:
            connection.execute(
                text("SELECT pg_advisory_xact_lock(hashtext(:key))"),
                {"key": f"{dataflow_uid}:{environment}"},
            )
            version_no = connection.execute(
                text(
                    "SELECT COALESCE(MAX(version_no), 0) + 1 "
                    "FROM public.dataflow_workflow_versions "
                    "WHERE dataflow_uid = CAST(:uid AS uuid) "
                    "AND environment = :environment"
                ),
                {"uid": dataflow_uid, "environment": environment},
            ).scalar_one()
            connection.execute(
                text(
                    "INSERT INTO public.dataflow_workflow_versions "
                    "(id, dataflow_uid, environment, version_no, "
                    "n8n_workflow_id, n8n_workflow_name, definition_hash, "
                    "definition_snapshot, status, created_by, engine_type, "
                    "engine_definition_id, engine_revision, "
                    "deployment_metadata, workflow_spec, schedule_plan, "
                    "business_domain, created_by_subject, write_authorized) "
                    "VALUES (CAST(:id AS uuid), CAST(:dataflow_uid AS uuid), "
                    ":environment, :version_no, NULL, NULL, :definition_hash, "
                    "CAST(:definition_snapshot AS jsonb), 'draft', NULL, "
                    "'kestra', :engine_definition_id, NULL, "
                    "CAST(:deployment_metadata AS jsonb), "
                    "CAST(:workflow_spec AS jsonb), "
                    "CAST(:schedule_plan AS jsonb), :business_domain, "
                    ":created_by_subject, FALSE)"
                ),
                {
                    "id": candidate_id,
                    "dataflow_uid": dataflow_uid,
                    "environment": environment,
                    "version_no": version_no,
                    "definition_hash": workflow_hash,
                    "definition_snapshot": _json(workflow_spec),
                    "engine_definition_id": placeholder,
                    "deployment_metadata": _json(metadata),
                    "workflow_spec": _json(workflow_spec),
                    "schedule_plan": _json(schedule_plan),
                    "business_domain": business_domain,
                    "created_by_subject": created_by,
                },
            )
            connection.execute(
                text(
                    "INSERT INTO public.workflow_engine_bindings "
                    "(id, workflow_version_id, dataflow_uid, environment, "
                    "engine_type, engine_definition_id, engine_revision, "
                    "role, status, metadata) "
                    "VALUES (CAST(:id AS uuid), CAST(:version_id AS uuid), "
                    "CAST(:dataflow_uid AS uuid), :environment, 'kestra', "
                    ":engine_definition_id, NULL, 'standby', 'disabled', "
                    "CAST(:metadata AS jsonb))"
                ),
                {
                    "id": binding_id,
                    "version_id": candidate_id,
                    "dataflow_uid": dataflow_uid,
                    "environment": environment,
                    "engine_definition_id": placeholder,
                    "metadata": _json(metadata),
                },
            )
            connection.execute(
                text(
                    "INSERT INTO public.workflow_schedules "
                    "(id, workflow_version_id, schedule_plan, schedule_hash, "
                    "timezone, status, created_by) "
                    "VALUES (CAST(:id AS uuid), CAST(:version_id AS uuid), "
                    "CAST(:schedule_plan AS jsonb), :schedule_hash, "
                    ":timezone, 'draft', NULL)"
                ),
                {
                    "id": schedule_id,
                    "version_id": candidate_id,
                    "schedule_plan": _json(schedule_plan),
                    "schedule_hash": _schedule_hash(schedule_plan),
                    "timezone": schedule_plan.get("timezone"),
                },
            )

        return {
            **record,
            "candidate_id": candidate_id,
            "version_no": int(version_no),
            "status": "candidate",
        }

    def get_candidate(self, candidate_id):
        with self.engine.connect() as connection:
            row = (
                connection.execute(
                    text(
                        "SELECT id::text AS candidate_id, "
                        "dataflow_uid::text AS dataflow_uid, business_domain, "
                        "environment, version_no, workflow_spec, schedule_plan, "
                        "definition_hash AS workflow_hash, status, "
                        "created_by_subject AS created_by, write_authorized, "
                        "deployment_metadata "
                        "FROM public.dataflow_workflow_versions "
                        "WHERE id = CAST(:id AS uuid) AND engine_type = 'kestra'"
                    ),
                    {"id": candidate_id},
                )
                .mappings()
                .one_or_none()
            )
        if row is None:
            raise KeyError(candidate_id)
        result = dict(row)
        metadata = dict(result.pop("deployment_metadata") or {})
        result["status"] = metadata.get("gateway_status", result["status"])
        result["trusted_write_authorized"] = bool(result.pop("write_authorized", False))
        return result

    def record_deployment(self, candidate_id, deployment):
        metadata = {
            "gateway_status": deployment["status"],
            "deployment": deployment,
        }
        with self.engine.begin() as connection:
            updated = connection.execute(
                text(
                    "UPDATE public.dataflow_workflow_versions "
                    "SET engine_definition_id = :flow_id, "
                    "engine_revision = :revision, "
                    "deployment_metadata = deployment_metadata || "
                    "CAST(:metadata AS jsonb), updated_at = CURRENT_TIMESTAMP "
                    "WHERE id = CAST(:id AS uuid) AND engine_type = 'kestra'"
                ),
                {
                    "id": candidate_id,
                    "flow_id": deployment["flow_id"],
                    "revision": deployment.get("engine_revision"),
                    "metadata": _json(metadata),
                },
            )
            if updated.rowcount != 1:
                raise KeyError(candidate_id)
            connection.execute(
                text(
                    "UPDATE public.workflow_engine_bindings "
                    "SET engine_definition_id = :flow_id, "
                    "engine_revision = :revision, metadata = CAST(:metadata AS jsonb), "
                    "updated_at = CURRENT_TIMESTAMP "
                    "WHERE workflow_version_id = CAST(:id AS uuid) "
                    "AND engine_type = 'kestra'"
                ),
                {
                    "id": candidate_id,
                    "flow_id": deployment["flow_id"],
                    "revision": deployment.get("engine_revision"),
                    "metadata": _json(deployment),
                },
            )
        return dict(deployment)

    def ensure_rule_deployment_identity(
        self,
        candidate_id,
        environment,
        *,
        status,
    ):
        if status not in {"disabled", "canary"}:
            raise ValueError("rule deployment status is invalid")
        deployment_id = new_governance_uid()
        with self.engine.begin() as connection:
            target = connection.execute(
                text(
                    """
                    SELECT w.id::text AS workflow_version_id,
                           v.id::text AS dataflow_version_id,
                           s.id::text AS schedule_plan_id
                    FROM public.dataflow_workflow_versions w
                    JOIN public.dataflow_versions v
                      ON v.dataflow_uid = w.dataflow_uid
                     AND v.version_no = w.version_no
                     AND v.status = 'released'
                    LEFT JOIN public.workflow_schedules s
                      ON s.workflow_version_id = w.id
                    WHERE w.id = CAST(:candidate_id AS uuid)
                      AND w.environment = :environment
                    """
                ),
                {
                    "candidate_id": candidate_id,
                    "environment": environment,
                },
            ).mappings().one_or_none()
            if target is None:
                raise RuntimeError(
                    "released dataflow version for workflow was not found"
                )
            existing = connection.execute(
                text(
                    """
                    SELECT id::text AS deployment_id,
                           workflow_version_id::text,
                           environment, status
                    FROM public.dataflow_deployments
                    WHERE dataflow_version_id =
                        CAST(:dataflow_version_id AS uuid)
                      AND environment = :environment
                    FOR UPDATE
                    """
                ),
                {
                    "dataflow_version_id": target[
                        "dataflow_version_id"
                    ],
                    "environment": environment,
                },
            ).mappings().one_or_none()
            if existing is None:
                existing = connection.execute(
                    text(
                        """
                        INSERT INTO public.dataflow_deployments (
                            id, dataflow_version_id, environment,
                            workflow_version_id, schedule_plan_id,
                            deployment_config, status, updated_at
                        ) VALUES (
                            CAST(:id AS uuid),
                            CAST(:dataflow_version_id AS uuid),
                            :environment,
                            CAST(:workflow_version_id AS uuid),
                            CAST(:schedule_plan_id AS uuid),
                            '{}'::jsonb, :status, CURRENT_TIMESTAMP
                        )
                        RETURNING id::text AS deployment_id,
                                  workflow_version_id::text,
                                  environment, status
                        """
                    ),
                    {
                        "id": deployment_id,
                        "dataflow_version_id": target[
                            "dataflow_version_id"
                        ],
                        "environment": environment,
                        "workflow_version_id": target[
                            "workflow_version_id"
                        ],
                        "schedule_plan_id": target["schedule_plan_id"],
                        "status": status,
                    },
                ).mappings().one()
            else:
                if existing["workflow_version_id"] is None:
                    if existing["status"] != "disabled":
                        raise RuntimeError(
                            "unbound deployment is not disabled"
                        )
                    existing = connection.execute(
                        text(
                            """
                            UPDATE public.dataflow_deployments
                            SET workflow_version_id =
                                    CAST(:workflow_version_id AS uuid),
                                schedule_plan_id =
                                    CAST(:schedule_plan_id AS uuid),
                                updated_at = CURRENT_TIMESTAMP
                            WHERE id = CAST(:id AS uuid)
                              AND workflow_version_id IS NULL
                              AND status = 'disabled'
                            RETURNING id::text AS deployment_id,
                                      workflow_version_id::text,
                                      environment, status
                            """
                        ),
                        {
                            "id": existing["deployment_id"],
                            "workflow_version_id": target[
                                "workflow_version_id"
                            ],
                            "schedule_plan_id": target[
                                "schedule_plan_id"
                            ],
                        },
                    ).mappings().one()
                elif (
                    existing["workflow_version_id"]
                    != target["workflow_version_id"]
                ):
                    raise RuntimeError(
                        "deployment belongs to another workflow version"
                    )
                if status == "canary" and existing["status"] == "disabled":
                    existing = connection.execute(
                        text(
                            """
                            UPDATE public.dataflow_deployments
                            SET status = 'canary',
                                updated_at = CURRENT_TIMESTAMP
                            WHERE id = CAST(:id AS uuid)
                              AND status = 'disabled'
                            RETURNING id::text AS deployment_id,
                                      workflow_version_id::text,
                                      environment, status
                            """
                        ),
                        {"id": existing["deployment_id"]},
                    ).mappings().one()
                elif existing["status"] not in {status, "canary"}:
                    raise RuntimeError(
                        "deployment state cannot enter canary"
                    )
        return dict(existing)

    def get_deployment(self, candidate_id):
        with self.engine.connect() as connection:
            metadata = connection.execute(
                text(
                    "SELECT deployment_metadata "
                    "FROM public.dataflow_workflow_versions "
                    "WHERE id = CAST(:id AS uuid) AND engine_type = 'kestra'"
                ),
                {"id": candidate_id},
            ).scalar_one_or_none()
        if metadata is None or not dict(metadata).get("deployment"):
            raise KeyError(candidate_id)
        return dict(metadata)["deployment"]

    def get_active_deployment(self, candidate_id):
        with self.engine.connect() as connection:
            row = (
                connection.execute(
                    text(
                        "SELECT active.id::text AS candidate_id, "
                        "active.deployment_metadata "
                        "FROM public.dataflow_workflow_versions candidate "
                        "JOIN public.dataflow_workflow_versions active "
                        "ON active.dataflow_uid = candidate.dataflow_uid "
                        "AND active.environment = candidate.environment "
                        "AND active.engine_type = 'kestra' "
                        "AND active.status = 'active' "
                        "AND active.id <> candidate.id "
                        "WHERE candidate.id = CAST(:id AS uuid) "
                        "ORDER BY active.version_no DESC LIMIT 1"
                    ),
                    {"id": candidate_id},
                )
                .mappings()
                .one_or_none()
            )
        if row is None:
            return None
        deployment = dict(row["deployment_metadata"] or {}).get("deployment")
        if not isinstance(deployment, dict):
            raise ValueError("active deployment is incomplete")
        return {**deployment, "candidate_id": row["candidate_id"]}

    def record_canary_execution(self, candidate_id, evidence):
        evidence_id = new_governance_uid()
        with self.engine.begin() as connection:
            connection.execute(
                text(
                    "INSERT INTO public.workflow_canary_evidence "
                    "(id, workflow_version_id, engine_execution_id, status, "
                    "sample_runs, verification_summary) "
                    "VALUES (CAST(:id AS uuid), CAST(:version_id AS uuid), "
                    ":execution_id, 'started', 0, '{}'::jsonb)"
                ),
                {
                    "id": evidence_id,
                    "version_id": candidate_id,
                    "execution_id": evidence["execution_id"],
                },
            )
        return dict(evidence)

    def get_canary_evidence(self, candidate_id):
        with self.engine.connect() as connection:
            row = (
                connection.execute(
                    text(
                        "SELECT engine_execution_id AS execution_id, status, "
                        "sample_runs, verified_by, verification_summary "
                        "FROM public.workflow_canary_evidence "
                        "WHERE workflow_version_id = CAST(:id AS uuid) "
                        "ORDER BY created_at DESC LIMIT 1"
                    ),
                    {"id": candidate_id},
                )
                .mappings()
                .one_or_none()
            )
        return dict(row) if row is not None else None

    def record_canary_verification(self, candidate_id, evidence):
        status = evidence.get("status")
        if status not in {"passed", "failed"}:
            raise ValueError("canary verification must be passed or failed")
        verified_by = _required_string(
            evidence.get("verified_by"), "canary verified_by", 200
        )
        summary = dict(evidence.get("verification_summary") or {})
        with self.engine.begin() as connection:
            updated = connection.execute(
                text(
                    "UPDATE public.workflow_canary_evidence "
                    "SET status = :status, sample_runs = :sample_runs, "
                    "verified_by = :verified_by, "
                    "verification_summary = CAST(:summary AS jsonb), "
                    "verified_at = CURRENT_TIMESTAMP "
                    "WHERE workflow_version_id = CAST(:version_id AS uuid) "
                    "AND engine_execution_id = :execution_id "
                    "AND status = 'started'"
                ),
                {
                    "status": status,
                    "sample_runs": int(evidence.get("sample_runs", 0)),
                    "verified_by": verified_by,
                    "summary": _json(summary),
                    "version_id": candidate_id,
                    "execution_id": evidence.get("execution_id"),
                },
            )
            if updated.rowcount != 1:
                raise ValueError("canary evidence is already finalized")
        return dict(evidence)

    def get_execution_record(self, execution_id):
        execution_id = _required_string(execution_id, "execution_id", 255)
        with self.engine.connect() as connection:
            row = (
                connection.execute(
                    text(
                        "SELECT run.id::text AS run_id, "
                        "run.workflow_version_id::text AS workflow_version_id, "
                        "run.schedule_id::text AS schedule_id, "
                        "run.engine_execution_id AS execution_id, "
                        "run.trigger_type, run.status, "
                        "run.correlation_id::text AS correlation_id, "
                        "run.attempt, run.run_metadata, "
                        "version.business_domain, version.environment, "
                        "(SELECT COUNT(*) FROM public.workflow_runs replay "
                        "WHERE replay.parent_run_id = run.id) AS retry_count "
                        "FROM public.workflow_runs run "
                        "JOIN public.dataflow_workflow_versions version "
                        "ON version.id = run.workflow_version_id "
                        "WHERE run.engine_type = 'kestra' "
                        "AND run.engine_execution_id = :execution_id"
                    ),
                    {"execution_id": execution_id},
                )
                .mappings()
                .one_or_none()
            )
        if row is None:
            raise KeyError(execution_id)
        return dict(row)

    def record_retry(self, execution_id, result):
        execution_id = _required_string(execution_id, "execution_id", 255)
        replay_execution_id = _required_string(
            (result or {}).get("id"), "replay execution id", 255
        )
        with self.engine.begin() as connection:
            original = (
                connection.execute(
                    text(
                        "SELECT id::text AS run_id, "
                        "workflow_version_id::text AS workflow_version_id, "
                        "schedule_id::text AS schedule_id, "
                        "correlation_id::text AS correlation_id, attempt "
                        "FROM public.workflow_runs "
                        "WHERE engine_type = 'kestra' "
                        "AND engine_execution_id = :execution_id "
                        "FOR UPDATE"
                    ),
                    {"execution_id": execution_id},
                )
                .mappings()
                .one_or_none()
            )
            if original is None:
                raise KeyError(execution_id)
            next_attempt = connection.execute(
                text(
                    "SELECT GREATEST(:original_attempt, "
                    "COALESCE(MAX(attempt), 0)) + 1 "
                    "FROM public.workflow_runs "
                    "WHERE parent_run_id = CAST(:run_id AS uuid)"
                ),
                {
                    "original_attempt": original["attempt"],
                    "run_id": original["run_id"],
                },
            ).scalar_one()
            connection.execute(
                text(
                    "INSERT INTO public.workflow_runs "
                    "(id, workflow_version_id, schedule_id, engine_type, "
                    "engine_execution_id, trigger_type, status, "
                    "correlation_id, parent_run_id, attempt, run_metadata) "
                    "VALUES (CAST(:id AS uuid), CAST(:version_id AS uuid), "
                    "CAST(:schedule_id AS uuid), 'kestra', :execution_id, "
                    "'replay', 'queued', CAST(:correlation_id AS uuid), "
                    "CAST(:parent_run_id AS uuid), :attempt, "
                    "CAST(:metadata AS jsonb))"
                ),
                {
                    "id": new_governance_uid(),
                    "version_id": original["workflow_version_id"],
                    "schedule_id": original["schedule_id"],
                    "execution_id": replay_execution_id,
                    "correlation_id": original["correlation_id"],
                    "parent_run_id": original["run_id"],
                    "attempt": int(next_attempt),
                    "metadata": _json({"source_execution_id": execution_id}),
                },
            )
        return dict(result)

    @staticmethod
    def _parse_timestamp(value, label):
        value = _required_string(value, label, 100)
        try:
            parsed = datetime.fromisoformat(value.replace("Z", "+00:00"))
        except ValueError as exc:
            raise ValueError(f"{label} must be an ISO 8601 timestamp") from exc
        if parsed.tzinfo is None:
            raise ValueError(f"{label} must include a timezone")
        return parsed

    def estimate_backfill_runs(self, candidate_id, start, end):
        candidate = self.get_candidate(candidate_id)
        schedule_plan = dict(candidate.get("schedule_plan") or {})
        start_at = self._parse_timestamp(start, "backfill start")
        end_at = self._parse_timestamp(end, "backfill end")
        if end_at <= start_at:
            raise ValueError("backfill end must be after start")
        try:
            timezone = ZoneInfo(schedule_plan.get("timezone") or "UTC")
        except (KeyError, ValueError) as exc:
            raise ValueError("schedule timezone is invalid") from exc
        start_local = start_at.astimezone(timezone)
        end_local = end_at.astimezone(timezone)
        total = 0
        for trigger in schedule_plan.get("triggers") or []:
            if trigger.get("type") != "cron":
                continue
            expression = _required_string(
                trigger.get("expression"), "cron expression", 100
            )
            if len(expression.split()) != 5:
                raise ValueError("backfill estimation supports five-field cron only")
            iterator = croniter(expression, start_local)
            occurrence = iterator.get_next(datetime)
            while occurrence <= end_local:
                total += 1
                if total > 1000:
                    return total
                occurrence = iterator.get_next(datetime)
        if total == 0:
            raise ValueError("candidate has no cron occurrences in window")
        return total

    def record_backfill(self, candidate_id, result):
        execution_id = _required_string(
            (result or {}).get("id"), "backfill execution id", 255
        )
        estimated_runs = int((result or {}).get("estimated_runs", 0))
        if estimated_runs < 1:
            raise ValueError("estimated_runs must be positive")
        with self.engine.begin() as connection:
            schedule_id = connection.execute(
                text(
                    "SELECT id::text FROM public.workflow_schedules "
                    "WHERE workflow_version_id = CAST(:id AS uuid) "
                    "ORDER BY created_at DESC LIMIT 1"
                ),
                {"id": candidate_id},
            ).scalar_one_or_none()
            if schedule_id is None:
                raise KeyError(candidate_id)
            connection.execute(
                text(
                    "INSERT INTO public.workflow_runs "
                    "(id, workflow_version_id, schedule_id, engine_type, "
                    "engine_execution_id, trigger_type, status, "
                    "correlation_id, attempt, run_metadata) "
                    "VALUES (CAST(:id AS uuid), CAST(:version_id AS uuid), "
                    "CAST(:schedule_id AS uuid), 'kestra', :execution_id, "
                    "'backfill', 'queued', CAST(:correlation_id AS uuid), "
                    "1, CAST(:metadata AS jsonb))"
                ),
                {
                    "id": new_governance_uid(),
                    "version_id": candidate_id,
                    "schedule_id": schedule_id,
                    "execution_id": execution_id,
                    "correlation_id": new_governance_uid(),
                    "metadata": _json({"estimated_runs": estimated_runs}),
                },
            )
        return dict(result)

    def _claim_operation(
        self,
        action,
        idempotency_key,
        candidate_id,
        *,
        actor_subject,
        correlation_id,
        result,
    ):
        operation_id = new_governance_uid()
        with self.engine.begin() as connection:
            inserted = connection.execute(
                text(
                    "INSERT INTO public.workflow_gateway_operations "
                    "(id, workflow_version_id, action, idempotency_key, status, "
                    "result, actor_subject, correlation_id) "
                    "VALUES (CAST(:id AS uuid), CAST(:version_id AS uuid), "
                    ":action, :key, 'claimed', CAST(:result AS jsonb), "
                    ":actor, CAST(:correlation_id AS uuid)) "
                    "ON CONFLICT (action, idempotency_key) DO NOTHING "
                    "RETURNING id::text"
                ),
                {
                    "id": operation_id,
                    "version_id": candidate_id,
                    "action": action,
                    "key": idempotency_key,
                    "result": _json(result),
                    "actor": actor_subject,
                    "correlation_id": correlation_id,
                },
            ).scalar_one_or_none()
            if inserted is not None:
                return True, result
            previous = (
                connection.execute(
                    text(
                        "SELECT workflow_version_id::text AS version_id, "
                        "result FROM public.workflow_gateway_operations "
                        "WHERE action = :action AND idempotency_key = :key"
                    ),
                    {"action": action, "key": idempotency_key},
                )
                .mappings()
                .one()
            )
            if previous["version_id"] != candidate_id:
                raise ValueError("idempotency key belongs to another workflow version")
            return False, dict(previous["result"])

    def claim_promotion(
        self,
        idempotency_key,
        candidate_id,
        *,
        actor_subject,
        correlation_id,
    ):
        deployment = self.get_deployment(candidate_id)
        result = {
            **deployment,
            "candidate_id": candidate_id,
            "status": "promoted",
        }
        return self._claim_operation(
            "promote_candidate",
            idempotency_key,
            candidate_id,
            actor_subject=actor_subject,
            correlation_id=correlation_id,
            result=result,
        )

    def _complete_operation(
        self, connection, action, idempotency_key, candidate_id, result
    ):
        updated = connection.execute(
            text(
                "UPDATE public.workflow_gateway_operations "
                "SET status = 'succeeded', result = CAST(:result AS jsonb), "
                "safe_error = NULL, updated_at = CURRENT_TIMESTAMP "
                "WHERE action = :action AND idempotency_key = :key "
                "AND workflow_version_id = CAST(:version_id AS uuid) "
                "AND status = 'claimed'"
            ),
            {
                "action": action,
                "key": idempotency_key,
                "version_id": candidate_id,
                "result": _json(result),
            },
        )
        if updated.rowcount != 1:
            raise ValueError("gateway operation is not in claimed state")

    def complete_promotion(self, idempotency_key, candidate_id, result):
        with self.engine.begin() as connection:
            row = connection.execute(
                text(
                    "SELECT dataflow_uid::text, environment "
                    "FROM public.dataflow_workflow_versions "
                    "WHERE id = CAST(:id AS uuid) FOR UPDATE"
                ),
                {"id": candidate_id},
            ).one_or_none()
            if row is None:
                raise KeyError(candidate_id)
            dataflow_uid, environment = row
            connection.execute(
                text("SELECT pg_advisory_xact_lock(hashtext(:key))"),
                {"key": f"{dataflow_uid}:{environment}"},
            )
            connection.execute(
                text(
                    "UPDATE public.workflow_engine_bindings "
                    "SET role = 'archived', status = 'disabled', "
                    "updated_at = CURRENT_TIMESTAMP "
                    "WHERE dataflow_uid = CAST(:dataflow_uid AS uuid) "
                    "AND environment = :environment "
                    "AND role = 'primary' AND status = 'enabled' "
                    "AND workflow_version_id <> CAST(:id AS uuid)"
                ),
                {
                    "dataflow_uid": dataflow_uid,
                    "environment": environment,
                    "id": candidate_id,
                },
            )
            connection.execute(
                text(
                    "UPDATE public.dataflow_workflow_versions "
                    "SET status = 'superseded', updated_at = CURRENT_TIMESTAMP "
                    "WHERE dataflow_uid = CAST(:dataflow_uid AS uuid) "
                    "AND environment = :environment AND status = 'active' "
                    "AND id <> CAST(:id AS uuid)"
                ),
                {
                    "dataflow_uid": dataflow_uid,
                    "environment": environment,
                    "id": candidate_id,
                },
            )
            connection.execute(
                text(
                    "UPDATE public.workflow_schedules SET status = 'retired', "
                    "updated_at = CURRENT_TIMESTAMP "
                    "WHERE workflow_version_id IN ("
                    "SELECT id FROM public.dataflow_workflow_versions "
                    "WHERE dataflow_uid = CAST(:dataflow_uid AS uuid) "
                    "AND environment = :environment "
                    "AND id <> CAST(:id AS uuid)) "
                    "AND status = 'active'"
                ),
                {
                    "dataflow_uid": dataflow_uid,
                    "environment": environment,
                    "id": candidate_id,
                },
            )
            connection.execute(
                text(
                    "UPDATE public.dataflow_workflow_versions "
                    "SET status = 'active', "
                    "deployment_metadata = deployment_metadata || "
                    '\'{"gateway_status":"promoted"}\'::jsonb, '
                    "updated_at = CURRENT_TIMESTAMP "
                    "WHERE id = CAST(:id AS uuid)"
                ),
                {"id": candidate_id},
            )
            connection.execute(
                text(
                    "UPDATE public.workflow_engine_bindings "
                    "SET role = 'primary', status = 'enabled', "
                    "updated_at = CURRENT_TIMESTAMP "
                    "WHERE workflow_version_id = CAST(:id AS uuid)"
                ),
                {"id": candidate_id},
            )
            connection.execute(
                text(
                    "UPDATE public.workflow_schedules "
                    "SET status = 'active', updated_at = CURRENT_TIMESTAMP "
                    "WHERE workflow_version_id = CAST(:id AS uuid)"
                ),
                {"id": candidate_id},
            )
            self._complete_operation(
                connection,
                "promote_candidate",
                idempotency_key,
                candidate_id,
                result,
            )
        return result

    def release_promotion(self, idempotency_key, candidate_id):
        self._release_operation("promote_candidate", idempotency_key, candidate_id)

    def mark_paused(self, candidate_id):
        with self.engine.begin() as connection:
            connection.execute(
                text(
                    "UPDATE public.workflow_engine_bindings "
                    "SET status = 'disabled', updated_at = CURRENT_TIMESTAMP "
                    "WHERE workflow_version_id = CAST(:id AS uuid)"
                ),
                {"id": candidate_id},
            )
            connection.execute(
                text(
                    "UPDATE public.workflow_schedules "
                    "SET status = 'paused', updated_at = CURRENT_TIMESTAMP "
                    "WHERE workflow_version_id = CAST(:id AS uuid)"
                ),
                {"id": candidate_id},
            )
            updated = connection.execute(
                text(
                    "UPDATE public.dataflow_workflow_versions "
                    "SET deployment_metadata = deployment_metadata || "
                    '\'{"gateway_status":"paused"}\'::jsonb, '
                    "updated_at = CURRENT_TIMESTAMP "
                    "WHERE id = CAST(:id AS uuid)"
                ),
                {"id": candidate_id},
            )
            if updated.rowcount != 1:
                raise KeyError(candidate_id)

    def get_previous_deployment(self, candidate_id):
        with self.engine.connect() as connection:
            row = (
                connection.execute(
                    text(
                        "SELECT previous.id::text AS candidate_id, "
                        "previous.deployment_metadata "
                        "FROM public.dataflow_workflow_versions current "
                        "JOIN LATERAL ("
                        "SELECT id, deployment_metadata "
                        "FROM public.dataflow_workflow_versions candidate "
                        "WHERE candidate.dataflow_uid = current.dataflow_uid "
                        "AND candidate.environment = current.environment "
                        "AND candidate.engine_type = 'kestra' "
                        "AND candidate.version_no < current.version_no "
                        "ORDER BY candidate.version_no DESC LIMIT 1"
                        ") previous ON TRUE "
                        "WHERE current.id = CAST(:id AS uuid)"
                    ),
                    {"id": candidate_id},
                )
                .mappings()
                .one_or_none()
            )
        if row is None:
            raise ValueError("previous deployment is not available")
        deployment = dict(row["deployment_metadata"] or {}).get("deployment")
        if not isinstance(deployment, dict):
            raise ValueError("previous deployment is incomplete")
        return {**deployment, "candidate_id": row["candidate_id"]}

    def claim_rollback(
        self,
        idempotency_key,
        candidate_id,
        *,
        actor_subject,
        correlation_id,
    ):
        previous = self.get_previous_deployment(candidate_id)
        result = {
            "candidate_id": candidate_id,
            "status": "rolled_back",
            "active_candidate_id": previous["candidate_id"],
        }
        return self._claim_operation(
            "rollback_to_previous_version",
            idempotency_key,
            candidate_id,
            actor_subject=actor_subject,
            correlation_id=correlation_id,
            result=result,
        )

    def complete_rollback(self, idempotency_key, candidate_id, result):
        previous = self.get_previous_deployment(candidate_id)
        previous_id = previous["candidate_id"]
        with self.engine.begin() as connection:
            connection.execute(
                text(
                    "UPDATE public.dataflow_workflow_versions "
                    "SET status = 'superseded', "
                    "deployment_metadata = deployment_metadata || "
                    '\'{"gateway_status":"rolled_back"}\'::jsonb, '
                    "updated_at = CURRENT_TIMESTAMP "
                    "WHERE id = CAST(:id AS uuid)"
                ),
                {"id": candidate_id},
            )
            connection.execute(
                text(
                    "UPDATE public.workflow_engine_bindings "
                    "SET role = 'archived', status = 'disabled', "
                    "updated_at = CURRENT_TIMESTAMP "
                    "WHERE workflow_version_id = CAST(:id AS uuid)"
                ),
                {"id": candidate_id},
            )
            connection.execute(
                text(
                    "UPDATE public.workflow_schedules "
                    "SET status = 'retired', updated_at = CURRENT_TIMESTAMP "
                    "WHERE workflow_version_id = CAST(:id AS uuid)"
                ),
                {"id": candidate_id},
            )
            connection.execute(
                text(
                    "UPDATE public.dataflow_workflow_versions "
                    "SET status = 'active', "
                    "deployment_metadata = deployment_metadata || "
                    '\'{"gateway_status":"promoted"}\'::jsonb, '
                    "updated_at = CURRENT_TIMESTAMP "
                    "WHERE id = CAST(:id AS uuid)"
                ),
                {"id": previous_id},
            )
            connection.execute(
                text(
                    "UPDATE public.workflow_engine_bindings "
                    "SET role = 'primary', status = 'enabled', "
                    "updated_at = CURRENT_TIMESTAMP "
                    "WHERE workflow_version_id = CAST(:id AS uuid)"
                ),
                {"id": previous_id},
            )
            connection.execute(
                text(
                    "UPDATE public.workflow_schedules "
                    "SET status = 'active', updated_at = CURRENT_TIMESTAMP "
                    "WHERE workflow_version_id = CAST(:id AS uuid)"
                ),
                {"id": previous_id},
            )
            self._complete_operation(
                connection,
                "rollback_to_previous_version",
                idempotency_key,
                candidate_id,
                result,
            )
        return result

    def release_rollback(self, idempotency_key, candidate_id):
        self._release_operation(
            "rollback_to_previous_version", idempotency_key, candidate_id
        )

    def _release_operation(self, action, idempotency_key, candidate_id):
        with self.engine.begin() as connection:
            connection.execute(
                text(
                    "DELETE FROM public.workflow_gateway_operations "
                    "WHERE action = :action AND idempotency_key = :key "
                    "AND workflow_version_id = CAST(:version_id AS uuid) "
                    "AND status = 'claimed'"
                ),
                {
                    "action": action,
                    "key": idempotency_key,
                    "version_id": candidate_id,
                },
            )


class PostgresContextRepository:
    """Secret-free PostgreSQL read model for the Context MCP server."""

    def __init__(self, engine, *, max_concurrency=10, clock=None):
        self.engine = engine
        self.max_concurrency = int(max_concurrency)
        self.clock = clock or (lambda: datetime.now(UTC))
        if self.max_concurrency < 1 or self.max_concurrency > 1000:
            raise ValueError("context max concurrency is invalid")

    @staticmethod
    def _dataflow_row(row):
        if row is None:
            return None
        value = dict(row)
        spec = dict(value.pop("workflow_spec") or {})
        return {
            "uid": value["uid"],
            "name": spec.get("name"),
            "description": spec.get("description"),
            "business_domain": value.get("business_domain"),
            "status": value.get("status"),
        }

    def list_dataflows(self):
        with self.engine.connect() as connection:
            rows = (
                connection.execute(
                    text(
                        "SELECT DISTINCT ON (dataflow_uid) "
                        "dataflow_uid::text AS uid, business_domain, status, "
                        "workflow_spec "
                        "FROM public.dataflow_workflow_versions "
                        "WHERE engine_type = 'kestra' "
                        "ORDER BY dataflow_uid, version_no DESC"
                    )
                )
                .mappings()
                .all()
            )
        return [self._dataflow_row(row) for row in rows]

    def describe_dataflow(self, dataflow_uid):
        with self.engine.connect() as connection:
            row = (
                connection.execute(
                    text(
                        "SELECT dataflow_uid::text AS uid, business_domain, "
                        "status, workflow_spec "
                        "FROM public.dataflow_workflow_versions "
                        "WHERE dataflow_uid = CAST(:uid AS uuid) "
                        "AND engine_type = 'kestra' "
                        "ORDER BY version_no DESC LIMIT 1"
                    ),
                    {"uid": dataflow_uid},
                )
                .mappings()
                .one_or_none()
            )
        if row is None:
            raise KeyError(dataflow_uid)
        return self._dataflow_row(row)

    def _latest_spec(self, dataflow_uid):
        with self.engine.connect() as connection:
            row = (
                connection.execute(
                    text(
                        "SELECT workflow_spec, business_domain "
                        "FROM public.dataflow_workflow_versions "
                        "WHERE dataflow_uid = CAST(:uid AS uuid) "
                        "AND engine_type = 'kestra' "
                        "ORDER BY version_no DESC LIMIT 1"
                    ),
                    {"uid": dataflow_uid},
                )
                .mappings()
                .one_or_none()
            )
        if row is None:
            raise KeyError(dataflow_uid)
        return dict(row["workflow_spec"] or {}), row["business_domain"]

    def get_dataflow_dependencies(self, dataflow_uid):
        spec, domain = self._latest_spec(dataflow_uid)
        dependencies = []
        seen = set()
        for node in spec.get("nodes") or []:
            if node.get("type") != "subflow":
                continue
            dependency_uid = (node.get("config") or {}).get("dataflow_uid")
            if not dependency_uid or dependency_uid in seen:
                continue
            seen.add(dependency_uid)
            try:
                dependency = self.describe_dataflow(dependency_uid)
            except (KeyError, ValueError):
                continue
            dependencies.append(
                {
                    "uid": dependency["uid"],
                    "name": dependency["name"],
                    "relation": "subflow",
                    "business_domain": domain,
                }
            )
        return dependencies

    def get_data_lineage(self, dataflow_uid, depth):
        del depth
        spec, _ = self._latest_spec(dataflow_uid)
        nodes = {node.get("id"): node for node in spec.get("nodes") or []}
        return [
            {
                "from_uid": (
                    nodes.get(edge.get("from"), {}).get("data_source_uid")
                    or f"{dataflow_uid}:{edge.get('from')}"
                ),
                "to_uid": (
                    nodes.get(edge.get("to"), {}).get("data_source_uid")
                    or f"{dataflow_uid}:{edge.get('to')}"
                ),
                "relation": "workflow_edge",
                "description": edge.get("condition"),
            }
            for edge in spec.get("edges") or []
        ]

    def list_datasource_capabilities(self, business_domain):
        with self.engine.connect() as connection:
            rows = (
                connection.execute(
                    text(
                        "WITH latest AS ("
                        "SELECT DISTINCT ON (dataflow_uid) workflow_spec "
                        "FROM public.dataflow_workflow_versions "
                        "WHERE engine_type = 'kestra' "
                        "AND business_domain = :domain "
                        "ORDER BY dataflow_uid, version_no DESC"
                        "), nodes AS ("
                        "SELECT jsonb_array_elements("
                        "COALESCE(workflow_spec->'nodes', '[]'::jsonb)) AS node "
                        "FROM latest"
                        "), sources AS ("
                        "SELECT node->>'data_source_uid' AS uid, "
                        "array_agg(DISTINCT COALESCE("
                        "node->>'purpose', 'read')) AS purposes "
                        "FROM nodes WHERE node ? 'data_source_uid' "
                        "GROUP BY node->>'data_source_uid'"
                        ") SELECT sources.uid, sources.purposes, "
                        "credentials.status AS credential_status "
                        "FROM sources LEFT JOIN "
                        "public.datasource_credentials credentials "
                        "ON credentials.data_source_uid = "
                        "CAST(sources.uid AS uuid) "
                        "AND credentials.status = 'active'"
                    ),
                    {"domain": business_domain},
                )
                .mappings()
                .all()
            )
        return [
            {
                "uid": row["uid"],
                "type": "external",
                "purposes": list(row["purposes"] or []),
                "health": (
                    "configured"
                    if row["credential_status"] == "active"
                    else "credential_unavailable"
                ),
                "capacity": {"available": row["credential_status"] == "active"},
            }
            for row in rows
        ]

    def get_datasource_pool_health(self, data_source_uid):
        with self.engine.connect() as connection:
            configured = bool(
                connection.execute(
                    text(
                        "SELECT EXISTS("
                        "SELECT 1 FROM public.datasource_credentials "
                        "WHERE data_source_uid = CAST(:uid AS uuid) "
                        "AND status = 'active')"
                    ),
                    {"uid": data_source_uid},
                ).scalar_one()
            )
        return {
            "uid": data_source_uid,
            "health": "configured" if configured else "credential_unavailable",
            "capacity": {"available": configured},
        }

    def get_execution_history(self, dataflow_uid, limit, days):
        with self.engine.connect() as connection:
            rows = (
                connection.execute(
                    text(
                        "SELECT run.status, run.correlation_id::text "
                        "AS correlation_id, run.started_at, run.finished_at, "
                        "run.run_metadata->>'safe_error' AS safe_error "
                        "FROM public.workflow_runs run "
                        "JOIN public.dataflow_workflow_versions version "
                        "ON version.id = run.workflow_version_id "
                        "WHERE version.dataflow_uid = CAST(:uid AS uuid) "
                        "AND run.created_at >= "
                        "CURRENT_TIMESTAMP - make_interval(days => :days) "
                        "ORDER BY run.created_at DESC LIMIT :limit"
                    ),
                    {
                        "uid": dataflow_uid,
                        "days": int(days),
                        "limit": int(limit),
                    },
                )
                .mappings()
                .all()
            )
        return [
            {
                **dict(row),
                "started_at": (
                    row["started_at"].isoformat()
                    if row["started_at"] is not None
                    else None
                ),
                "finished_at": (
                    row["finished_at"].isoformat()
                    if row["finished_at"] is not None
                    else None
                ),
            }
            for row in rows
        ]

    def get_sla_constraints(self, dataflow_uid):
        with self.engine.connect() as connection:
            plan = connection.execute(
                text(
                    "SELECT schedule_plan "
                    "FROM public.dataflow_workflow_versions "
                    "WHERE dataflow_uid = CAST(:uid AS uuid) "
                    "AND engine_type = 'kestra' "
                    "ORDER BY version_no DESC LIMIT 1"
                ),
                {"uid": dataflow_uid},
            ).scalar_one_or_none()
        if plan is None:
            raise KeyError(dataflow_uid)
        plan = dict(plan)
        return {
            "timezone": plan.get("timezone"),
            "max_duration_seconds": plan.get("timeout_seconds"),
        }

    def estimate_schedule_capacity(self, business_domain, schedule_plan):
        with self.engine.connect() as connection:
            running = int(
                connection.execute(
                    text(
                        "SELECT COUNT(*) FROM public.workflow_runs run "
                        "JOIN public.dataflow_workflow_versions version "
                        "ON version.id = run.workflow_version_id "
                        "WHERE version.business_domain = :domain "
                        "AND run.status IN ('queued', 'running')"
                    ),
                    {"domain": business_domain},
                ).scalar_one()
            )
        available = max(0, self.max_concurrency - running)
        required = int(schedule_plan.get("max_concurrency", 1))
        return {
            "feasible": required <= available,
            "required_slots": required,
            "available_slots": available,
            "warnings": (
                []
                if required <= available
                else ["requested concurrency exceeds current capacity"]
            ),
        }

    def simulate_schedule(self, business_domain, schedule_plan):
        del business_domain
        now = self.clock()
        timezone = ZoneInfo(schedule_plan.get("timezone") or "UTC")
        local_now = now.astimezone(timezone)
        next_runs = []
        warnings = []
        for trigger in schedule_plan.get("triggers") or []:
            kind = trigger.get("type")
            if kind == "cron":
                expression = trigger.get("expression", "")
                if len(str(expression).split()) != 5:
                    warnings.append("simulation supports five-field cron only")
                    continue
                iterator = croniter(expression, local_now)
                next_runs.extend(
                    iterator.get_next(datetime).isoformat() for _ in range(5)
                )
            elif kind == "at":
                next_runs.append(str(trigger.get("at")))
            elif kind in {"manual", "event"}:
                warnings.append(f"{kind} trigger has no deterministic next run")
        return {
            "next_runs": sorted(next_runs)[:20],
            "warnings": warnings,
        }

    def compare_execution_results(self, baseline_execution_id, candidate_execution_id):
        with self.engine.connect() as connection:
            rows = (
                connection.execute(
                    text(
                        "SELECT engine_execution_id, status, started_at, "
                        "finished_at, run_metadata "
                        "FROM public.workflow_runs "
                        "WHERE engine_type = 'kestra' "
                        "AND engine_execution_id IN (:baseline, :candidate)"
                    ),
                    {
                        "baseline": baseline_execution_id,
                        "candidate": candidate_execution_id,
                    },
                )
                .mappings()
                .all()
            )
        by_id = {row["engine_execution_id"]: dict(row) for row in rows}
        if baseline_execution_id not in by_id or candidate_execution_id not in by_id:
            raise KeyError("execution comparison input is unavailable")
        baseline = by_id[baseline_execution_id]
        candidate = by_id[candidate_execution_id]

        def duration(row):
            if row["started_at"] is None or row["finished_at"] is None:
                return None
            return (row["finished_at"] - row["started_at"]).total_seconds()

        baseline_duration = duration(baseline)
        candidate_duration = duration(candidate)
        baseline_metadata = dict(baseline["run_metadata"] or {})
        candidate_metadata = dict(candidate["run_metadata"] or {})
        baseline_hash = baseline_metadata.get("output_hash")
        candidate_hash = candidate_metadata.get("output_hash")
        hash_match = (
            baseline_hash == candidate_hash
            if baseline_hash is not None and candidate_hash is not None
            else None
        )
        duration_delta = (
            candidate_duration - baseline_duration
            if baseline_duration is not None and candidate_duration is not None
            else None
        )
        equivalent = (
            baseline["status"] == "success"
            and candidate["status"] == "success"
            and hash_match is not False
        )
        return {
            "equivalent": equivalent,
            "metrics": {
                "duration_delta_seconds": duration_delta,
                "output_hash_match": hash_match,
            },
            "differences": (
                [] if equivalent else ["execution status or output hash differs"]
            ),
        }


class PostgresAuditSink:
    def __init__(self, engine):
        self.engine = engine

    def record(self, event):
        detail = dict(event.get("detail") or {})
        candidate_id = detail.get("candidate_id")
        try:
            workflow_version_id = str(uuid.UUID(candidate_id)) if candidate_id else None
            correlation_id = str(uuid.UUID(event["correlation_id"]))
        except (TypeError, ValueError) as exc:
            raise ValueError("audit identifiers must be UUIDs") from exc
        decision = event.get("decision")
        if decision not in {
            "allowed",
            "rejected",
            "failed",
            "recorded",
            "idempotent_replay",
        }:
            raise ValueError("unsupported audit decision")
        bounded_detail = {
            **detail,
            "roles": list(event.get("roles") or [])[:20],
            "business_domain": event.get("business_domain"),
            "environment": event.get("environment"),
        }
        encoded_detail = _json(bounded_detail)
        if len(encoded_detail.encode("utf-8")) > 1048576:
            raise ValueError("audit detail exceeds size limit")

        def optional_string(name, maximum):
            value = event.get(name)
            return (
                _required_string(value, f"audit {name}", maximum)
                if value is not None
                else None
            )

        def optional_hash(value, label):
            if value is None:
                return None
            normalized = _required_string(value, label, 64).lower()
            if len(normalized) != 64 or any(
                character not in "0123456789abcdef" for character in normalized
            ):
                raise ValueError(f"{label} must be a SHA-256 hex digest")
            return normalized

        schema_version = _required_string(
            event.get("schema_version", "v53"),
            "audit schema_version",
            40,
        )
        context_hash = optional_hash(
            event.get("context_hash"),
            "audit context_hash",
        )
        candidate_hash = optional_hash(
            detail.get("workflow_hash"),
            "audit candidate_hash",
        )
        with self.engine.begin() as connection:
            connection.execute(
                text(
                    "INSERT INTO public.workflow_plan_audits "
                    "(id, workflow_version_id, actor_uid, actor_subject, action, "
                    "model_provider, model_name, prompt_version, schema_version, "
                    "context_hash, candidate_hash, decision, decision_detail, "
                    "correlation_id) "
                    "VALUES (CAST(:id AS uuid), "
                    "CAST(:workflow_version_id AS uuid), NULL, :actor_subject, "
                    ":action, :model_provider, :model_name, :prompt_version, "
                    ":schema_version, :context_hash, :candidate_hash, :decision, "
                    "CAST(:detail AS jsonb), CAST(:correlation_id AS uuid))"
                ),
                {
                    "id": new_governance_uid(),
                    "workflow_version_id": workflow_version_id,
                    "actor_subject": _required_string(
                        event.get("subject"), "audit subject", 200
                    ),
                    "action": _required_string(event.get("action"), "audit action", 80),
                    "model_provider": optional_string("model_provider", 80),
                    "model_name": optional_string("model_name", 120),
                    "prompt_version": optional_string("prompt_version", 80),
                    "schema_version": schema_version,
                    "context_hash": context_hash,
                    "candidate_hash": candidate_hash,
                    "decision": decision,
                    "detail": encoded_detail,
                    "correlation_id": correlation_id,
                },
            )