refactor(db)!: finalize pooled WNDB v2 migration

2026-08-27 17:26:22 +08:00
parent fa188af0b1
commit b74799a39d
105 changed files with 4988 additions and 5565 deletions
@@ -175,6 +175,53 @@ class AnalysisResultsRepository:
            await cur.execute(query, (run_id, link_id, start_time, end_time))
            return await cur.fetchall()

+    @staticmethod
+    async def get_series_by_ids(
+        conn: AsyncConnection,
+        run_id: UUID,
+        element_type: str,
+        element_ids: list[str],
+        start_time: datetime,
+        end_time: datetime,
+        field: str,
+    ) -> dict[str, list[dict[str, Any]]]:
+        if element_type == "node":
+            table_name, id_column, valid_fields = (
+                "node_results", "node_id", AnalysisResultsRepository.NODE_FIELDS
+            )
+        elif element_type == "link":
+            table_name, id_column, valid_fields = (
+                "link_results", "link_id", AnalysisResultsRepository.LINK_FIELDS
+            )
+        else:
+            raise ValueError(f"invalid analysis element type: {element_type}")
+        if field not in valid_fields:
+            raise ValueError(f"invalid {element_type} result field: {field}")
+
+        result: dict[str, list[dict[str, Any]]] = {
+            element_id: [] for element_id in element_ids
+        }
+        if not element_ids:
+            return result
+        query = sql.SQL(
+            "SELECT {} AS element_id, time, {} AS value FROM analysis.{} "
+            "WHERE run_id = %s AND {} = ANY(%s) AND time BETWEEN %s AND %s "
+            "ORDER BY {}, time"
+        ).format(
+            sql.Identifier(id_column),
+            sql.Identifier(field),
+            sql.Identifier(table_name),
+            sql.Identifier(id_column),
+            sql.Identifier(id_column),
+        )
+        async with conn.cursor() as cur:
+            await cur.execute(query, (run_id, element_ids, start_time, end_time))
+            for row in await cur.fetchall():
+                result.setdefault(str(row["element_id"]), []).append(
+                    {"time": row["time"], "value": row["value"]}
+                )
+        return result
+
    @staticmethod
    async def get_values_at_time(
        conn: AsyncConnection,
@@ -7,6 +7,19 @@ from app.services.time_api import parse_utc_time

 class RealtimeRepository:

+    @staticmethod
+    def _batch_time(data: List[dict]) -> datetime:
+        """Return one normalized timestamp shared by every row in a snapshot."""
+        if not data:
+            raise ValueError("Realtime batch must not be empty")
+        times = {
+            parse_utc_time(item["time"], field_name="time")
+            for item in data
+        }
+        if len(times) != 1:
+            raise ValueError("Realtime batch must contain exactly one timestamp")
+        return times.pop()
+
    # --- Link Simulation ---

    @staticmethod
@@ -15,8 +28,7 @@ class RealtimeRepository:
        if not data:
            return

-        # 假设同一批次的数据时间是相同的
-        target_time = data[0]["time"]
+        target_time = RealtimeRepository._batch_time(data)

        # 使用事务确保原子性
        async with conn.transaction():
@@ -38,7 +50,7 @@ class RealtimeRepository:
                    for item in data:
                        await copy.write_row(
                            (
-                                item["time"],
+                                target_time,
                                item["id"],
                                item.get("flow"),
                                item.get("friction"),
@@ -57,8 +69,7 @@ class RealtimeRepository:
        if not data:
            return

-        # 假设同一批次的数据时间是相同的
-        target_time = data[0]["time"]
+        target_time = RealtimeRepository._batch_time(data)

        # 使用事务确保原子性
        with conn.transaction():
@@ -80,7 +91,7 @@ class RealtimeRepository:
                    for item in data:
                        copy.write_row(
                            (
-                                item["time"],
+                                target_time,
                                item["id"],
                                item.get("flow"),
                                item.get("friction"),
@@ -99,7 +110,8 @@ class RealtimeRepository:
    ) -> List[dict]:
        async with conn.cursor() as cur:
            await cur.execute(
-                "SELECT * FROM realtime.link_results WHERE time >= %s AND time <= %s AND link_id = %s",
+                "SELECT * FROM realtime.link_results WHERE time >= %s AND time <= %s "
+                "AND link_id = %s ORDER BY time",
                (start_time, end_time, link_id),
            )
            return await cur.fetchall()
@@ -112,7 +124,8 @@ class RealtimeRepository:
        normalized_end_time = parse_utc_time(end_time, field_name="end_time")
        async with conn.cursor() as cur:
            await cur.execute(
-                "SELECT * FROM realtime.link_results WHERE time >= %s AND time <= %s",
+                "SELECT * FROM realtime.link_results WHERE time >= %s AND time <= %s "
+                "ORDER BY time, link_id",
                (normalized_start_time, normalized_end_time),
            )
            return await cur.fetchall()
@@ -140,7 +153,8 @@ class RealtimeRepository:
            raise ValueError(f"Invalid field: {field}")

        query = sql.SQL(
-            "SELECT time, {} FROM realtime.link_results WHERE time >= %s AND time <= %s AND link_id = %s"
+            "SELECT time, {} FROM realtime.link_results WHERE time >= %s "
+            "AND time <= %s AND link_id = %s ORDER BY time"
        ).format(sql.Identifier(field))

        async with conn.cursor() as cur:
@@ -150,6 +164,36 @@ class RealtimeRepository:
                {"time": row["time"].isoformat(), "value": row[field]} for row in rows
            ]

+    @staticmethod
+    async def get_link_fields_by_ids_time_range(
+        conn: AsyncConnection,
+        start_time: datetime,
+        end_time: datetime,
+        link_ids: list[str],
+        field: str,
+    ) -> dict[str, list[dict[str, Any]]]:
+        valid_fields = {
+            "flow", "friction", "headloss", "quality", "reaction",
+            "setting", "status", "velocity",
+        }
+        if field not in valid_fields:
+            raise ValueError(f"Invalid field: {field}")
+        result = {link_id: [] for link_id in link_ids}
+        if not link_ids:
+            return result
+        query = sql.SQL(
+            "SELECT link_id, time, {} FROM realtime.link_results "
+            "WHERE time BETWEEN %s AND %s AND link_id = ANY(%s) "
+            "ORDER BY link_id, time"
+        ).format(sql.Identifier(field))
+        async with conn.cursor() as cur:
+            await cur.execute(query, (start_time, end_time, link_ids))
+            for row in await cur.fetchall():
+                result.setdefault(str(row["link_id"]), []).append(
+                    {"time": row["time"].isoformat(), "value": row[field]}
+                )
+        return result
+
    @staticmethod
    async def get_links_field_by_time_range(
        conn: AsyncConnection,
@@ -172,7 +216,8 @@ class RealtimeRepository:
            raise ValueError(f"Invalid field: {field}")

        query = sql.SQL(
-            "SELECT link_id, time, {} FROM realtime.link_results WHERE time >= %s AND time <= %s"
+            "SELECT link_id, time, {} FROM realtime.link_results "
+            "WHERE time >= %s AND time <= %s ORDER BY link_id, time"
        ).format(sql.Identifier(field))

        async with conn.cursor() as cur:
@@ -230,8 +275,7 @@ class RealtimeRepository:
        if not data:
            return

-        # 假设同一批次的数据时间是相同的
-        target_time = data[0]["time"]
+        target_time = RealtimeRepository._batch_time(data)

        # 使用事务确保原子性
        async with conn.transaction():
@@ -253,7 +297,7 @@ class RealtimeRepository:
                    for item in data:
                        await copy.write_row(
                            (
-                                item["time"],
+                                target_time,
                                item["id"],
                                item.get("actual_demand"),
                                item.get("total_head"),
@@ -267,8 +311,7 @@ class RealtimeRepository:
        if not data:
            return

-        # 假设同一批次的数据时间是相同的
-        target_time = data[0]["time"]
+        target_time = RealtimeRepository._batch_time(data)

        # 使用事务确保原子性
        with conn.transaction():
@@ -290,7 +333,7 @@ class RealtimeRepository:
                    for item in data:
                        copy.write_row(
                            (
-                                item["time"],
+                                target_time,
                                item["id"],
                                item.get("actual_demand"),
                                item.get("total_head"),
@@ -305,7 +348,8 @@ class RealtimeRepository:
    ) -> List[dict]:
        async with conn.cursor() as cur:
            await cur.execute(
-                "SELECT * FROM realtime.node_results WHERE time >= %s AND time <= %s AND node_id = %s",
+                "SELECT * FROM realtime.node_results WHERE time >= %s AND time <= %s "
+                "AND node_id = %s ORDER BY time",
                (start_time, end_time, node_id),
            )
            return await cur.fetchall()
@@ -318,7 +362,8 @@ class RealtimeRepository:
        normalized_end_time = parse_utc_time(end_time, field_name="end_time")
        async with conn.cursor() as cur:
            await cur.execute(
-                "SELECT * FROM realtime.node_results WHERE time >= %s AND time <= %s",
+                "SELECT * FROM realtime.node_results WHERE time >= %s AND time <= %s "
+                "ORDER BY time, node_id",
                (normalized_start_time, normalized_end_time),
            )
            return await cur.fetchall()
@@ -336,7 +381,8 @@ class RealtimeRepository:
            raise ValueError(f"Invalid field: {field}")

        query = sql.SQL(
-            "SELECT time, {} FROM realtime.node_results WHERE time >= %s AND time <= %s AND node_id = %s"
+            "SELECT time, {} FROM realtime.node_results WHERE time >= %s "
+            "AND time <= %s AND node_id = %s ORDER BY time"
        ).format(sql.Identifier(field))

        async with conn.cursor() as cur:
@@ -346,6 +392,33 @@ class RealtimeRepository:
                {"time": row["time"].isoformat(), "value": row[field]} for row in rows
            ]

+    @staticmethod
+    async def get_node_fields_by_ids_time_range(
+        conn: AsyncConnection,
+        start_time: datetime,
+        end_time: datetime,
+        node_ids: list[str],
+        field: str,
+    ) -> dict[str, list[dict[str, Any]]]:
+        valid_fields = {"actual_demand", "total_head", "pressure", "quality"}
+        if field not in valid_fields:
+            raise ValueError(f"Invalid field: {field}")
+        result = {node_id: [] for node_id in node_ids}
+        if not node_ids:
+            return result
+        query = sql.SQL(
+            "SELECT node_id, time, {} FROM realtime.node_results "
+            "WHERE time BETWEEN %s AND %s AND node_id = ANY(%s) "
+            "ORDER BY node_id, time"
+        ).format(sql.Identifier(field))
+        async with conn.cursor() as cur:
+            await cur.execute(query, (start_time, end_time, node_ids))
+            for row in await cur.fetchall():
+                result.setdefault(str(row["node_id"]), []).append(
+                    {"time": row["time"].isoformat(), "value": row[field]}
+                )
+        return result
+
    @staticmethod
    async def get_nodes_field_by_time_range(
        conn: AsyncConnection, start_time: datetime, end_time: datetime, field: str
@@ -355,7 +428,8 @@ class RealtimeRepository:
            raise ValueError(f"Invalid field: {field}")

        query = sql.SQL(
-            "SELECT node_id, time, {} FROM realtime.node_results WHERE time >= %s AND time <= %s"
+            "SELECT node_id, time, {} FROM realtime.node_results "
+            "WHERE time >= %s AND time <= %s ORDER BY node_id, time"
        ).format(sql.Identifier(field))

        async with conn.cursor() as cur:
@@ -459,6 +533,19 @@ class RealtimeRepository:
        # transactions (savepoints), while this outer transaction guarantees
        # that a link write failure also rolls back the node replacement.
        async with conn.transaction():
+            async with conn.cursor() as cur:
+                await cur.execute(
+                    "SELECT pg_advisory_xact_lock(hashtextextended(%s::text, 1))",
+                    (simulation_time,),
+                )
+                await cur.execute(
+                    "DELETE FROM realtime.node_results WHERE time = %s",
+                    (simulation_time,),
+                )
+                await cur.execute(
+                    "DELETE FROM realtime.link_results WHERE time = %s",
+                    (simulation_time,),
+                )
            if node_data:
                await RealtimeRepository.insert_nodes_batch(conn, node_data)

@@ -525,6 +612,19 @@ class RealtimeRepository:
        # transactions (savepoints), while this outer transaction guarantees
        # that a link write failure also rolls back the node replacement.
        with conn.transaction():
+            with conn.cursor() as cur:
+                cur.execute(
+                    "SELECT pg_advisory_xact_lock(hashtextextended(%s::text, 1))",
+                    (simulation_time,),
+                )
+                cur.execute(
+                    "DELETE FROM realtime.node_results WHERE time = %s",
+                    (simulation_time,),
+                )
+                cur.execute(
+                    "DELETE FROM realtime.link_results WHERE time = %s",
+                    (simulation_time,),
+                )
            if node_data:
                RealtimeRepository.insert_nodes_batch_sync(conn, node_data)

@@ -35,7 +35,8 @@ class ScadaRepository:
    ) -> List[dict]:
        async with conn.cursor() as cur:
            await cur.execute(
-                "SELECT * FROM scada.measurements WHERE device_id = ANY(%s) AND time >= %s AND time <= %s",
+                "SELECT * 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()
@@ -49,7 +50,8 @@ class ScadaRepository:
    ) -> List[dict]:
        with conn.cursor(row_factory=dict_row) as cur:
            cur.execute(
-                "SELECT * FROM scada.measurements WHERE device_id = ANY(%s) AND time >= %s AND time <= %s",
+                "SELECT * 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()
@@ -88,7 +90,9 @@ class ScadaRepository:
            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)"
+            "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:
@@ -122,6 +126,31 @@ class ScadaRepository:
            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."""
+        valid_fields = {"monitored_value", "cleaned_value"}
+        if field not in valid_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