from datetime import datetime, timezone from typing import Any from uuid import UUID, uuid4 from psycopg.rows import dict_row from psycopg.types.json import Jsonb from app.infra.db.postgresql.analysis import AnalysisRepository from app.native.wndb.core.connection import project_connection RUN_TYPE = "sensor_placement" RESULT_TYPE = "sensor_placement" def _placement_row(row: dict[str, Any]) -> dict[str, Any]: payload = row.get("payload") if isinstance(row.get("payload"), dict) else {} locations = [str(item) for item in payload.get("sensor_locations", [])] return { "run_id": row["run_id"], "name": row["name"], "sensor_count": len(locations), "min_diameter": int(payload.get("minimum_diameter", 0)), "created_by": row["created_by"], "created_at": row["created_at"], "status": row["status"], "sensor_locations": locations, } def get_all_sensor_placements(name: str) -> list[dict[str, Any]]: with project_connection(name) as conn, conn.cursor(row_factory=dict_row) as cur: cur.execute( """ SELECT r.run_id, r.name, r.created_by, r.created_at, r.status, result.payload FROM analysis.runs AS r JOIN LATERAL ( SELECT payload FROM analysis.results WHERE run_id = r.run_id AND result_type = %s ORDER BY created_at DESC, result_id DESC LIMIT 1 ) AS result ON true WHERE r.run_type = %s ORDER BY r.created_at DESC, r.run_id """, (RESULT_TYPE, RUN_TYPE), ) return [_placement_row(row) for row in cur.fetchall()] def create_sensor_placement( name: str, *, run_name: str, min_diameter: int, created_by: str, sensor_locations: list[str], ) -> dict[str, Any]: run_id = uuid4() payload = { "sensor_number": len(sensor_locations), "minimum_diameter": min_diameter, "sensor_locations": sensor_locations, } with project_connection(name) as conn, conn.transaction(): created = AnalysisRepository.create_run_sync( conn, run_id=run_id, name=run_name, run_type=RUN_TYPE, created_by=created_by, started_at=datetime.now(timezone.utc), status="completed", parameters={}, ) AnalysisRepository.insert_result_sync( conn, run_id, result_type=RESULT_TYPE, payload=payload, ) return _placement_row(dict(created) | {"payload": payload}) def get_sensor_placement(name: str, run_id: UUID) -> dict[str, Any] | None: with project_connection(name) as conn, conn.cursor(row_factory=dict_row) as cur: cur.execute( """ SELECT r.run_id, r.name, r.created_by, r.created_at, r.status, result.payload FROM analysis.runs AS r JOIN LATERAL ( SELECT payload FROM analysis.results WHERE run_id = r.run_id AND result_type = %s ORDER BY created_at DESC, result_id DESC LIMIT 1 ) AS result ON true WHERE r.run_id = %s AND r.run_type = %s """, (RESULT_TYPE, run_id, RUN_TYPE), ) row = cur.fetchone() return _placement_row(row) if row else None def get_sensor_placement_nodes( name: str, node_ids: list[str], ) -> list[dict[str, Any]]: if not node_ids: return [] with project_connection(name) as conn, conn.cursor(row_factory=dict_row) as cur: cur.execute( """ WITH incident_pipe_diameters AS ( SELECT node_id, MAX(diameter) AS max_pipe_diameter FROM ( SELECT l.start_node_id AS node_id, p.diameter FROM network.pipes AS p JOIN network.links AS l ON l.id = p.link_id WHERE l.start_node_id = ANY(%s) UNION ALL SELECT l.end_node_id AS node_id, p.diameter FROM network.pipes AS p JOIN network.links AS l ON l.id = p.link_id WHERE l.end_node_id = ANY(%s) ) AS incident_pipes GROUP BY node_id ) SELECT g.id AS node_id, ipd.max_pipe_diameter, g.elevation, g.x AS project_x, g.y AS project_y, ST_X(g.geom) AS map_x, ST_Y(g.geom) AS map_y FROM gis.junctions AS g LEFT JOIN incident_pipe_diameters AS ipd ON ipd.node_id = g.id WHERE g.id = ANY(%s) ORDER BY g.id """, (node_ids, node_ids, node_ids), ) return list(cur.fetchall()) def update_sensor_placement( name: str, run_id: UUID, *, expected_sensor_locations: list[str], sensor_locations: list[str], ) -> dict[str, Any] | None: with project_connection(name) as conn, conn.transaction(): with conn.cursor(row_factory=dict_row) as cur: cur.execute( """ SELECT result_id, payload FROM analysis.results WHERE run_id = %s AND result_type = %s ORDER BY created_at DESC, result_id DESC LIMIT 1 FOR UPDATE """, (run_id, RESULT_TYPE), ) result = cur.fetchone() if result is None: return None payload = result["payload"] if isinstance(result["payload"], dict) else {} current = [str(item) for item in payload.get("sensor_locations", [])] if current != expected_sensor_locations: return None payload = { **payload, "sensor_number": len(sensor_locations), "sensor_locations": sensor_locations, } cur.execute( "UPDATE analysis.results SET payload = %s WHERE result_id = %s", (Jsonb(payload), result["result_id"]), ) return get_sensor_placement(name, run_id)