|
|
@@ -0,0 +1,1317 @@
|
|
|
+"""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"),
|
|
|
+ }
|
|
|
+ with self.engine.begin() as connection:
|
|
|
+ connection.execute(
|
|
|
+ text(
|
|
|
+ "INSERT INTO public.workflow_plan_audits "
|
|
|
+ "(id, workflow_version_id, actor_uid, actor_subject, action, "
|
|
|
+ "schema_version, candidate_hash, decision, decision_detail, "
|
|
|
+ "correlation_id) "
|
|
|
+ "VALUES (CAST(:id AS uuid), "
|
|
|
+ "CAST(:workflow_version_id AS uuid), NULL, :actor_subject, "
|
|
|
+ ":action, 'v53', :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),
|
|
|
+ "candidate_hash": detail.get("workflow_hash"),
|
|
|
+ "decision": decision,
|
|
|
+ "detail": _json(bounded_detail),
|
|
|
+ "correlation_id": correlation_id,
|
|
|
+ },
|
|
|
+ )
|