108 lines
4.1 KiB
Python
108 lines
4.1 KiB
Python
"""PostgreSQL persistence for task-produced alert events and current state."""
|
|
from __future__ import annotations
|
|
|
|
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
|
|
|
|
|
|
class TaskAlertRepository:
|
|
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 alert storage is not provisioned",
|
|
status=503,
|
|
) from exc
|
|
|
|
async def lock_alert(self, alert_id: str) -> None:
|
|
await self.connection.execute(
|
|
"SELECT pg_advisory_xact_lock(hashtextextended(%s, 0))",
|
|
(f"tjwater-alert:{self.project_id}:{alert_id}",),
|
|
)
|
|
|
|
async def get_event(self, producer_id: str, event_id: str) -> dict | None:
|
|
cursor = await self.connection.execute(
|
|
"""SELECT document FROM workbench.alert_events
|
|
WHERE project_id=%s AND producer_id=%s AND event_id=%s""",
|
|
(self.project_id, producer_id, event_id),
|
|
)
|
|
row = await cursor.fetchone()
|
|
return row["document"] if row else None
|
|
|
|
async def get_alert(self, alert_id: str) -> tuple[int, dict] | None:
|
|
cursor = await self.connection.execute(
|
|
"""SELECT revision,document FROM workbench.alerts
|
|
WHERE project_id=%s AND alert_id=%s""",
|
|
(self.project_id, alert_id),
|
|
)
|
|
row = await cursor.fetchone()
|
|
return (row["revision"], row["document"]) if row else None
|
|
|
|
async def save_event(self, event: dict, resulting_state: dict) -> None:
|
|
document = event | {"resulting_state": resulting_state}
|
|
await self.connection.execute(
|
|
"""INSERT INTO workbench.alert_events
|
|
(project_id,producer_id,event_id,alert_id,event_type,occurred_at,document)
|
|
VALUES (%s,%s,%s,%s,%s,%s,%s)""",
|
|
(
|
|
self.project_id,
|
|
event["producer_id"],
|
|
event["event_id"],
|
|
event["alert_id"],
|
|
event["event_type"],
|
|
event["occurred_at"],
|
|
Jsonb(document),
|
|
),
|
|
)
|
|
|
|
async def save_alert(self, state: dict) -> None:
|
|
await self.connection.execute(
|
|
"""INSERT INTO workbench.alerts
|
|
(project_id,alert_id,revision,status,occurred_at,updated_at,document)
|
|
VALUES (%s,%s,%s,%s,%s,%s,%s)
|
|
ON CONFLICT(project_id,alert_id) DO UPDATE SET
|
|
revision=excluded.revision,status=excluded.status,
|
|
updated_at=excluded.updated_at,document=excluded.document""",
|
|
(
|
|
self.project_id,
|
|
state["alert_id"],
|
|
state["revision"],
|
|
state["status"],
|
|
state["occurred_at"],
|
|
state["updated_at"],
|
|
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"]}',),
|
|
)
|