补全 realtime scheme 中的复合 存储、查询方法

timescaledb/schemas/scheme.py

@@ -1,5 +1,5 @@
-from typing import List, Any, Optional
-from datetime import datetime
+from typing import List, Any, Dict
+from datetime import datetime, timedelta
 from psycopg import AsyncConnection, sql
 
 
@@ -286,3 +286,158 @@ class SchemeRepository:
                 "DELETE FROM scheme.node_simulation WHERE scheme = %s AND time >= %s AND time <= %s",
                 (scheme, start_time, end_time),
             )
+
+    # --- 复合查询 ---
+
+    @staticmethod
+    async def store_scheme_simulation_result(
+        conn: AsyncConnection,
+        scheme: str,
+        node_result_list: List[Dict[str, any]],
+        link_result_list: List[Dict[str, any]],
+        result_start_time: str,
+    ):
+        """
+        Store scheme simulation results to TimescaleDB.
+
+        Args:
+            conn: Database connection
+            scheme: Scheme name
+            node_result_list: List of node simulation results
+            link_result_list: List of link simulation results
+            result_start_time: Start time for the results (ISO format string)
+        """
+        # Convert result_start_time string to datetime if needed
+        if isinstance(result_start_time, str):
+            simulation_time = datetime.fromisoformat(
+                result_start_time.replace("Z", "+00:00")
+            )
+        else:
+            simulation_time = result_start_time
+
+        # Prepare node data for batch insert
+        node_data = []
+        for node_result in node_result_list:
+            node_data.append(
+                {
+                    "time": simulation_time,
+                    "scheme": scheme,
+                    "id": node_result.get("id"),
+                    "actual_demand": node_result.get("actual_demand"),
+                    "total_head": node_result.get("total_head"),
+                    "pressure": node_result.get("pressure"),
+                    "quality": node_result.get("quality"),
+                }
+            )
+
+        # Prepare link data for batch insert
+        link_data = []
+        for link_result in link_result_list:
+            link_data.append(
+                {
+                    "time": simulation_time,
+                    "scheme": scheme,
+                    "id": link_result.get("id"),
+                    "flow": link_result.get("flow"),
+                    "friction": link_result.get("friction"),
+                    "headloss": link_result.get("headloss"),
+                    "quality": link_result.get("quality"),
+                    "reaction": link_result.get("reaction"),
+                    "setting": link_result.get("setting"),
+                    "status": link_result.get("status"),
+                    "velocity": link_result.get("velocity"),
+                }
+            )
+
+        # Insert data using batch methods
+        if node_data:
+            await SchemeRepository.insert_nodes_batch(conn, node_data)
+
+        if link_data:
+            await SchemeRepository.insert_links_batch(conn, link_data)
+
+    @staticmethod
+    async def query_all_record_by_scheme_time_property(
+        conn: AsyncConnection,
+        scheme: str,
+        query_time: str,
+        type: str,
+        property: str,
+    ) -> list:
+        """
+        Query all records by scheme, time and property from TimescaleDB.
+
+        Args:
+            conn: Database connection
+            scheme: Scheme name
+            query_time: Time to query (ISO format string)
+            type: Type of data ("node" or "link")
+            property: Property/field to query
+
+        Returns:
+            List of records matching the criteria
+        """
+        # Convert query_time string to datetime
+        if isinstance(query_time, str):
+            target_time = datetime.fromisoformat(query_time.replace("Z", "+00:00"))
+        else:
+            target_time = query_time
+
+        # Create time range: query_time ± 1 second
+        start_time = target_time - timedelta(seconds=1)
+        end_time = target_time + timedelta(seconds=1)
+
+        # Query based on type
+        if type.lower() == "node":
+            return await SchemeRepository.get_nodes_field_by_scheme_and_time_range(
+                conn, scheme, start_time, end_time, property
+            )
+        elif type.lower() == "link":
+            return await SchemeRepository.get_links_field_by_scheme_and_time_range(
+                conn, scheme, start_time, end_time, property
+            )
+        else:
+            raise ValueError(f"Invalid type: {type}. Must be 'node' or 'link'")
+
+    @staticmethod
+    async def query_scheme_simulation_result_by_ID_time(
+        conn: AsyncConnection,
+        scheme: str,
+        ID: str,
+        type: str,
+        query_time: str,
+    ) -> list[dict]:
+        """
+        Query scheme simulation results by ID and time from TimescaleDB.
+
+        Args:
+            conn: Database connection
+            scheme: Scheme name
+            ID: The ID of the node or link
+            type: Type of data ("node" or "link")
+            query_time: Time to query (ISO format string)
+
+        Returns:
+            List of records matching the criteria
+        """
+        # Convert query_time string to datetime
+        if isinstance(query_time, str):
+            target_time = datetime.fromisoformat(query_time.replace("Z", "+00:00"))
+        else:
+            target_time = query_time
+
+        # Create time range: query_time ± 1 second
+        start_time = target_time - timedelta(seconds=1)
+        end_time = target_time + timedelta(seconds=1)
+
+        # Query based on type
+        if type.lower() == "node":
+            return await SchemeRepository.get_node_by_scheme_and_time_range(
+                conn, scheme, start_time, end_time, ID
+            )
+        elif type.lower() == "link":
+            return await SchemeRepository.get_link_by_scheme_and_time_range(
+                conn, scheme, start_time, end_time, ID
+            )
+        else:
+            raise ValueError(f"Invalid type: {type}. Must be 'node' or 'link'")