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

from dataclasses import dataclass
from types import MappingProxyType
from typing import Any, Mapping

from psycopg import AsyncConnection

from app.native.wndb.core.database import read_all, try_read


_SCADA_VIEW_SELECT = """
    SELECT id AS device_id, device_type, node_id, link_id, api_query_id,
           measurement_unit, transmission_mode, transmission_frequency, reliability, x, y,
           ST_X(ST_Transform(geom, 4326)) AS longitude,
           ST_Y(ST_Transform(geom, 4326)) AS latitude
    FROM gis.scada_devices
"""


@dataclass(frozen=True)
class ScadaElementMappings:
    reservoirs: Mapping[str, str]
    tanks: Mapping[str, str]
    fixed_pumps: Mapping[str, str]
    variable_pumps: Mapping[str, str]
    pressure: Mapping[str, str]
    demand: Mapping[str, str]
    quality: Mapping[str, str]


def _empty_mapping_groups() -> dict[str, dict[str, str]]:
    return {
        "reservoir_liquid_level": {},
        "tank_liquid_level": {},
        "fixed_pump": {},
        "variable_pump": {},
        "pressure": {},
        "demand": {},
        "quality": {},
    }

def _optional_text(value: Any) -> str | None:
    return str(value).strip() if value is not None else None


def _optional_float(value: Any) -> float | None:
    return float(value) if value is not None else None


def _optional_int(value: Any) -> int | None:
    return int(value) if value is not None else None


def _device(record: dict[str, Any]) -> dict[str, Any]:
    return {
        "device_id": str(record["device_id"]).strip(),
        "device_type": str(record["device_type"]).strip().lower(),
        "node_id": _optional_text(record["node_id"]),
        "link_id": _optional_text(record["link_id"]),
        "api_query_id": _optional_text(record["api_query_id"]),
        "measurement_unit": str(record["measurement_unit"]).strip(),
        "transmission_mode": record["transmission_mode"],
        "transmission_frequency": record["transmission_frequency"],
        "reliability": _optional_int(record["reliability"]),
        "x": _optional_float(record["x"]),
        "y": _optional_float(record["y"]),
        "longitude": _optional_float(record["longitude"]),
        "latitude": _optional_float(record["latitude"]),
    }


class ScadaInfoRepository:
    """Read SCADA metadata from the current project's business database."""

    @staticmethod
    async def get_scadas(conn: AsyncConnection) -> list[dict[str, Any]]:
        async with conn.cursor() as cur:
            await cur.execute(
                _SCADA_VIEW_SELECT + " ORDER BY device_id"
            )
            records = await cur.fetchall()

        return [_device(record) for record in records]

    @staticmethod
    async def get_scada(
        conn: AsyncConnection, device_id: str
    ) -> dict[str, Any] | None:
        async with conn.cursor() as cur:
            await cur.execute(
                _SCADA_VIEW_SELECT + " WHERE id = %s",
                (device_id,),
            )
            record = await cur.fetchone()
        return _device(record) if record else None

    @staticmethod
    async def get_existing_device_ids(
        conn: AsyncConnection, device_ids: list[str]
    ) -> set[str]:
        if not device_ids:
            return set()
        async with conn.cursor() as cur:
            await cur.execute(
                "SELECT device_id FROM asset.scada_devices WHERE device_id = ANY(%s)",
                (device_ids,),
            )
            return {str(row["device_id"]).strip() for row in await cur.fetchall()}

    @staticmethod
    async def get_scadas_for_elements(
        conn: AsyncConnection,
        node_ids: list[str],
        link_ids: list[str],
    ) -> list[dict[str, Any]]:
        if not node_ids and not link_ids:
            return []
        async with conn.cursor() as cur:
            await cur.execute(
                _SCADA_VIEW_SELECT
                + " WHERE node_id = ANY(%s) OR link_id = ANY(%s) ORDER BY device_id",
                (node_ids, link_ids),
            )
            return [_device(record) for record in await cur.fetchall()]


def get_scada_info_schema(name: str) -> dict[str, dict[str, Any]]:
    return {
        "device_id": {"type": "str", "optional": False, "readonly": True},
        "device_type": {"type": "str", "optional": False, "readonly": True},
        "node_id": {"type": "str", "optional": True, "readonly": True},
        "link_id": {"type": "str", "optional": True, "readonly": True},
        "api_query_id": {"type": "str", "optional": True, "readonly": True},
        "measurement_unit": {"type": "str", "optional": False, "readonly": True},
        "transmission_mode": {"type": "str", "optional": False, "readonly": True},
        "transmission_frequency": {"type": "str", "optional": False, "readonly": True},
        "reliability": {"type": "int", "optional": False, "readonly": True},
        "x": {"type": "float", "optional": True, "readonly": True},
        "y": {"type": "float", "optional": True, "readonly": True},
        "longitude": {"type": "float", "optional": True, "readonly": True},
        "latitude": {"type": "float", "optional": True, "readonly": True},
    }


def get_scada_info(name: str, device_id: str) -> dict[str, Any]:
    row = try_read(
        name,
        _SCADA_VIEW_SELECT + " WHERE id = %s",
        (device_id,),
    )
    return _device(row) if row else {}


def get_all_scada_info(name: str) -> list[dict[str, Any]]:
    return [
        _device(row)
        for row in read_all(name, _SCADA_VIEW_SELECT + " ORDER BY device_id")
    ]


def load_realtime_element_mappings(name: str) -> ScadaElementMappings:
    """Load one project-local immutable SCADA-to-model mapping snapshot."""
    groups = _empty_mapping_groups()
    rows = read_all(
        name,
        """
        SELECT device_type, COALESCE(node_id, link_id) AS element_id,
               api_query_id
        FROM asset.scada_devices
        WHERE transmission_mode = 'realtime'
          AND api_query_id IS NOT NULL
        """,
    )
    for row in rows:
        group = groups.get(str(row["device_type"]).strip().lower())
        if group is not None:
            group[str(row["element_id"]).strip()] = str(row["api_query_id"]).strip()
    immutable = {
        name: MappingProxyType(values.copy()) for name, values in groups.items()
    }
    return ScadaElementMappings(
        reservoirs=immutable["reservoir_liquid_level"],
        tanks=immutable["tank_liquid_level"],
        fixed_pumps=immutable["fixed_pump"],
        variable_pumps=immutable["variable_pump"],
        pressure=immutable["pressure"],
        demand=immutable["demand"],
        quality=immutable["quality"],
    )