删除旧文件；更新数据库查询方法

timescaledb/schemas/scheme.py

@@ -1,13 +1,15 @@
-from typing import List, Any
+from typing import List, Any, Optional
 from datetime import datetime
 from psycopg import AsyncConnection, sql
 
+
 class SchemeRepository:
-    
+
     # --- Link Simulation ---
-    
+
     @staticmethod
     async def insert_links_batch(conn: AsyncConnection, data: List[dict]):
+        """Batch insert for scheme.link_simulation using COPY for performance."""
         if not data:
             return
 
@@ -16,31 +18,150 @@ class SchemeRepository:
                 "COPY scheme.link_simulation (time, scheme, id, flow, friction, headloss, quality, reaction, setting, status, velocity) FROM STDIN"
             ) as copy:
                 for item in data:
-                    await copy.write_row((
-                        item['time'], item['scheme'], item['id'], item.get('flow'), item.get('friction'), item.get('headloss'), 
-                        item.get('quality'), item.get('reaction'), item.get('setting'), item.get('status'), item.get('velocity')
-                    ))
+                    await copy.write_row(
+                        (
+                            item["time"],
+                            item["scheme"],
+                            item["id"],
+                            item.get("flow"),
+                            item.get("friction"),
+                            item.get("headloss"),
+                            item.get("quality"),
+                            item.get("reaction"),
+                            item.get("setting"),
+                            item.get("status"),
+                            item.get("velocity"),
+                        )
+                    )
 
     @staticmethod
-    async def get_links_by_scheme_and_time(conn: AsyncConnection, scheme: str, start_time: datetime, end_time: datetime) -> List[dict]:
+    async def get_link_by_scheme_and_time_range(
+        conn: AsyncConnection,
+        scheme: str,
+        start_time: datetime,
+        end_time: datetime,
+        link_id: str,
+    ) -> List[dict]:
         async with conn.cursor() as cur:
             await cur.execute(
-                "SELECT * FROM scheme.link_simulation WHERE scheme = %s AND time >= %s AND time <= %s",
-                (scheme, start_time, end_time)
+                "SELECT * FROM scheme.link_simulation WHERE scheme = %s AND time >= %s AND time <= %s AND id = %s",
+                (scheme, start_time, end_time, link_id),
             )
             return await cur.fetchall()
 
     @staticmethod
-    async def update_link_field(conn: AsyncConnection, time: datetime, scheme: str, link_id: str, field: str, value: Any):
-        valid_fields = {"flow", "friction", "headloss", "quality", "reaction", "setting", "status", "velocity"}
+    async def get_links_by_scheme_and_time_range(
+        conn: AsyncConnection, scheme: str, start_time: datetime, end_time: datetime
+    ) -> List[dict]:
+        async with conn.cursor() as cur:
+            await cur.execute(
+                "SELECT * FROM scheme.link_simulation WHERE scheme = %s AND time >= %s AND time <= %s",
+                (scheme, start_time, end_time),
+            )
+            return await cur.fetchall()
+
+    @staticmethod
+    async def get_link_field_by_scheme_and_time_range(
+        conn: AsyncConnection,
+        scheme: str,
+        start_time: datetime,
+        end_time: datetime,
+        link_id: str,
+        field: str,
+    ) -> Any:
+        # Validate field name to prevent SQL injection
+        valid_fields = {
+            "flow",
+            "friction",
+            "headloss",
+            "quality",
+            "reaction",
+            "setting",
+            "status",
+            "velocity",
+        }
         if field not in valid_fields:
             raise ValueError(f"Invalid field: {field}")
 
-        query = sql.SQL("UPDATE scheme.link_simulation SET {} = %s WHERE time = %s AND scheme = %s AND id = %s").format(sql.Identifier(field))
-        
+        query = sql.SQL(
+            "SELECT {} FROM scheme.link_simulation WHERE scheme = %s AND time >= %s AND time <= %s AND id = %s"
+        ).format(sql.Identifier(field))
+
+        async with conn.cursor() as cur:
+            await cur.execute(query, (scheme, start_time, end_time, link_id))
+            row = await cur.fetchone()
+            return row[field] if row else None
+
+    @staticmethod
+    async def get_links_field_by_scheme_and_time_range(
+        conn: AsyncConnection,
+        scheme: str,
+        start_time: datetime,
+        end_time: datetime,
+        field: str,
+    ) -> Any:
+        # Validate field name to prevent SQL injection
+        valid_fields = {
+            "flow",
+            "friction",
+            "headloss",
+            "quality",
+            "reaction",
+            "setting",
+            "status",
+            "velocity",
+        }
+        if field not in valid_fields:
+            raise ValueError(f"Invalid field: {field}")
+
+        query = sql.SQL(
+            "SELECT {} FROM scheme.link_simulation WHERE scheme = %s AND time >= %s AND time <= %s"
+        ).format(sql.Identifier(field))
+
+        async with conn.cursor() as cur:
+            await cur.execute(query, (scheme, start_time, end_time))
+            row = await cur.fetchone()
+            return row[field] if row else None
+
+    @staticmethod
+    async def update_link_field(
+        conn: AsyncConnection,
+        time: datetime,
+        scheme: str,
+        link_id: str,
+        field: str,
+        value: Any,
+    ):
+        valid_fields = {
+            "flow",
+            "friction",
+            "headloss",
+            "quality",
+            "reaction",
+            "setting",
+            "status",
+            "velocity",
+        }
+        if field not in valid_fields:
+            raise ValueError(f"Invalid field: {field}")
+
+        query = sql.SQL(
+            "UPDATE scheme.link_simulation SET {} = %s WHERE time = %s AND scheme = %s AND id = %s"
+        ).format(sql.Identifier(field))
+
         async with conn.cursor() as cur:
             await cur.execute(query, (value, time, scheme, link_id))
 
+    @staticmethod
+    async def delete_links_by_scheme_and_time_range(
+        conn: AsyncConnection, scheme: str, start_time: datetime, end_time: datetime
+    ):
+        async with conn.cursor() as cur:
+            await cur.execute(
+                "DELETE FROM scheme.link_simulation WHERE scheme = %s AND time >= %s AND time <= %s",
+                (scheme, start_time, end_time),
+            )
+
     # --- Node Simulation ---
 
     @staticmethod
@@ -53,26 +174,115 @@ class SchemeRepository:
                 "COPY scheme.node_simulation (time, scheme, id, actual_demand, total_head, pressure, quality) FROM STDIN"
             ) as copy:
                 for item in data:
-                    await copy.write_row((
-                        item['time'], item['scheme'], item['id'], item.get('actual_demand'), item.get('total_head'), item.get('pressure'), item.get('quality')
-                    ))
+                    await copy.write_row(
+                        (
+                            item["time"],
+                            item["scheme"],
+                            item["id"],
+                            item.get("actual_demand"),
+                            item.get("total_head"),
+                            item.get("pressure"),
+                            item.get("quality"),
+                        )
+                    )
 
     @staticmethod
-    async def get_nodes_by_scheme_and_time(conn: AsyncConnection, scheme: str, start_time: datetime, end_time: datetime) -> List[dict]:
+    async def get_node_by_scheme_and_time_range(
+        conn: AsyncConnection,
+        scheme: str,
+        start_time: datetime,
+        end_time: datetime,
+        node_id: str,
+    ) -> List[dict]:
         async with conn.cursor() as cur:
             await cur.execute(
-                "SELECT * FROM scheme.node_simulation WHERE scheme = %s AND time >= %s AND time <= %s",
-                (scheme, start_time, end_time)
+                "SELECT * FROM scheme.node_simulation WHERE scheme = %s AND time >= %s AND time <= %s AND id = %s",
+                (scheme, start_time, end_time, node_id),
             )
             return await cur.fetchall()
 
     @staticmethod
-    async def update_node_field(conn: AsyncConnection, time: datetime, scheme: str, node_id: str, field: str, value: Any):
+    async def get_nodes_by_scheme_and_time_range(
+        conn: AsyncConnection, scheme: str, start_time: datetime, end_time: datetime
+    ) -> List[dict]:
+        async with conn.cursor() as cur:
+            await cur.execute(
+                "SELECT * FROM scheme.node_simulation WHERE scheme = %s AND time >= %s AND time <= %s",
+                (scheme, start_time, end_time),
+            )
+            return await cur.fetchall()
+
+    @staticmethod
+    async def get_node_field_by_scheme_and_time_range(
+        conn: AsyncConnection,
+        scheme: str,
+        start_time: datetime,
+        end_time: datetime,
+        node_id: str,
+        field: str,
+    ) -> Any:
+        # Validate field name to prevent SQL injection
         valid_fields = {"actual_demand", "total_head", "pressure", "quality"}
         if field not in valid_fields:
             raise ValueError(f"Invalid field: {field}")
 
-        query = sql.SQL("UPDATE scheme.node_simulation SET {} = %s WHERE time = %s AND scheme = %s AND id = %s").format(sql.Identifier(field))
-        
+        query = sql.SQL(
+            "SELECT {} FROM scheme.node_simulation WHERE scheme = %s AND time >= %s AND time <= %s AND id = %s"
+        ).format(sql.Identifier(field))
+
+        async with conn.cursor() as cur:
+            await cur.execute(query, (scheme, start_time, end_time, node_id))
+            row = await cur.fetchone()
+            return row[field] if row else None
+
+    @staticmethod
+    async def get_nodes_field_by_scheme_and_time_range(
+        conn: AsyncConnection,
+        scheme: str,
+        start_time: datetime,
+        end_time: datetime,
+        field: str,
+    ) -> Any:
+        # Validate field name to prevent SQL injection
+        valid_fields = {"actual_demand", "total_head", "pressure", "quality"}
+        if field not in valid_fields:
+            raise ValueError(f"Invalid field: {field}")
+
+        query = sql.SQL(
+            "SELECT {} FROM scheme.node_simulation WHERE scheme = %s AND time >= %s AND time <= %s"
+        ).format(sql.Identifier(field))
+
+        async with conn.cursor() as cur:
+            await cur.execute(query, (scheme, start_time, end_time))
+            row = await cur.fetchone()
+            return row[field] if row else None
+
+    @staticmethod
+    async def update_node_field(
+        conn: AsyncConnection,
+        time: datetime,
+        scheme: str,
+        node_id: str,
+        field: str,
+        value: Any,
+    ):
+        valid_fields = {"actual_demand", "total_head", "pressure", "quality"}
+        if field not in valid_fields:
+            raise ValueError(f"Invalid field: {field}")
+
+        query = sql.SQL(
+            "UPDATE scheme.node_simulation SET {} = %s WHERE time = %s AND scheme = %s AND id = %s"
+        ).format(sql.Identifier(field))
+
         async with conn.cursor() as cur:
             await cur.execute(query, (value, time, scheme, node_id))
+
+    @staticmethod
+    async def delete_nodes_by_scheme_and_time_range(
+        conn: AsyncConnection, scheme: str, start_time: datetime, end_time: datetime
+    ):
+        async with conn.cursor() as cur:
+            await cur.execute(
+                "DELETE FROM scheme.node_simulation WHERE scheme = %s AND time >= %s AND time <= %s",
+                (scheme, start_time, end_time),
+            )