next-tjwater-server/app/infra/db/postgresql/analysis.py

from datetime import datetime
from typing import Any
from uuid import UUID, uuid4

from psycopg import AsyncConnection, Connection
from psycopg.types.json import Jsonb
from app.infra.db.postgresql.workbench_events import lock_event_transaction_sync


class AnalysisRepository:
    @staticmethod
    def create_run_sync(
        conn: Connection,
        *,
        name: str,
        run_type: str,
        created_by: str,
        started_at: datetime,
        status: str,
        parameters: dict[str, Any],
        run_id: UUID | None = None,
    ) -> dict[str, Any]:
        lock_event_transaction_sync(conn)
        execution_id = run_id or uuid4()
        with conn.cursor() as cur:
            cur.execute(
                """
                INSERT INTO analysis.runs
                    (run_id, name, run_type, created_by, created_at,
                     started_at, status, parameters)
                VALUES (%s, %s, %s, %s, now(), %s, %s, %s)
                RETURNING run_id, name, run_type, created_by, created_at,
                          started_at, status, parameters
                """,
                (
                    execution_id,
                    name,
                    run_type,
                    created_by,
                    started_at,
                    status,
                    Jsonb(parameters),
                ),
            )
            created = cur.fetchone()
        if created is None:
            raise RuntimeError("analysis run insert returned no row")
        return created

    @staticmethod
    def update_run_sync(
        conn: Connection,
        run_id: UUID,
        *,
        status: str,
        created_by: str,
        parameters: dict[str, Any],
    ) -> int:
        lock_event_transaction_sync(conn)
        with conn.cursor() as cur:
            cur.execute(
                """
                UPDATE analysis.runs
                SET created_by = %s, status = %s, parameters = %s,
                    revision = revision + 1
                WHERE run_id = %s
                RETURNING revision
                """,
                (created_by, status, Jsonb(parameters), run_id),
            )
            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(
        conn: Connection,
        run_id: UUID,
        *,
        result_type: str,
        payload: dict[str, Any],
        node_id: str | None = None,
        link_id: str | None = None,
    ) -> None:
        with conn.cursor() as cur:
            cur.execute(
                """
                INSERT INTO analysis.results
                    (run_id, result_type, node_id, link_id, payload)
                VALUES (%s, %s, %s, %s, %s)
                """,
                (run_id, result_type, node_id, link_id, Jsonb(payload)),
            )

    @staticmethod
    def get_run_sync(conn: Connection, run_id: UUID) -> dict[str, Any] | None:
        with conn.cursor() as cur:
            cur.execute(
                """
                SELECT run_id, name, run_type, created_by, created_at,
                       started_at, status, parameters
                FROM analysis.runs
                WHERE run_id = %s
                """,
                (run_id,),
            )
            return cur.fetchone()

    @staticmethod
    async def list_runs(
        conn: AsyncConnection, *, limit: int = 100, offset: int = 0
    ) -> list[dict]:
        if not 1 <= limit <= 100 or offset < 0:
            raise ValueError("analysis run page is out of range")
        async with conn.cursor() as cur:
            await cur.execute(
                """
                SELECT run_id, name, run_type, created_by, created_at,
                       started_at, status, parameters
                FROM analysis.runs
                ORDER BY created_at DESC, run_id
                LIMIT %s OFFSET %s
                """,
                (limit, offset),
            )
            return await cur.fetchall()

    @staticmethod
    async def list_recent_runs(
        conn: AsyncConnection, *, limit: int = 100
    ) -> list[dict]:
        if not 1 <= limit <= 100:
            raise ValueError("recent analysis run limit must be between 1 and 100")
        async with conn.cursor() as cur:
            await cur.execute(
                """
                SELECT run_id, name, run_type, created_by, created_at,
                       started_at, status, parameters
                FROM analysis.runs
                ORDER BY created_at DESC, run_id
                LIMIT %s
                """,
                (limit,),
            )
            return await cur.fetchall()

    @staticmethod
    async def list_run_summaries(
        conn: AsyncConnection, *, limit: int = 21, offset: int = 0
    ) -> list[dict]:
        if not 1 <= limit <= 101 or offset < 0:
            raise ValueError("analysis run summary page is out of range")
        async with conn.cursor() as cur:
            await cur.execute(
                """
                SELECT r.run_id, r.name, r.run_type, r.created_by,
                       r.created_at, r.started_at, r.status,
                       EXISTS (
                           SELECT 1
                           FROM analysis.results AS result
                           WHERE result.run_id = r.run_id
                       ) AS has_results
                FROM analysis.runs AS r
                ORDER BY r.created_at DESC, r.run_id
                LIMIT %s OFFSET %s
                """,
                (limit, offset),
            )
            return await cur.fetchall()

    @staticmethod
    async def get_run(conn: AsyncConnection, run_id: UUID) -> dict | None:
        async with conn.cursor() as cur:
            await cur.execute(
                """
                SELECT run_id, name, run_type, created_by, created_at,
                       started_at, status, parameters
                FROM analysis.runs
                WHERE run_id = %s
                """,
                (run_id,),
            )
            return await cur.fetchone()

    @staticmethod
    async def list_results(
        conn: AsyncConnection,
        run_id: UUID,
        result_type: str | None = None,
        *,
        limit: int = 100,
        offset: int = 0,
        lookahead: bool = False,
    ) -> list[dict]:
        if not 1 <= limit <= 500 or offset < 0:
            raise ValueError("analysis result page is out of range")
        query = """
            SELECT result_id, run_id, result_type, node_id, link_id,
                   payload, created_at
            FROM analysis.results
            WHERE run_id = %s
        """
        params: tuple = (run_id,)
        if result_type is not None:
            query += " AND result_type = %s"
            params = (run_id, result_type)
        query += " ORDER BY created_at, result_id LIMIT %s OFFSET %s"
        params += (limit + int(lookahead), offset)

        async with conn.cursor() as cur:
            await cur.execute(query, params)
            return await cur.fetchall()

    @staticmethod
    async def list_recent_results(
        conn: AsyncConnection, run_id: UUID, *, limit: int = 100
    ) -> list[dict]:
        if not 1 <= limit <= 100:
            raise ValueError("recent analysis result limit must be between 1 and 100")
        async with conn.cursor() as cur:
            await cur.execute(
                """
                SELECT result_id, run_id, result_type, node_id, link_id,
                       payload, created_at
                FROM (
                    SELECT result_id, run_id, result_type, node_id, link_id,
                           payload, created_at
                    FROM analysis.results
                    WHERE run_id = %s
                    ORDER BY created_at DESC, result_id DESC
                    LIMIT %s
                ) AS recent_results
                ORDER BY created_at, result_id
                """,
                (run_id, limit),
            )
            return await cur.fetchall()

    @staticmethod
    async def list_junction_map_points(
        conn: AsyncConnection, node_ids: list[str]
    ) -> list[dict]:
        if not node_ids:
            return []
        if len(node_ids) > 500:
            raise ValueError("analysis map projection supports at most 500 junctions")
        async with conn.cursor() as cur:
            await cur.execute(
                """
                SELECT g.id AS node_id,
                       ST_X(ST_Transform(g.geom, 4326)) AS longitude,
                       ST_Y(ST_Transform(g.geom, 4326)) AS latitude
                FROM gis.junctions AS g
                WHERE g.id = ANY(%s)
                ORDER BY g.id
                """,
                (node_ids,),
            )
            return await cur.fetchall()