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
@@ -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)