TJWaterServerBinary/app/infra/db/postgresql/sensor_placement.py

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)