TJWaterServerBinary/app/infra/db/timescaledb/repositories/scada.py

from typing import List, Any
from datetime import datetime
from collections import defaultdict
from psycopg import AsyncConnection, Connection, sql
from psycopg.rows import dict_row


class ScadaRepository:
    VALUE_FIELDS = frozenset({"monitored_value", "cleaned_value"})
    RESULT_COLUMNS = "time, device_id, monitored_value, cleaned_value"

    @staticmethod
    async def insert_scada_batch(conn: AsyncConnection, data: List[dict]):
        if not data:
            return

        async with conn.cursor() as cur:
            async with cur.copy(
                "COPY scada.measurements (time, device_id, monitored_value, cleaned_value) FROM STDIN"
            ) as copy:
                for item in data:
                    await copy.write_row(
                        (
                            item["time"],
                            item["device_id"],
                            item.get("monitored_value"),
                            item.get("cleaned_value"),
                        )
                    )

    @staticmethod
    async def get_scada_by_ids_time_range(
        conn: AsyncConnection,
        device_ids: List[str],
        start_time: datetime,
        end_time: datetime,
    ) -> List[dict]:
        async with conn.cursor() as cur:
            await cur.execute(
                f"SELECT {ScadaRepository.RESULT_COLUMNS} FROM scada.measurements "
                "WHERE device_id = ANY(%s) "
                "AND time >= %s AND time <= %s ORDER BY device_id, time",
                (device_ids, start_time, end_time),
            )
            return await cur.fetchall()

    @staticmethod
    def get_scada_by_ids_time_range_sync(
        conn: Connection,
        device_ids: List[str],
        start_time: datetime,
        end_time: datetime,
    ) -> List[dict]:
        with conn.cursor(row_factory=dict_row) as cur:
            cur.execute(
                f"SELECT {ScadaRepository.RESULT_COLUMNS} FROM scada.measurements "
                "WHERE device_id = ANY(%s) "
                "AND time >= %s AND time <= %s ORDER BY device_id, time",
                (device_ids, start_time, end_time),
            )
            return cur.fetchall()

    @staticmethod
    def get_latest_scada_time_sync(
        conn: Connection,
        device_ids: List[str],
        before_time: datetime | None = None,
    ) -> datetime | None:
        with conn.cursor(row_factory=dict_row) as cur:
            if before_time is None:
                cur.execute(
                    "SELECT max(time) AS time FROM scada.measurements WHERE device_id = ANY(%s)",
                    (device_ids,),
                )
            else:
                cur.execute(
                    "SELECT max(time) AS time FROM scada.measurements "
                    "WHERE device_id = ANY(%s) AND time <= %s",
                    (device_ids, before_time),
                )
            row = cur.fetchone()
            return row["time"] if row else None

    @staticmethod
    async def get_scada_field_by_id_time_range(
        conn: AsyncConnection,
        device_ids: List[str],
        start_time: datetime,
        end_time: datetime,
        field: str,
    ) -> dict:
        if field not in ScadaRepository.VALUE_FIELDS:
            raise ValueError(f"Invalid field: {field}")

        query = sql.SQL(
            "SELECT device_id, time, {} FROM scada.measurements "
            "WHERE time >= %s AND time <= %s AND device_id = ANY(%s) "
            "ORDER BY device_id, time"
        ).format(sql.Identifier(field))

        async with conn.cursor() as cur:
            await cur.execute(query, (start_time, end_time, device_ids))
            rows = await cur.fetchall()
            result = defaultdict(list)
            for row in rows:
                result[row["device_id"]].append({
                    "time": row["time"].isoformat(),
                    "value": row[field]
                })
            return dict(result)

    @staticmethod
    async def update_scada_field(
        conn: AsyncConnection, time: datetime, device_id: str, field: str, value: Any
    ):
        if field not in ScadaRepository.VALUE_FIELDS:
            raise ValueError(f"Invalid field: {field}")

        update_query = sql.SQL(
            "UPDATE scada.measurements SET {} = %s WHERE time = %s AND device_id = %s"
        ).format(sql.Identifier(field))
        insert_query = sql.SQL(
            "INSERT INTO scada.measurements (time, device_id, {}) VALUES (%s, %s, %s)"
        ).format(sql.Identifier(field))

        async with conn.cursor() as cur:
            await cur.execute(update_query, (value, time, device_id))
            if cur.rowcount == 0:
                await cur.execute(insert_query, (time, device_id, value))

    @staticmethod
    async def update_scada_field_batch(
        conn: AsyncConnection,
        rows: list[tuple[datetime, str, float | None]],
        field: str,
    ) -> int:
        """Update existing SCADA samples in one set-based statement."""
        if field not in ScadaRepository.VALUE_FIELDS:
            raise ValueError(f"Invalid field: {field}")
        if not rows:
            return 0

        query = sql.SQL(
            "UPDATE scada.measurements AS measurement SET {} = batch.value "
            "FROM unnest(%s::timestamptz[], %s::text[], %s::double precision[]) "
            "AS batch(time, device_id, value) "
            "WHERE measurement.time = batch.time "
            "AND measurement.device_id = batch.device_id"
        ).format(sql.Identifier(field))
        times, device_ids, values = zip(*rows)
        async with conn.cursor() as cur:
            await cur.execute(query, (list(times), list(device_ids), list(values)))
            return cur.rowcount

    @staticmethod
    async def delete_scada_by_id_time_range(
        conn: AsyncConnection, device_id: str, start_time: datetime, end_time: datetime
    ):
        async with conn.cursor() as cur:
            await cur.execute(
                "DELETE FROM scada.measurements WHERE device_id = %s AND time >= %s AND time <= %s",
                (device_id, start_time, end_time),
            )