from datetime import datetime from typing import Any from uuid import UUID, uuid4 from psycopg import AsyncConnection, Connection from psycopg.types.json import Jsonb 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]: 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], ) -> None: with conn.cursor() as cur: cur.execute( """ UPDATE analysis.runs SET created_by = %s, status = %s, parameters = %s WHERE run_id = %s """, (created_by, status, Jsonb(parameters), run_id), ) if cur.rowcount != 1: raise LookupError(f"analysis run {run_id} does not exist") @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, ) -> 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, 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()