from datetime import date, datetime from typing import Any from uuid import UUID, uuid4 from psycopg.types.json import Jsonb from app.native.wndb.core.connection import project_connection from app.services.time_api import parse_utc_time def scheme_name_exists(name: str, scheme_name: str) -> bool: with project_connection(name) as conn, conn.cursor() as cur: cur.execute( "select exists(select 1 from analysis.runs where name = %s)", (scheme_name,), ) row = cur.fetchone() return bool(row and row[0]) def store_scheme_info( name: str, scheme_name: str, scheme_type: str, username: str, scheme_start_time: datetime | str, scheme_detail: dict, ) -> UUID: """Create one completed, immutable analysis run.""" return create_analysis_run( name=name, scheme_name=scheme_name, scheme_type=scheme_type, username=username, scheme_start_time=scheme_start_time, scheme_detail=scheme_detail, status="completed", ) def create_analysis_run( name: str, scheme_name: str, scheme_type: str, username: str, scheme_start_time: datetime | str, scheme_detail: dict, *, status: str = "running", ) -> UUID: """Create a distinct execution record; names are labels, not identities.""" started_at = parse_utc_time(scheme_start_time, field_name="scheme_start_time") run_id = uuid4() with project_connection(name) as conn, 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) """, ( run_id, scheme_name, scheme_type, username, started_at, status, Jsonb(scheme_detail), ), ) return run_id def update_analysis_run( name: str, run_id: UUID, *, status: str, username: str, scheme_detail: dict, ) -> None: """Update lifecycle state and metadata for one execution identity.""" with project_connection(name) as conn, conn.cursor() as cur: cur.execute( """ update analysis.runs set created_by = %s, status = %s, parameters = %s where run_id = %s """, (username, status, Jsonb(scheme_detail), run_id), ) if cur.rowcount != 1: raise LookupError(f"analysis run {run_id} does not exist") def _run_row(row: dict[str, Any]) -> dict[str, Any]: parameters = row.get("parameters") if isinstance(row.get("parameters"), dict) else {} return { "run_id": row["run_id"], "name": row["name"], "run_type": row["run_type"], "created_by": row["created_by"], "created_at": row["created_at"], "started_at": row["started_at"], "status": row["status"], "parameters": parameters, } def _list_runs( name: str, run_type: str | None = None, query_date: date | None = None, ) -> list[dict[str, Any]]: clauses: list[str] = [] params: list[Any] = [] if run_type: clauses.append("run_type = %s") params.append(run_type) if query_date is not None: clauses.append("created_at::date = %s") params.append(query_date) where = f"where {' and '.join(clauses)}" if clauses else "" with project_connection(name) as conn, conn.cursor() as cur: cur.execute( f"select run_id, name, run_type, created_by, created_at, started_at, status, parameters from analysis.runs {where} order by created_at desc", params, ) return [_run_row(row) for row in cur.fetchall()] def query_scheme_list( name: str, scheme_type: str | None = None, query_date: date | None = None, ) -> list[dict[str, Any]]: return _list_runs(name, scheme_type, query_date) def _get_run_by_name( name: str, run_name: str, run_type: str | None = None, ) -> dict[str, Any]: params: list[Any] = [run_name] type_clause = "" if run_type: type_clause = "and run_type = %s" params.append(run_type) with project_connection(name) as conn, conn.cursor() as cur: cur.execute( f""" select run_id, name, run_type, created_by, created_at, started_at, status, parameters from analysis.runs where name = %s {type_clause} order by created_at desc limit 1 """, params, ) row = cur.fetchone() return _run_row(row) if row else {} def get_analysis_run(name: str, run_id: UUID) -> dict[str, Any]: with project_connection(name) as conn, 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,), ) row = cur.fetchone() return _run_row(row) if row else {} def query_scheme_detail( name: str, scheme_name: str, scheme_type: str | None = None, ) -> dict[str, Any]: return _get_run_by_name(name, scheme_name, scheme_type) def store_leakage_identify_result( name: str, scheme_name: str, network: str, sensor_nodes: list[str], result_rows: list[dict], node_area_map: dict[str, str], areas: list[dict], drawing_payload: dict | None = None, run_status: str = "completed", error_message: str | None = None, ) -> None: run = _get_run_by_name(name, scheme_name, "dma_leak_identification") if not run: raise LookupError(f"analysis run {scheme_name!r} does not exist") payload = { "network": network, "run_status": run_status, "error_message": error_message, "sensor_nodes": sensor_nodes, "rows": result_rows, "node_area_map": node_area_map, "areas": areas, "drawing_payload": drawing_payload or {}, } with project_connection(name) as conn, conn.cursor() as cur: cur.execute( "insert into analysis.results (run_id, result_type, payload) values (%s, 'leakage_identification', %s)", (run["run_id"], Jsonb(payload)), ) def _list_typed_runs( name: str, network: str, run_type: str, query_date: date | None, ) -> list[dict[str, Any]]: rows = _list_runs(name, run_type, query_date) return [ row for row in rows if not network or row["parameters"].get("network") in (None, network) ] def _typed_run_detail(name: str, run_name: str, run_type: str) -> dict[str, Any]: run = _get_run_by_name(name, run_name, run_type) if not run: return {} with project_connection(name) as conn, conn.cursor() as cur: cur.execute( "select result_type, payload, created_at from analysis.results where run_id = %s order by created_at, result_id", (run["run_id"],), ) results = [dict(row) for row in cur.fetchall()] return run | {"results": results} def query_leakage_identify_schemes( name: str, network: str, scheme_type: str = "dma_leak_identification", query_date: date | None = None, ) -> list[dict[str, Any]]: return _list_typed_runs(name, network, scheme_type, query_date) def query_leakage_identify_scheme_detail(name: str, scheme_name: str) -> dict[str, Any]: return _typed_run_detail(name, scheme_name, "dma_leak_identification") def query_burst_location_schemes( name: str, network: str, scheme_type: str = "burst_location", query_date: date | None = None, ) -> list[dict[str, Any]]: return _list_typed_runs(name, network, scheme_type, query_date) def query_burst_location_scheme_detail(name: str, scheme_name: str) -> dict[str, Any]: return _typed_run_detail(name, scheme_name, "burst_location") def query_burst_detection_schemes( name: str, network: str, scheme_type: str = "burst_detection", query_date: date | None = None, ) -> list[dict[str, Any]]: return _list_typed_runs(name, network, scheme_type, query_date) def query_burst_detection_scheme_detail(name: str, scheme_name: str) -> dict[str, Any]: return _typed_run_detail(name, scheme_name, "burst_detection")