删除旧文件;更新数据库查询方法
This commit is contained in:
@@ -2,10 +2,11 @@ from typing import List, Any, Optional
|
||||
from datetime import datetime
|
||||
from psycopg import AsyncConnection, sql
|
||||
|
||||
|
||||
class RealtimeRepository:
|
||||
|
||||
|
||||
# --- Link Simulation ---
|
||||
|
||||
|
||||
@staticmethod
|
||||
async def insert_links_batch(conn: AsyncConnection, data: List[dict]):
|
||||
"""Batch insert for realtime.link_simulation using COPY for performance."""
|
||||
@@ -17,51 +18,141 @@ class RealtimeRepository:
|
||||
"COPY realtime.link_simulation (time, id, flow, friction, headloss, quality, reaction, setting, status, velocity) FROM STDIN"
|
||||
) as copy:
|
||||
for item in data:
|
||||
await copy.write_row((
|
||||
item['time'], 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["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_time(conn: AsyncConnection, start_time: datetime, end_time: datetime) -> List[dict]:
|
||||
async def get_link_by_time_range(
|
||||
conn: AsyncConnection, start_time: datetime, end_time: datetime, link_id: str
|
||||
) -> List[dict]:
|
||||
async with conn.cursor() as cur:
|
||||
await cur.execute(
|
||||
"SELECT * FROM realtime.link_simulation WHERE time >= %s AND time <= %s",
|
||||
(start_time, end_time)
|
||||
"SELECT * FROM realtime.link_simulation WHERE time >= %s AND time <= %s AND id = %s",
|
||||
(start_time, end_time, link_id),
|
||||
)
|
||||
return await cur.fetchall()
|
||||
|
||||
@staticmethod
|
||||
async def get_link_field(conn: AsyncConnection, time: datetime, link_id: str, field: str) -> Any:
|
||||
async def get_links_by_time_range(
|
||||
conn: AsyncConnection, start_time: datetime, end_time: datetime
|
||||
) -> List[dict]:
|
||||
async with conn.cursor() as cur:
|
||||
await cur.execute(
|
||||
"SELECT * FROM realtime.link_simulation WHERE time >= %s AND time <= %s",
|
||||
(start_time, end_time),
|
||||
)
|
||||
return await cur.fetchall()
|
||||
|
||||
@staticmethod
|
||||
async def get_link_field_by_time_range(
|
||||
conn: AsyncConnection,
|
||||
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"}
|
||||
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 realtime.link_simulation WHERE time = %s AND id = %s").format(sql.Identifier(field))
|
||||
|
||||
query = sql.SQL(
|
||||
"SELECT {} FROM realtime.link_simulation WHERE time >= %s AND time <= %s AND id = %s"
|
||||
).format(sql.Identifier(field))
|
||||
|
||||
async with conn.cursor() as cur:
|
||||
await cur.execute(query, (time, link_id))
|
||||
await cur.execute(query, (start_time, end_time, link_id))
|
||||
row = await cur.fetchone()
|
||||
return row[field] if row else None
|
||||
|
||||
@staticmethod
|
||||
async def update_link_field(conn: AsyncConnection, time: datetime, link_id: str, field: str, value: Any):
|
||||
valid_fields = {"flow", "friction", "headloss", "quality", "reaction", "setting", "status", "velocity"}
|
||||
async def get_links_field_by_time_range(
|
||||
conn: AsyncConnection,
|
||||
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 realtime.link_simulation SET {} = %s WHERE time = %s AND id = %s").format(sql.Identifier(field))
|
||||
|
||||
query = sql.SQL(
|
||||
"SELECT {} FROM realtime.link_simulation WHERE time >= %s AND time <= %s"
|
||||
).format(sql.Identifier(field))
|
||||
|
||||
async with conn.cursor() as cur:
|
||||
await cur.execute(query, (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,
|
||||
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 realtime.link_simulation SET {} = %s WHERE time = %s AND id = %s"
|
||||
).format(sql.Identifier(field))
|
||||
|
||||
async with conn.cursor() as cur:
|
||||
await cur.execute(query, (value, time, link_id))
|
||||
|
||||
@staticmethod
|
||||
async def delete_links_by_time(conn: AsyncConnection, start_time: datetime, end_time: datetime):
|
||||
async def delete_links_by_time_range(
|
||||
conn: AsyncConnection, start_time: datetime, end_time: datetime
|
||||
):
|
||||
async with conn.cursor() as cur:
|
||||
await cur.execute(
|
||||
"DELETE FROM realtime.link_simulation WHERE time >= %s AND time <= %s",
|
||||
(start_time, end_time)
|
||||
(start_time, end_time),
|
||||
)
|
||||
|
||||
# --- Node Simulation ---
|
||||
@@ -76,39 +167,102 @@ class RealtimeRepository:
|
||||
"COPY realtime.node_simulation (time, id, actual_demand, total_head, pressure, quality) FROM STDIN"
|
||||
) as copy:
|
||||
for item in data:
|
||||
await copy.write_row((
|
||||
item['time'], item['id'], item.get('actual_demand'), item.get('total_head'), item.get('pressure'), item.get('quality')
|
||||
))
|
||||
await copy.write_row(
|
||||
(
|
||||
item["time"],
|
||||
item["id"],
|
||||
item.get("actual_demand"),
|
||||
item.get("total_head"),
|
||||
item.get("pressure"),
|
||||
item.get("quality"),
|
||||
)
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
async def get_nodes_by_time(conn: AsyncConnection, start_time: datetime, end_time: datetime) -> List[dict]:
|
||||
async def get_node_by_time_range(
|
||||
conn: AsyncConnection, start_time: datetime, end_time: datetime, node_id: str
|
||||
) -> List[dict]:
|
||||
async with conn.cursor() as cur:
|
||||
await cur.execute(
|
||||
"SELECT * FROM realtime.node_simulation WHERE time >= %s AND time <= %s AND id = %s",
|
||||
(start_time, end_time, node_id),
|
||||
)
|
||||
return await cur.fetchall()
|
||||
|
||||
@staticmethod
|
||||
async def get_nodes_by_time_range(
|
||||
conn: AsyncConnection, start_time: datetime, end_time: datetime
|
||||
) -> List[dict]:
|
||||
async with conn.cursor() as cur:
|
||||
await cur.execute(
|
||||
"SELECT * FROM realtime.node_simulation WHERE time >= %s AND time <= %s",
|
||||
(start_time, end_time)
|
||||
(start_time, end_time),
|
||||
)
|
||||
return await cur.fetchall()
|
||||
|
||||
|
||||
@staticmethod
|
||||
async def get_node_field(conn: AsyncConnection, time: datetime, node_id: str, field: str) -> Any:
|
||||
async def get_node_field_by_time_range(
|
||||
conn: AsyncConnection,
|
||||
start_time: datetime,
|
||||
end_time: datetime,
|
||||
node_id: str,
|
||||
field: str,
|
||||
) -> Any:
|
||||
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 realtime.node_simulation WHERE time = %s AND id = %s").format(sql.Identifier(field))
|
||||
|
||||
query = sql.SQL(
|
||||
"SELECT {} FROM realtime.node_simulation WHERE time >= %s AND time <= %s AND id = %s"
|
||||
).format(sql.Identifier(field))
|
||||
|
||||
async with conn.cursor() as cur:
|
||||
await cur.execute(query, (time, node_id))
|
||||
await cur.execute(query, (start_time, end_time, node_id))
|
||||
row = await cur.fetchone()
|
||||
return row[field] if row else None
|
||||
|
||||
@staticmethod
|
||||
async def update_node_field(conn: AsyncConnection, time: datetime, node_id: str, field: str, value: Any):
|
||||
async def get_nodes_field_by_time_range(
|
||||
conn: AsyncConnection, start_time: datetime, end_time: datetime, field: str
|
||||
) -> 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 realtime.node_simulation SET {} = %s WHERE time = %s AND id = %s").format(sql.Identifier(field))
|
||||
|
||||
query = sql.SQL(
|
||||
"SELECT {} FROM realtime.node_simulation WHERE time >= %s AND time <= %s"
|
||||
).format(sql.Identifier(field))
|
||||
|
||||
async with conn.cursor() as cur:
|
||||
await cur.execute(query, (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,
|
||||
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 realtime.node_simulation SET {} = %s WHERE time = %s AND id = %s"
|
||||
).format(sql.Identifier(field))
|
||||
|
||||
async with conn.cursor() as cur:
|
||||
await cur.execute(query, (value, time, node_id))
|
||||
|
||||
@staticmethod
|
||||
async def delete_nodes_by_time_range(
|
||||
conn: AsyncConnection, start_time: datetime, end_time: datetime
|
||||
):
|
||||
async with conn.cursor() as cur:
|
||||
await cur.execute(
|
||||
"DELETE FROM realtime.node_simulation WHERE time >= %s AND time <= %s",
|
||||
(start_time, end_time),
|
||||
)
|
||||
|
||||
Reference in New Issue
Block a user