feat(workbench): persist operational workflows and events
This commit is contained in:
@@ -53,18 +53,52 @@ class AnalysisRepository:
|
||||
status: str,
|
||||
created_by: str,
|
||||
parameters: dict[str, Any],
|
||||
) -> None:
|
||||
) -> int:
|
||||
with conn.cursor() as cur:
|
||||
cur.execute(
|
||||
"""
|
||||
UPDATE analysis.runs
|
||||
SET created_by = %s, status = %s, parameters = %s
|
||||
SET created_by = %s, status = %s, parameters = %s,
|
||||
revision = revision + 1
|
||||
WHERE run_id = %s
|
||||
RETURNING revision
|
||||
""",
|
||||
(created_by, status, Jsonb(parameters), run_id),
|
||||
)
|
||||
if cur.rowcount != 1:
|
||||
row = cur.fetchone()
|
||||
if row is None:
|
||||
raise LookupError(f"analysis run {run_id} does not exist")
|
||||
return int(row["revision"])
|
||||
|
||||
@staticmethod
|
||||
def append_run_changed_event_sync(
|
||||
conn: Connection,
|
||||
*,
|
||||
project_id: UUID,
|
||||
run_id: UUID,
|
||||
revision: int,
|
||||
status: str,
|
||||
occurred_at: datetime,
|
||||
results_committed: bool,
|
||||
) -> None:
|
||||
with conn.cursor() as cur:
|
||||
cur.execute(
|
||||
"""INSERT INTO workbench.domain_event_log
|
||||
(project_id,event_id,topic,aggregate_type,aggregate_id,
|
||||
aggregate_revision,occurred_at,metadata)
|
||||
VALUES (%s,%s,'analysis.run.changed','analysis_run',%s,%s,%s,%s)
|
||||
RETURNING event_seq""",
|
||||
(
|
||||
project_id, uuid4(), str(run_id), revision, occurred_at,
|
||||
Jsonb({"run_id": str(run_id), "status": status,
|
||||
"results_committed": results_committed}),
|
||||
),
|
||||
)
|
||||
row = cur.fetchone()
|
||||
cur.execute(
|
||||
"SELECT pg_notify('workbench_events', %s)",
|
||||
(f'{project_id}:{row["event_seq"]}',),
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def insert_result_sync(
|
||||
|
||||
@@ -0,0 +1,589 @@
|
||||
"""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
|
||||
@@ -6,6 +6,7 @@ from contextlib import asynccontextmanager
|
||||
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
|
||||
|
||||
@@ -86,3 +87,21 @@ class TaskAlertRepository:
|
||||
Jsonb({key: value for key, value in state.items() if key != "revision"}),
|
||||
),
|
||||
)
|
||||
|
||||
async def append_changed_event(self, state: dict) -> None:
|
||||
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,'workbench.alert.changed','workbench_alert',%s,%s,%s,%s)
|
||||
RETURNING event_seq""",
|
||||
(
|
||||
self.project_id, uuid4(), state["alert_id"], state["revision"],
|
||||
state["updated_at"], Jsonb({"alert_id": state["alert_id"]}),
|
||||
),
|
||||
)
|
||||
row = await cursor.fetchone()
|
||||
await self.connection.execute(
|
||||
"SELECT pg_notify('workbench_events', %s)",
|
||||
(f'{self.project_id}:{row["event_seq"]}',),
|
||||
)
|
||||
|
||||
@@ -0,0 +1,200 @@
|
||||
"""Per-process, project-scoped PostgreSQL LISTEN hubs for workbench SSE."""
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import contextlib
|
||||
import logging
|
||||
from contextlib import asynccontextmanager
|
||||
from uuid import UUID
|
||||
|
||||
from psycopg import AsyncConnection
|
||||
from psycopg_pool import AsyncConnectionPool
|
||||
from psycopg.rows import dict_row
|
||||
|
||||
from app.infra.db.postgresql.workbench_events import WorkbenchEventRepository
|
||||
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class WorkbenchEventHub:
|
||||
"""Own one LISTEN connection and fan wake-ups out to SSE subscribers."""
|
||||
|
||||
def __init__(self, project_id: UUID, dsn: str) -> None:
|
||||
self.project_id = project_id
|
||||
self.dsn = dsn
|
||||
self.subscribers = 0
|
||||
self._connection: AsyncConnection | None = None
|
||||
self._pool: AsyncConnectionPool | None = None
|
||||
self._connection_lock = asyncio.Lock()
|
||||
self._condition = asyncio.Condition()
|
||||
self._generation = 0
|
||||
self._ready = asyncio.Event()
|
||||
self._closed = False
|
||||
self._listener_task: asyncio.Task | None = None
|
||||
self._idle_task: asyncio.Task | None = None
|
||||
|
||||
async def start(self) -> None:
|
||||
if self._listener_task is None or self._listener_task.done():
|
||||
self._closed = False
|
||||
self._ready.clear()
|
||||
self._listener_task = asyncio.create_task(
|
||||
self._listen_loop(), name=f"workbench-events:{self.project_id}"
|
||||
)
|
||||
await asyncio.wait_for(self._ready.wait(), timeout=10)
|
||||
|
||||
async def _connect(self) -> AsyncConnection:
|
||||
if self._pool is None:
|
||||
self._pool = AsyncConnectionPool(
|
||||
conninfo=self.dsn,
|
||||
min_size=0,
|
||||
max_size=1,
|
||||
open=False,
|
||||
kwargs={"autocommit": True, "row_factory": dict_row},
|
||||
)
|
||||
await self._pool.open()
|
||||
connection = await self._pool.getconn()
|
||||
await connection.execute("LISTEN workbench_events")
|
||||
return connection
|
||||
|
||||
async def _listen_loop(self) -> None:
|
||||
retry_seconds = 1
|
||||
while not self._closed:
|
||||
try:
|
||||
if self._connection is None or self._connection.closed:
|
||||
async with self._connection_lock:
|
||||
if self._connection is None or self._connection.closed:
|
||||
self._connection = await self._connect()
|
||||
self._ready.set()
|
||||
retry_seconds = 1
|
||||
notified = False
|
||||
async with self._connection_lock:
|
||||
connection = self._connection
|
||||
if connection is None:
|
||||
continue
|
||||
async for notification in connection.notifies(timeout=1, stop_after=1):
|
||||
if notification.payload.startswith(f"{self.project_id}:"):
|
||||
notified = True
|
||||
break
|
||||
if notified:
|
||||
await self._wake()
|
||||
await asyncio.sleep(0)
|
||||
except asyncio.CancelledError:
|
||||
raise
|
||||
except Exception:
|
||||
logger.exception("Workbench event listener failed for project %s", self.project_id)
|
||||
await self._drop_connection()
|
||||
await self._wake()
|
||||
await asyncio.sleep(retry_seconds)
|
||||
retry_seconds = min(retry_seconds * 2, 30)
|
||||
|
||||
async def _wake(self) -> None:
|
||||
async with self._condition:
|
||||
self._generation += 1
|
||||
self._condition.notify_all()
|
||||
|
||||
async def wait_for_wake(self, generation: int, timeout: float) -> int:
|
||||
async with self._condition:
|
||||
if self._generation != generation:
|
||||
return self._generation
|
||||
with contextlib.suppress(asyncio.TimeoutError):
|
||||
await asyncio.wait_for(
|
||||
self._condition.wait_for(lambda: self._generation != generation),
|
||||
timeout=timeout,
|
||||
)
|
||||
return self._generation
|
||||
|
||||
@property
|
||||
def generation(self) -> int:
|
||||
return self._generation
|
||||
|
||||
async def oldest_cursor(self) -> int | None:
|
||||
return await self._read(lambda repository: repository.oldest_cursor())
|
||||
|
||||
async def current_cursor(self) -> int:
|
||||
return await self._read(lambda repository: repository.current_cursor())
|
||||
|
||||
async def read_after(self, after: int) -> list[dict]:
|
||||
return await self._read(lambda repository: repository.read_after(after))
|
||||
|
||||
async def _read(self, operation):
|
||||
await self.start()
|
||||
async with self._connection_lock:
|
||||
if self._connection is None or self._connection.closed:
|
||||
self._connection = await self._connect()
|
||||
self._ready.set()
|
||||
repository = WorkbenchEventRepository(self._connection, self.project_id)
|
||||
return await operation(repository)
|
||||
|
||||
async def _drop_connection(self) -> None:
|
||||
async with self._connection_lock:
|
||||
connection, self._connection = self._connection, None
|
||||
if connection is not None and self._pool is not None:
|
||||
with contextlib.suppress(Exception):
|
||||
await self._pool.putconn(connection)
|
||||
|
||||
async def close(self) -> None:
|
||||
self._closed = True
|
||||
if self._idle_task is not None:
|
||||
self._idle_task.cancel()
|
||||
self._idle_task = None
|
||||
if self._listener_task is not None:
|
||||
self._listener_task.cancel()
|
||||
with contextlib.suppress(asyncio.CancelledError):
|
||||
await self._listener_task
|
||||
self._listener_task = None
|
||||
await self._drop_connection()
|
||||
if self._pool is not None:
|
||||
await self._pool.close()
|
||||
self._pool = None
|
||||
|
||||
|
||||
class WorkbenchEventHubManager:
|
||||
def __init__(self, idle_seconds: float = 60) -> None:
|
||||
self.idle_seconds = idle_seconds
|
||||
self._hubs: dict[tuple[UUID, str], WorkbenchEventHub] = {}
|
||||
self._lock = asyncio.Lock()
|
||||
|
||||
@asynccontextmanager
|
||||
async def subscribe(self, project_id: UUID, dsn: str):
|
||||
key = (project_id, dsn)
|
||||
async with self._lock:
|
||||
hub = self._hubs.get(key)
|
||||
if hub is None:
|
||||
hub = WorkbenchEventHub(project_id, dsn)
|
||||
self._hubs[key] = hub
|
||||
if hub._idle_task is not None:
|
||||
hub._idle_task.cancel()
|
||||
hub._idle_task = None
|
||||
hub.subscribers += 1
|
||||
try:
|
||||
await hub.start()
|
||||
yield hub
|
||||
finally:
|
||||
async with self._lock:
|
||||
hub.subscribers = max(0, hub.subscribers - 1)
|
||||
if hub.subscribers == 0 and hub._idle_task is None:
|
||||
hub._idle_task = asyncio.create_task(self._close_when_idle(key, hub))
|
||||
|
||||
async def _close_when_idle(
|
||||
self, key: tuple[UUID, str], hub: WorkbenchEventHub
|
||||
) -> None:
|
||||
try:
|
||||
await asyncio.sleep(self.idle_seconds)
|
||||
async with self._lock:
|
||||
if hub.subscribers or self._hubs.get(key) is not hub:
|
||||
return
|
||||
self._hubs.pop(key, None)
|
||||
hub._idle_task = None
|
||||
await hub.close()
|
||||
except asyncio.CancelledError:
|
||||
return
|
||||
|
||||
async def close_all(self) -> None:
|
||||
async with self._lock:
|
||||
hubs = list(self._hubs.values())
|
||||
self._hubs.clear()
|
||||
await asyncio.gather(*(hub.close() for hub in hubs), return_exceptions=True)
|
||||
|
||||
|
||||
workbench_event_hubs = WorkbenchEventHubManager()
|
||||
@@ -0,0 +1,37 @@
|
||||
"""Cursor reads for the durable project-scoped workbench event log."""
|
||||
from __future__ import annotations
|
||||
|
||||
from psycopg import AsyncConnection
|
||||
|
||||
|
||||
class WorkbenchEventRepository:
|
||||
def __init__(self, connection: AsyncConnection, project_id):
|
||||
self.connection = connection
|
||||
self.project_id = str(project_id)
|
||||
|
||||
async def current_cursor(self) -> int:
|
||||
cursor = await self.connection.execute(
|
||||
"SELECT COALESCE(max(event_seq),0) AS cursor FROM workbench.domain_event_log WHERE project_id=%s",
|
||||
(self.project_id,),
|
||||
)
|
||||
row = await cursor.fetchone()
|
||||
return int(row["cursor"])
|
||||
|
||||
async def oldest_cursor(self) -> int | None:
|
||||
cursor = await self.connection.execute(
|
||||
"SELECT min(event_seq) AS cursor FROM workbench.domain_event_log WHERE project_id=%s",
|
||||
(self.project_id,),
|
||||
)
|
||||
row = await cursor.fetchone()
|
||||
return int(row["cursor"]) if row and row["cursor"] is not None else None
|
||||
|
||||
async def read_after(self, after: int, limit: int = 200) -> list[dict]:
|
||||
cursor = await self.connection.execute(
|
||||
"""SELECT event_seq,event_id,project_id,topic,aggregate_type,aggregate_id,
|
||||
aggregate_revision,occurred_at,metadata
|
||||
FROM workbench.domain_event_log
|
||||
WHERE project_id=%s AND event_seq>%s
|
||||
ORDER BY event_seq LIMIT %s""",
|
||||
(self.project_id, after, limit),
|
||||
)
|
||||
return [dict(row) for row in await cursor.fetchall()]
|
||||
Reference in New Issue
Block a user