79 lines
3.5 KiB
Python
79 lines
3.5 KiB
Python
from typing import List, Any
|
|
from datetime import datetime
|
|
from psycopg import AsyncConnection, sql
|
|
|
|
class SchemeRepository:
|
|
|
|
# --- Link Simulation ---
|
|
|
|
@staticmethod
|
|
async def insert_links_batch(conn: AsyncConnection, data: List[dict]):
|
|
if not data:
|
|
return
|
|
|
|
async with conn.cursor() as cur:
|
|
async with cur.copy(
|
|
"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')
|
|
))
|
|
|
|
@staticmethod
|
|
async def get_links_by_scheme_and_time(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 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))
|
|
|
|
# --- Node Simulation ---
|
|
|
|
@staticmethod
|
|
async def insert_nodes_batch(conn: AsyncConnection, data: List[dict]):
|
|
if not data:
|
|
return
|
|
|
|
async with conn.cursor() as cur:
|
|
async with cur.copy(
|
|
"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')
|
|
))
|
|
|
|
@staticmethod
|
|
async def get_nodes_by_scheme_and_time(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 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))
|