"""PostgreSQL persistence for durable operational tasks.""" from __future__ import annotations from contextlib import asynccontextmanager from typing import Any from psycopg import AsyncConnection from psycopg.errors import UndefinedTable from psycopg.types.json import Jsonb from uuid import uuid4 from app.domain.domain_errors import DomainError class OperationalTaskRepository: def __init__(self, connection: AsyncConnection, project_id): self.connection = connection self.project_id = str(project_id) @asynccontextmanager async def transaction(self): try: async with self.connection.transaction(): yield self except UndefinedTable as exc: raise DomainError( "DEPENDENCY_UNAVAILABLE", "Project operational-task storage is not provisioned", status=503, ) from exc async def lock(self, key: str) -> None: await self.connection.execute( "SELECT pg_advisory_xact_lock(hashtextextended(%s, 0))", (f"tjwater-operational:{self.project_id}:{key}",), ) async def get_request(self, scope: str, request_id: str) -> dict | None: cursor = await self.connection.execute( """SELECT request_digest,response FROM workbench.operational_task_requests WHERE project_id=%s AND scope=%s AND request_id=%s""", (self.project_id, scope, request_id), ) row = await cursor.fetchone() return dict(row) if row else None async def save_request( self, scope: str, request_id: str, request_digest: str, response: dict ) -> None: cursor = await self.connection.execute( """INSERT INTO workbench.operational_task_requests (project_id,scope,request_id,request_digest,response) VALUES (%s,%s,%s,%s,%s) ON CONFLICT(project_id,scope,request_id) DO UPDATE SET response=excluded.response WHERE operational_task_requests.request_digest=excluded.request_digest""", (self.project_id, scope, request_id, request_digest, Jsonb(response)), ) if cursor.rowcount != 1: raise DomainError( "IDEMPOTENCY_CONFLICT", "request_id was used with different content", status=409, ) async def get_task(self, task_id: str) -> dict | None: cursor = await self.connection.execute( """SELECT document FROM workbench.operational_tasks WHERE project_id=%s AND task_id=%s""", (self.project_id, task_id), ) row = await cursor.fetchone() return row["document"] if row else None async def get_active_run(self, task_id: str) -> dict | None: cursor = await self.connection.execute( """SELECT run_id,task_id,generation,execution_request_id,status,revision, agent_session_id,lease_expires_at,heartbeat_at,started_at,finished_at, failure_code,failure_detail,dispatch_attempts,next_dispatch_at, created_at,updated_at FROM workbench.operational_task_runs WHERE project_id=%s AND task_id=%s AND status IN ('queued','claimed','running') ORDER BY generation DESC LIMIT 1""", (self.project_id, task_id), ) row = await cursor.fetchone() return dict(row) if row else None async def get_run(self, run_id: str) -> dict | None: cursor = await self.connection.execute( """SELECT run_id,task_id,generation,execution_request_id,status,revision, agent_session_id,lease_expires_at,heartbeat_at,started_at,finished_at, failure_code,failure_detail,dispatch_attempts,next_dispatch_at, created_at,updated_at FROM workbench.operational_task_runs WHERE project_id=%s AND run_id=%s""", (self.project_id, run_id), ) row = await cursor.fetchone() return dict(row) if row else None async def insert_run(self, run: dict) -> None: await self.connection.execute( """INSERT INTO workbench.operational_task_runs (project_id,run_id,task_id,generation,execution_request_id,status,revision, agent_session_id,lease_expires_at,heartbeat_at,started_at,finished_at, failure_code,failure_detail,dispatch_attempts,next_dispatch_at, created_at,updated_at) VALUES (%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s)""", ( self.project_id, run["run_id"], run["task_id"], run["generation"], run["execution_request_id"], run["status"], run["revision"], run.get("agent_session_id"), run.get("lease_expires_at"), run.get("heartbeat_at"), run.get("started_at"), run.get("finished_at"), run.get("failure_code"), run.get("failure_detail"), run.get("dispatch_attempts", 0), run["next_dispatch_at"], run["created_at"], run["updated_at"], ), ) async def update_run(self, run: dict, expected_revision: int) -> None: cursor = await self.connection.execute( """UPDATE workbench.operational_task_runs SET status=%s,revision=%s,agent_session_id=%s,lease_expires_at=%s, heartbeat_at=%s,started_at=%s,finished_at=%s,failure_code=%s, failure_detail=%s,dispatch_attempts=%s,next_dispatch_at=%s,updated_at=%s WHERE project_id=%s AND run_id=%s AND revision=%s""", ( run["status"], run["revision"], run.get("agent_session_id"), run.get("lease_expires_at"), run.get("heartbeat_at"), run.get("started_at"), run.get("finished_at"), run.get("failure_code"), run.get("failure_detail"), run.get("dispatch_attempts", 0), run["next_dispatch_at"], run["updated_at"], self.project_id, run["run_id"], expected_revision, ), ) if cursor.rowcount != 1: raise DomainError("VERSION_CONFLICT", "Run revision changed", status=409) async def list_runs( self, *, status: str | None, limit: int, cursor_created_at=None, cursor_run_id=None ) -> list[dict]: result = await self.connection.execute( """SELECT run_id,task_id,generation,execution_request_id,status,revision, agent_session_id,lease_expires_at,heartbeat_at,started_at,finished_at, failure_code,failure_detail,dispatch_attempts,next_dispatch_at, created_at,updated_at FROM workbench.operational_task_runs WHERE project_id=%s AND (%s::text IS NULL OR status=%s::text) AND (%s::timestamptz IS NULL OR (created_at,run_id)>(%s::timestamptz,%s::uuid)) AND (status<>'queued' OR next_dispatch_at<=now()) ORDER BY created_at,run_id LIMIT %s""", ( self.project_id, status, status, cursor_created_at, cursor_created_at, cursor_run_id, limit, ), ) return [dict(row) for row in await result.fetchall()] async def claim_expired_runs(self, limit: int) -> list[dict]: cursor = await self.connection.execute( """SELECT run_id,task_id,generation,execution_request_id,status,revision, agent_session_id,lease_expires_at,heartbeat_at,started_at,finished_at, failure_code,failure_detail,dispatch_attempts,next_dispatch_at, created_at,updated_at FROM workbench.operational_task_runs WHERE project_id=%s AND status IN ('claimed','running') AND lease_expires_at IS NOT NULL AND lease_expires_at<=now() ORDER BY lease_expires_at,run_id FOR UPDATE SKIP LOCKED LIMIT %s""", (self.project_id, limit), ) return [dict(row) for row in await cursor.fetchall()] async def consecutive_timed_out_runs(self, task_id: str) -> int: cursor = await self.connection.execute( """SELECT status FROM workbench.operational_task_runs WHERE project_id=%s AND task_id=%s ORDER BY generation DESC""", (self.project_id, task_id), ) count = 0 for row in await cursor.fetchall(): if row["status"] != "timed_out": break count += 1 return count async def append_event( self, *, topic: str, aggregate_type: str, aggregate_id: str, aggregate_revision: int, occurred_at, metadata: dict | None = None, ) -> int: event_id = uuid4() cursor = await self.connection.execute( """INSERT INTO workbench.domain_event_log (project_id,event_id,topic,aggregate_type,aggregate_id, aggregate_revision,occurred_at,metadata) VALUES (%s,%s,%s,%s,%s,%s,%s,%s) RETURNING event_seq""", ( self.project_id, event_id, topic, aggregate_type, aggregate_id, aggregate_revision, occurred_at, Jsonb(metadata or {}), ), ) row = await cursor.fetchone() event_seq = int(row["event_seq"]) await self.connection.execute( "SELECT pg_notify('workbench_events', %s)", (f'{self.project_id}:{event_seq}',), ) return event_seq async def get_task_by_create_request(self, request_id: str) -> dict | None: cursor = await self.connection.execute( """SELECT document FROM workbench.operational_tasks WHERE project_id=%s AND create_request_id=%s""", (self.project_id, request_id), ) row = await cursor.fetchone() return row["document"] if row else None async def get_task_by_external_source( self, source_kind: str, external_id: str ) -> dict | None: cursor = await self.connection.execute( """SELECT document FROM workbench.operational_tasks WHERE project_id=%s AND document#>>'{source,kind}'=%s AND document#>>'{source,external_id}'=%s""", (self.project_id, source_kind, external_id), ) row = await cursor.fetchone() return row["document"] if row else None async def list_tasks( self, *, task_type: str | None, status: str | None, limit: int, offset: int, ) -> list[dict]: cursor = await self.connection.execute( """SELECT task.document, CASE WHEN active_run.run_id IS NULL THEN NULL ELSE to_jsonb(active_run) - 'project_id' END AS active_run FROM workbench.operational_tasks AS task LEFT JOIN LATERAL ( SELECT run.* FROM workbench.operational_task_runs AS run WHERE run.project_id=task.project_id AND run.task_id=task.task_id AND run.status IN ('queued','claimed','running') ORDER BY run.generation DESC LIMIT 1 ) AS active_run ON true WHERE task.project_id=%s AND (%s::text IS NULL OR task.task_type=%s::text) AND (%s::text IS NULL OR task.status=%s::text) ORDER BY task.updated_at DESC,task.task_id LIMIT %s OFFSET %s""", ( self.project_id, task_type, task_type, status, status, limit, offset, ), ) return [ {**row["document"], "active_run": row["active_run"]} for row in await cursor.fetchall() ] async def insert_task(self, task: dict, create_request_id: str) -> None: await self.connection.execute( """INSERT INTO workbench.operational_tasks (project_id,task_id,create_request_id,task_type,workflow,status,lifecycle_status,priority, revision,execution_generation,execution_request_id,agent_session_id, agent_lease_expires_at,current_plan_revision,current_report_revision, scheduled_for,created_at,updated_at,document) VALUES (%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s)""", self._task_params(task, create_request_id), ) async def update_task(self, task: dict) -> None: cursor = await self.connection.execute( """UPDATE workbench.operational_tasks SET status=%s,lifecycle_status=%s,revision=%s,execution_generation=%s,execution_request_id=%s, agent_session_id=%s,agent_lease_expires_at=%s,current_plan_revision=%s, current_report_revision=%s,updated_at=%s,document=%s WHERE project_id=%s AND task_id=%s AND revision=%s""", ( task["status"], task["lifecycle_status"], task["revision"], task["execution_generation"], task["execution_request_id"], task.get("agent_session_id"), task.get("agent_lease_expires_at"), task.get("current_plan_revision"), task.get("current_report_revision"), task["updated_at"], Jsonb(task), self.project_id, task["task_id"], task["revision"] - 1, ), ) if cursor.rowcount != 1: raise DomainError("VERSION_CONFLICT", "Task revision changed", status=409) def _task_params(self, task: dict, create_request_id: str) -> tuple: return ( self.project_id, task["task_id"], create_request_id, task["task_type"], task["workflow"], task["status"], task["lifecycle_status"], task["priority"], task["revision"], task["execution_generation"], task["execution_request_id"], task.get("agent_session_id"), task.get("agent_lease_expires_at"), task.get("current_plan_revision"), task.get("current_report_revision"), task.get("scheduled_for"), task["created_at"], task["updated_at"], Jsonb(task), ) async def get_evidence(self, task_id: str, evidence_id: str) -> dict | None: cursor = await self.connection.execute( """SELECT document FROM workbench.operational_task_evidence WHERE project_id=%s AND task_id=%s AND evidence_id=%s""", (self.project_id, task_id, evidence_id), ) row = await cursor.fetchone() return row["document"] if row else None async def list_evidence(self, task_id: str) -> list[dict]: cursor = await self.connection.execute( """SELECT document FROM workbench.operational_task_evidence WHERE project_id=%s AND task_id=%s ORDER BY created_at,evidence_id""", (self.project_id, task_id), ) return [row["document"] for row in await cursor.fetchall()] async def insert_evidence(self, task_id: str, evidence: dict) -> None: await self.connection.execute( """INSERT INTO workbench.operational_task_evidence (project_id,task_id,evidence_id,document,created_at) VALUES (%s,%s,%s,%s,%s)""", (self.project_id, task_id, evidence["evidence_id"], Jsonb(evidence), evidence["recorded_at"]), ) async def get_latest_plan(self, task_id: str) -> dict | None: return await self._latest_artifact("operational_task_plans", task_id) async def get_latest_report(self, task_id: str) -> dict | None: return await self._latest_artifact("operational_task_reports", task_id) async def _latest_artifact(self, table: str, task_id: str) -> dict | None: if table not in {"operational_task_plans", "operational_task_reports"}: raise ValueError("unsupported artifact table") cursor = await self.connection.execute( f"""SELECT document FROM workbench.{table} WHERE project_id=%s AND task_id=%s ORDER BY revision DESC LIMIT 1""", (self.project_id, task_id), ) row = await cursor.fetchone() return row["document"] if row else None async def insert_plan(self, task_id: str, snapshot: dict) -> None: await self._insert_artifact("operational_task_plans", task_id, snapshot) async def insert_report(self, task_id: str, snapshot: dict) -> None: await self._insert_artifact("operational_task_reports", task_id, snapshot) async def _insert_artifact(self, table: str, task_id: str, snapshot: dict) -> None: if table not in {"operational_task_plans", "operational_task_reports"}: raise ValueError("unsupported artifact table") await self.connection.execute( f"""INSERT INTO workbench.{table} (project_id,task_id,revision,document,created_at) VALUES (%s,%s,%s,%s,%s)""", (self.project_id, task_id, snapshot["revision"], Jsonb(snapshot), snapshot["published_at"]), ) async def get_action(self, task_id: str, action_id: str) -> dict | None: cursor = await self.connection.execute( """SELECT document FROM workbench.operational_task_actions WHERE project_id=%s AND task_id=%s AND action_id=%s""", (self.project_id, task_id, action_id), ) row = await cursor.fetchone() return row["document"] if row else None async def list_actions(self, task_id: str) -> list[dict]: cursor = await self.connection.execute( """SELECT document FROM workbench.operational_task_actions WHERE project_id=%s AND task_id=%s ORDER BY created_at,action_id""", (self.project_id, task_id), ) return [row["document"] for row in await cursor.fetchall()] async def insert_action(self, task_id: str, action: dict) -> None: await self.connection.execute( """INSERT INTO workbench.operational_task_actions (project_id,task_id,action_id,version,approval_status,execution_status, document,created_at,updated_at) VALUES (%s,%s,%s,%s,%s,%s,%s,%s,%s)""", (self.project_id, task_id, action["action_id"], action["version"], action["approval_status"], action["execution_status"], Jsonb(action), action["proposed_at"], action["proposed_at"]), ) await self._insert_action_revision(task_id, action) async def update_action(self, task_id: str, action: dict) -> None: await self._insert_action_revision(task_id, action) await self.connection.execute( """UPDATE workbench.operational_task_actions SET version=%s,approval_status=%s,execution_status=%s,document=%s,updated_at=%s WHERE project_id=%s AND task_id=%s AND action_id=%s""", (action["version"], action["approval_status"], action["execution_status"], Jsonb(action), action.get("executed_at") or action.get("decided_at") or action["proposed_at"], self.project_id, task_id, action["action_id"]), ) async def _insert_action_revision(self, task_id: str, action: dict) -> None: await self.connection.execute( """INSERT INTO workbench.operational_task_action_revisions (project_id,task_id,action_id,version,document,recorded_at) VALUES (%s,%s,%s,%s,%s,%s)""", ( self.project_id, task_id, action["action_id"], action["version"], Jsonb(action), action.get("executed_at") or action.get("decided_at") or action["proposed_at"], ), ) async def get_feedback(self, task_id: str, source_event_id: str) -> dict | None: cursor = await self.connection.execute( """SELECT source_revision,document,accepted_at FROM workbench.operational_task_external_feedback WHERE project_id=%s AND task_id=%s AND source_event_id=%s""", (self.project_id, task_id, source_event_id), ) row = await cursor.fetchone() return dict(row) if row else None async def upsert_feedback(self, task_id: str, feedback: dict, accepted_at) -> None: await self.connection.execute( """INSERT INTO workbench.operational_task_external_feedback (project_id,task_id,source_event_id,source_revision,document,accepted_at) VALUES (%s,%s,%s,%s,%s,%s) ON CONFLICT(project_id,task_id,source_event_id) DO UPDATE SET source_revision=excluded.source_revision,document=excluded.document, accepted_at=excluded.accepted_at""", (self.project_id, task_id, feedback["source_event_id"], feedback["source_revision"], Jsonb(feedback), accepted_at), ) async def insert_notification(self, notification: dict) -> None: await self.connection.execute( """INSERT INTO workbench.operational_task_notifications (project_id,notification_id,task_id,created_at,document) VALUES (%s,%s,%s,%s,%s)""", (self.project_id, notification["notification_id"], notification["task_id"], notification["created_at"], Jsonb(notification)), ) async def get_notification(self, notification_id: str, user_id: str) -> dict | None: cursor = await self.connection.execute( """SELECT n.document || jsonb_build_object( 'status',CASE WHEN r.read_at IS NULL THEN 'unread' ELSE 'read' END, 'version',COALESCE(r.version,1), 'read_at',r.read_at,'read_by', CASE WHEN r.user_id IS NULL THEN NULL ELSE 'principal:' || r.user_id::text END ) AS document FROM workbench.operational_task_notifications n LEFT JOIN workbench.operational_task_notification_reads r ON r.project_id=n.project_id AND r.notification_id=n.notification_id AND r.user_id=%s WHERE n.project_id=%s AND n.notification_id=%s""", (user_id, self.project_id, notification_id), ) row = await cursor.fetchone() return row["document"] if row else None async def mark_notification_read( self, notification_id: str, user_id: str, version: int, read_at ) -> None: await self.connection.execute( """INSERT INTO workbench.operational_task_notification_reads (project_id,notification_id,user_id,version,read_at) VALUES (%s,%s,%s,%s,%s) ON CONFLICT(project_id,notification_id,user_id) DO UPDATE SET version=excluded.version,read_at=excluded.read_at""", (self.project_id, notification_id, user_id, version, read_at), ) async def list_notifications( self, *, user_id: str, allowed_permissions: list[str], unread_only: bool, limit: int, offset: int ) -> list[dict]: cursor = await self.connection.execute( """SELECT n.document || jsonb_build_object( 'status',CASE WHEN r.read_at IS NULL THEN 'unread' ELSE 'read' END, 'version',COALESCE(r.version,1), 'read_at',r.read_at,'read_by', CASE WHEN r.user_id IS NULL THEN NULL ELSE 'principal:' || r.user_id::text END ) AS document FROM workbench.operational_task_notifications n LEFT JOIN workbench.operational_task_notification_reads r ON r.project_id=n.project_id AND r.notification_id=n.notification_id AND r.user_id=%s WHERE n.project_id=%s AND (%s=false OR r.read_at IS NULL) AND (n.document->>'required_permission' IS NULL OR n.document->>'required_permission'=ANY(%s)) ORDER BY n.created_at DESC,n.notification_id LIMIT %s OFFSET %s""", (user_id, self.project_id, unread_only, allowed_permissions, limit, offset), ) return [row["document"] for row in await cursor.fetchall()] async def get_schedule_by_request(self, request_id: str) -> dict | None: cursor = await self.connection.execute( """SELECT document FROM workbench.operational_task_schedules WHERE project_id=%s AND create_request_id=%s""", (self.project_id, request_id), ) row = await cursor.fetchone() return row["document"] if row else None async def insert_schedule(self, schedule: dict, request_id: str) -> None: await self.connection.execute( """INSERT INTO workbench.operational_task_schedules (project_id,schedule_id,create_request_id,schedule_key,next_run_at, interval_seconds,enabled,version,document,created_at,updated_at) VALUES (%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s)""", (self.project_id, schedule["schedule_id"], request_id, schedule["schedule_key"], schedule["next_run_at"], schedule["interval_seconds"], schedule["enabled"], schedule["version"], Jsonb(schedule), schedule["created_at"], schedule["updated_at"]), ) async def list_schedules(self) -> list[dict]: cursor = await self.connection.execute( """SELECT document FROM workbench.operational_task_schedules WHERE project_id=%s ORDER BY next_run_at,schedule_id""", (self.project_id,), ) return [row["document"] for row in await cursor.fetchall()] async def get_schedule(self, schedule_id: str) -> dict | None: cursor = await self.connection.execute( """SELECT document FROM workbench.operational_task_schedules WHERE project_id=%s AND schedule_id=%s""", (self.project_id, schedule_id), ) row = await cursor.fetchone() return row["document"] if row else None async def update_schedule(self, schedule: dict) -> None: cursor = await self.connection.execute( """UPDATE workbench.operational_task_schedules SET enabled=%s,version=%s,updated_at=%s,document=%s WHERE project_id=%s AND schedule_id=%s AND version=%s""", ( schedule["enabled"], schedule["version"], schedule["updated_at"], Jsonb(schedule), self.project_id, schedule["schedule_id"], schedule["version"] - 1, ), ) if cursor.rowcount != 1: raise DomainError("VERSION_CONFLICT", "Schedule version changed", status=409) async def claim_due_schedules(self, limit: int) -> list[dict]: cursor = await self.connection.execute( """SELECT document FROM workbench.operational_task_schedules WHERE project_id=%s AND enabled=true AND next_run_at <= now() ORDER BY next_run_at,schedule_id FOR UPDATE SKIP LOCKED LIMIT %s""", (self.project_id, limit), ) return [row["document"] for row in await cursor.fetchall()] async def advance_schedule(self, schedule: dict) -> None: await self.connection.execute( """UPDATE workbench.operational_task_schedules SET next_run_at=%s,version=%s,updated_at=%s,document=%s WHERE project_id=%s AND schedule_id=%s""", (schedule["next_run_at"], schedule["version"], schedule["updated_at"], Jsonb(schedule), self.project_id, schedule["schedule_id"]), ) async def insert_occurrence(self, schedule_id: str, due_at, task_id: str, created_at) -> bool: cursor = await self.connection.execute( """INSERT INTO workbench.operational_schedule_occurrences (project_id,schedule_id,due_at,task_id,created_at) VALUES (%s,%s,%s,%s,%s) ON CONFLICT(project_id,schedule_id,due_at) DO NOTHING RETURNING task_id""", (self.project_id, schedule_id, due_at, task_id, created_at), ) return await cursor.fetchone() is not None