"""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 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, }, )