OrgTJWater/TJWaterServerBinary
2
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

refactor(db)!: adopt project-routed pooled databases

Browse Source 
Reorganize WNDB by responsibility and remove legacy scheme endpoints.\n\nRoute analysis and time-series access through project pools, preserve transactional realtime replacement, and refresh GIS materialized views after writes.\n\nAdd database architecture documentation, live pooling coverage, API contract updates, and executable container verification.\n\nBREAKING CHANGE: legacy scheme APIs and flat app.native.wndb module imports are removed.
This commit is contained in:
jiang
2026-08-25 18:35:05 +08:00
parent fdbcc5c033
commit fa188af0b1
181 changed files with 8446 additions and 33546 deletions
Download Patch File Download Diff File Expand all files Collapse all files
app/infra/db/timescaledb/repositories/analysis.py
+268
View File
@@ -0,0 +1,268 @@
from datetime import datetime, timedelta
from typing import Any
from uuid import UUID
from psycopg import AsyncConnection, Connection, sql
from app.services.time_api import parse_utc_time
class AnalysisResultsRepository:
NODE_FIELDS = {"actual_demand", "total_head", "pressure", "quality"}
LINK_FIELDS = {
"flow",
"friction",
"headloss",
"quality",
"reaction",
"setting",
"status",
"velocity",
}
@staticmethod
def prepare_simulation_rows(
node_results: list[dict[str, Any]],
link_results: list[dict[str, Any]],
result_start_time: str,
num_periods: int,
result_timestep_seconds: int,
) -> tuple[list[dict[str, Any]], list[dict[str, Any]]]:
start_time = parse_utc_time(
result_start_time, field_name="result_start_time"
)
timestep = timedelta(seconds=result_timestep_seconds)
node_rows: list[dict[str, Any]] = []
for node_result in node_results:
for period_index, values in enumerate(
node_result.get("result", [])[:num_periods]
):
node_rows.append(
{
"time": start_time + timestep * period_index,
"node_id": node_result["node"],
"actual_demand": values.get("demand"),
"total_head": values.get("head"),
"pressure": values.get("pressure"),
"quality": values.get("quality"),
}
)
link_rows: list[dict[str, Any]] = []
for link_result in link_results:
for period_index, values in enumerate(
link_result.get("result", [])[:num_periods]
):
link_rows.append(
{
"time": start_time + timestep * period_index,
"link_id": link_result["link"],
**{field: values.get(field) for field in AnalysisResultsRepository.LINK_FIELDS},
}
)
return node_rows, link_rows
@staticmethod
async def store_results(
conn: AsyncConnection,
run_id: UUID,
node_rows: list[dict[str, Any]],
link_rows: list[dict[str, Any]],
) -> None:
async with conn.transaction(), conn.cursor() as cur:
await AnalysisResultsRepository._lock_run(cur, run_id)
await AnalysisResultsRepository._assert_run_is_empty(cur, run_id)
if node_rows:
async with cur.copy(
"COPY analysis.node_results "
"(time, run_id, node_id, actual_demand, total_head, pressure, quality) "
"FROM STDIN"
) as copy:
for row in node_rows:
await copy.write_row(
(
row["time"],
run_id,
row["node_id"],
row.get("actual_demand"),
row.get("total_head"),
row.get("pressure"),
row.get("quality"),
)
)
if link_rows:
async with cur.copy(
"COPY analysis.link_results "
"(time, run_id, link_id, flow, friction, headloss, quality, "
"reaction, setting, status, velocity) FROM STDIN"
) as copy:
for row in link_rows:
await copy.write_row(
(
row["time"],
run_id,
row["link_id"],
row.get("flow"),
row.get("friction"),
row.get("headloss"),
row.get("quality"),
row.get("reaction"),
row.get("setting"),
row.get("status"),
row.get("velocity"),
)
)
@staticmethod
async def _assert_run_is_empty(cur, run_id: UUID) -> None:
await cur.execute(
"""
SELECT EXISTS (
SELECT 1 FROM analysis.node_results WHERE run_id = %s
UNION ALL
SELECT 1 FROM analysis.link_results WHERE run_id = %s
) AS exists
""",
(run_id, run_id),
)
row = await cur.fetchone()
if row and row["exists"]:
raise ValueError(f"analysis results already exist for run {run_id}")
@staticmethod
async def _lock_run(cur, run_id: UUID) -> None:
await cur.execute(
"SELECT pg_advisory_xact_lock(hashtextextended(%s::text, 0))",
(run_id,),
)
@staticmethod
async def get_node_series(
conn: AsyncConnection,
run_id: UUID,
node_id: str,
start_time: datetime,
end_time: datetime,
field: str,
) -> list[dict[str, Any]]:
if field not in AnalysisResultsRepository.NODE_FIELDS:
raise ValueError(f"invalid node result field: {field}")
query = sql.SQL(
"SELECT time, {} AS value FROM analysis.node_results "
"WHERE run_id = %s AND node_id = %s AND time BETWEEN %s AND %s "
"ORDER BY time"
).format(sql.Identifier(field))
async with conn.cursor() as cur:
await cur.execute(query, (run_id, node_id, start_time, end_time))
return await cur.fetchall()
@staticmethod
async def get_link_series(
conn: AsyncConnection,
run_id: UUID,
link_id: str,
start_time: datetime,
end_time: datetime,
field: str,
) -> list[dict[str, Any]]:
if field not in AnalysisResultsRepository.LINK_FIELDS:
raise ValueError(f"invalid link result field: {field}")
query = sql.SQL(
"SELECT time, {} AS value FROM analysis.link_results "
"WHERE run_id = %s AND link_id = %s AND time BETWEEN %s AND %s "
"ORDER BY time"
).format(sql.Identifier(field))
async with conn.cursor() as cur:
await cur.execute(query, (run_id, link_id, start_time, end_time))
return await cur.fetchall()
@staticmethod
async def get_values_at_time(
conn: AsyncConnection,
run_id: UUID,
element_type: str,
result_time: datetime,
field: str,
) -> dict[str, Any]:
if element_type == "node":
table, id_column, fields = (
"node_results",
"node_id",
AnalysisResultsRepository.NODE_FIELDS,
)
elif element_type == "link":
table, id_column, fields = (
"link_results",
"link_id",
AnalysisResultsRepository.LINK_FIELDS,
)
else:
raise ValueError("element_type must be node or link")
if field not in fields:
raise ValueError(f"invalid {element_type} result field: {field}")
query = sql.SQL(
"SELECT {id_column}, {field} AS value FROM analysis.{table} "
"WHERE run_id = %s AND time = %s ORDER BY {id_column}"
).format(
id_column=sql.Identifier(id_column),
field=sql.Identifier(field),
table=sql.Identifier(table),
)
async with conn.cursor() as cur:
await cur.execute(query, (run_id, result_time))
return {row[id_column]: row["value"] for row in await cur.fetchall()}
@staticmethod
def store_results_sync(
conn: Connection,
run_id: UUID,
node_rows: list[dict[str, Any]],
link_rows: list[dict[str, Any]],
) -> None:
with conn.transaction(), conn.cursor() as cur:
cur.execute(
"SELECT pg_advisory_xact_lock(hashtextextended(%s::text, 0))",
(run_id,),
)
cur.execute(
"""
SELECT EXISTS (
SELECT 1 FROM analysis.node_results WHERE run_id = %s
UNION ALL
SELECT 1 FROM analysis.link_results WHERE run_id = %s
) AS exists
""",
(run_id, run_id),
)
row = cur.fetchone()
if row and row["exists"]:
raise ValueError(f"analysis results already exist for run {run_id}")
if node_rows:
with cur.copy(
"COPY analysis.node_results "
"(time, run_id, node_id, actual_demand, total_head, pressure, quality) "
"FROM STDIN"
) as copy:
for item in node_rows:
copy.write_row(
(
item["time"], run_id, item["node_id"],
item.get("actual_demand"), item.get("total_head"),
item.get("pressure"), item.get("quality"),
)
)
if link_rows:
with cur.copy(
"COPY analysis.link_results "
"(time, run_id, link_id, flow, friction, headloss, quality, "
"reaction, setting, status, velocity) FROM STDIN"
) as copy:
for item in link_rows:
copy.write_row(
(
item["time"], run_id, item["link_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"),
)
)
app/infra/db/timescaledb/repositories/realtime.py
+56 -34
View File
@@ -11,7 +11,7 @@ class RealtimeRepository:
@staticmethod
async def insert_links_batch(conn: AsyncConnection, data: List[dict]):
"""Batch insert for realtime.link_simulation using DELETE then COPY for performance."""
"""Batch insert for realtime.link_results using DELETE then COPY."""
if not data:
return
@@ -21,15 +21,19 @@ class RealtimeRepository:
# 使用事务确保原子性
async with conn.transaction():
async with conn.cursor() as cur:
await cur.execute(
"SELECT pg_advisory_xact_lock(hashtextextended(%s::text, 1))",
(target_time,),
)
# 1. 先删除该时间点的旧数据
await cur.execute(
"DELETE FROM realtime.link_simulation WHERE time = %s",
"DELETE FROM realtime.link_results WHERE time = %s",
(target_time,),
)
# 2. 使用 COPY 快速写入新数据
async with cur.copy(
"COPY realtime.link_simulation (time, id, flow, friction, headloss, quality, reaction, setting, status, velocity) FROM STDIN"
"COPY realtime.link_results (time, link_id, flow, friction, headloss, quality, reaction, setting, status, velocity) FROM STDIN"
) as copy:
for item in data:
await copy.write_row(
@@ -49,7 +53,7 @@ class RealtimeRepository:
@staticmethod
def insert_links_batch_sync(conn: Connection, data: List[dict]):
"""Batch insert for realtime.link_simulation using DELETE then COPY for performance (sync version)."""
"""Synchronous batch insert for realtime.link_results."""
if not data:
return
@@ -59,15 +63,19 @@ class RealtimeRepository:
# 使用事务确保原子性
with conn.transaction():
with conn.cursor() as cur:
cur.execute(
"SELECT pg_advisory_xact_lock(hashtextextended(%s::text, 1))",
(target_time,),
)
# 1. 先删除该时间点的旧数据
cur.execute(
"DELETE FROM realtime.link_simulation WHERE time = %s",
"DELETE FROM realtime.link_results WHERE time = %s",
(target_time,),
)
# 2. 使用 COPY 快速写入新数据
with cur.copy(
"COPY realtime.link_simulation (time, id, flow, friction, headloss, quality, reaction, setting, status, velocity) FROM STDIN"
"COPY realtime.link_results (time, link_id, flow, friction, headloss, quality, reaction, setting, status, velocity) FROM STDIN"
) as copy:
for item in data:
copy.write_row(
@@ -91,7 +99,7 @@ class RealtimeRepository:
) -> List[dict]:
async with conn.cursor() as cur:
await cur.execute(
"SELECT * FROM realtime.link_simulation WHERE time >= %s AND time <= %s AND id = %s",
"SELECT * FROM realtime.link_results WHERE time >= %s AND time <= %s AND link_id = %s",
(start_time, end_time, link_id),
)
return await cur.fetchall()
@@ -104,7 +112,7 @@ 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_simulation WHERE time >= %s AND time <= %s",
"SELECT * FROM realtime.link_results WHERE time >= %s AND time <= %s",
(normalized_start_time, normalized_end_time),
)
return await cur.fetchall()
@@ -132,7 +140,7 @@ class RealtimeRepository:
raise ValueError(f"Invalid field: {field}")
query = sql.SQL(
"SELECT time, {} FROM realtime.link_simulation WHERE time >= %s AND time <= %s AND id = %s"
"SELECT time, {} FROM realtime.link_results WHERE time >= %s AND time <= %s AND link_id = %s"
).format(sql.Identifier(field))
async with conn.cursor() as cur:
@@ -164,7 +172,7 @@ class RealtimeRepository:
raise ValueError(f"Invalid field: {field}")
query = sql.SQL(
"SELECT id, time, {} FROM realtime.link_simulation WHERE time >= %s AND time <= %s"
"SELECT link_id, time, {} FROM realtime.link_results WHERE time >= %s AND time <= %s"
).format(sql.Identifier(field))
async with conn.cursor() as cur:
@@ -172,7 +180,7 @@ class RealtimeRepository:
rows = await cur.fetchall()
result = defaultdict(list)
for row in rows:
result[row["id"]].append(
result[row["link_id"]].append(
{"time": row["time"].isoformat(), "value": row[field]}
)
return dict(result)
@@ -199,7 +207,7 @@ class RealtimeRepository:
raise ValueError(f"Invalid field: {field}")
query = sql.SQL(
"UPDATE realtime.link_simulation SET {} = %s WHERE time = %s AND id = %s"
"UPDATE realtime.link_results SET {} = %s WHERE time = %s AND link_id = %s"
).format(sql.Identifier(field))
async with conn.cursor() as cur:
@@ -211,7 +219,7 @@ class RealtimeRepository:
):
async with conn.cursor() as cur:
await cur.execute(
"DELETE FROM realtime.link_simulation WHERE time >= %s AND time <= %s",
"DELETE FROM realtime.link_results WHERE time >= %s AND time <= %s",
(start_time, end_time),
)
@@ -228,15 +236,19 @@ class RealtimeRepository:
# 使用事务确保原子性
async with conn.transaction():
async with conn.cursor() as cur:
await cur.execute(
"SELECT pg_advisory_xact_lock(hashtextextended(%s::text, 1))",
(target_time,),
)
# 1. 先删除该时间点的旧数据
await cur.execute(
"DELETE FROM realtime.node_simulation WHERE time = %s",
"DELETE FROM realtime.node_results WHERE time = %s",
(target_time,),
)
# 2. 使用 COPY 快速写入新数据
async with cur.copy(
"COPY realtime.node_simulation (time, id, actual_demand, total_head, pressure, quality) FROM STDIN"
"COPY realtime.node_results (time, node_id, actual_demand, total_head, pressure, quality) FROM STDIN"
) as copy:
for item in data:
await copy.write_row(
@@ -261,15 +273,19 @@ class RealtimeRepository:
# 使用事务确保原子性
with conn.transaction():
with conn.cursor() as cur:
cur.execute(
"SELECT pg_advisory_xact_lock(hashtextextended(%s::text, 1))",
(target_time,),
)
# 1. 先删除该时间点的旧数据
cur.execute(
"DELETE FROM realtime.node_simulation WHERE time = %s",
"DELETE FROM realtime.node_results WHERE time = %s",
(target_time,),
)
# 2. 使用 COPY 快速写入新数据
with cur.copy(
"COPY realtime.node_simulation (time, id, actual_demand, total_head, pressure, quality) FROM STDIN"
"COPY realtime.node_results (time, node_id, actual_demand, total_head, pressure, quality) FROM STDIN"
) as copy:
for item in data:
copy.write_row(
@@ -289,7 +305,7 @@ class RealtimeRepository:
) -> 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",
"SELECT * FROM realtime.node_results WHERE time >= %s AND time <= %s AND node_id = %s",
(start_time, end_time, node_id),
)
return await cur.fetchall()
@@ -302,7 +318,7 @@ 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_simulation WHERE time >= %s AND time <= %s",
"SELECT * FROM realtime.node_results WHERE time >= %s AND time <= %s",
(normalized_start_time, normalized_end_time),
)
return await cur.fetchall()
@@ -320,7 +336,7 @@ class RealtimeRepository:
raise ValueError(f"Invalid field: {field}")
query = sql.SQL(
"SELECT time, {} FROM realtime.node_simulation WHERE time >= %s AND time <= %s AND id = %s"
"SELECT time, {} FROM realtime.node_results WHERE time >= %s AND time <= %s AND node_id = %s"
).format(sql.Identifier(field))
async with conn.cursor() as cur:
@@ -339,7 +355,7 @@ class RealtimeRepository:
raise ValueError(f"Invalid field: {field}")
query = sql.SQL(
"SELECT id, time, {} FROM realtime.node_simulation WHERE time >= %s AND time <= %s"
"SELECT node_id, time, {} FROM realtime.node_results WHERE time >= %s AND time <= %s"
).format(sql.Identifier(field))
async with conn.cursor() as cur:
@@ -347,7 +363,7 @@ class RealtimeRepository:
rows = await cur.fetchall()
result = defaultdict(list)
for row in rows:
result[row["id"]].append(
result[row["node_id"]].append(
{"time": row["time"].isoformat(), "value": row[field]}
)
return dict(result)
@@ -365,7 +381,7 @@ class RealtimeRepository:
raise ValueError(f"Invalid field: {field}")
query = sql.SQL(
"UPDATE realtime.node_simulation SET {} = %s WHERE time = %s AND id = %s"
"UPDATE realtime.node_results SET {} = %s WHERE time = %s AND node_id = %s"
).format(sql.Identifier(field))
async with conn.cursor() as cur:
@@ -377,7 +393,7 @@ class RealtimeRepository:
):
async with conn.cursor() as cur:
await cur.execute(
"DELETE FROM realtime.node_simulation WHERE time >= %s AND time <= %s",
"DELETE FROM realtime.node_results WHERE time >= %s AND time <= %s",
(start_time, end_time),
)
@@ -439,12 +455,15 @@ class RealtimeRepository:
}
)
# Insert data using batch methods
if node_data:
await RealtimeRepository.insert_nodes_batch(conn, node_data)
# Keep node and link replacement atomic. The batch helpers use nested
# transactions (savepoints), while this outer transaction guarantees
# that a link write failure also rolls back the node replacement.
async with conn.transaction():
if node_data:
await RealtimeRepository.insert_nodes_batch(conn, node_data)
if link_data:
await RealtimeRepository.insert_links_batch(conn, link_data)
if link_data:
await RealtimeRepository.insert_links_batch(conn, link_data)
@staticmethod
def store_realtime_simulation_result_sync(
@@ -502,12 +521,15 @@ class RealtimeRepository:
}
)
# Insert data using batch methods
if node_data:
RealtimeRepository.insert_nodes_batch_sync(conn, node_data)
# Keep node and link replacement atomic. The batch helpers use nested
# transactions (savepoints), while this outer transaction guarantees
# that a link write failure also rolls back the node replacement.
with conn.transaction():
if node_data:
RealtimeRepository.insert_nodes_batch_sync(conn, node_data)
if link_data:
RealtimeRepository.insert_links_batch_sync(conn, link_data)
if link_data:
RealtimeRepository.insert_links_batch_sync(conn, link_data)
@staticmethod
async def query_all_record_by_time_property(
app/infra/db/timescaledb/repositories/scada.py
+9 -9
View File
@@ -14,7 +14,7 @@ class ScadaRepository:
async with conn.cursor() as cur:
async with cur.copy(
"COPY scada.scada_data (time, device_id, monitored_value, cleaned_value) FROM STDIN"
"COPY scada.measurements (time, device_id, monitored_value, cleaned_value) FROM STDIN"
) as copy:
for item in data:
await copy.write_row(
@@ -35,7 +35,7 @@ class ScadaRepository:
) -> List[dict]:
async with conn.cursor() as cur:
await cur.execute(
"SELECT * FROM scada.scada_data WHERE device_id = ANY(%s) AND time >= %s AND time <= %s",
"SELECT * FROM scada.measurements WHERE device_id = ANY(%s) AND time >= %s AND time <= %s",
(device_ids, start_time, end_time),
)
return await cur.fetchall()
@@ -49,7 +49,7 @@ class ScadaRepository:
) -> List[dict]:
with conn.cursor(row_factory=dict_row) as cur:
cur.execute(
"SELECT * FROM scada.scada_data WHERE device_id = ANY(%s) AND time >= %s AND time <= %s",
"SELECT * FROM scada.measurements WHERE device_id = ANY(%s) AND time >= %s AND time <= %s",
(device_ids, start_time, end_time),
)
return cur.fetchall()
@@ -63,12 +63,12 @@ class ScadaRepository:
with conn.cursor(row_factory=dict_row) as cur:
if before_time is None:
cur.execute(
"SELECT max(time) AS time FROM scada.scada_data WHERE device_id = ANY(%s)",
"SELECT max(time) AS time FROM scada.measurements WHERE device_id = ANY(%s)",
(device_ids,),
)
else:
cur.execute(
"SELECT max(time) AS time FROM scada.scada_data "
"SELECT max(time) AS time FROM scada.measurements "
"WHERE device_id = ANY(%s) AND time <= %s",
(device_ids, before_time),
)
@@ -88,7 +88,7 @@ class ScadaRepository:
raise ValueError(f"Invalid field: {field}")
query = sql.SQL(
"SELECT device_id, time, {} FROM scada.scada_data WHERE time >= %s AND time <= %s AND device_id = ANY(%s)"
"SELECT device_id, time, {} FROM scada.measurements WHERE time >= %s AND time <= %s AND device_id = ANY(%s)"
).format(sql.Identifier(field))
async with conn.cursor() as cur:
@@ -111,10 +111,10 @@ class ScadaRepository:
raise ValueError(f"Invalid field: {field}")
update_query = sql.SQL(
"UPDATE scada.scada_data SET {} = %s WHERE time = %s AND device_id = %s"
"UPDATE scada.measurements SET {} = %s WHERE time = %s AND device_id = %s"
).format(sql.Identifier(field))
insert_query = sql.SQL(
"INSERT INTO scada.scada_data (time, device_id, {}) VALUES (%s, %s, %s)"
"INSERT INTO scada.measurements (time, device_id, {}) VALUES (%s, %s, %s)"
).format(sql.Identifier(field))
async with conn.cursor() as cur:
@@ -128,6 +128,6 @@ class ScadaRepository:
):
async with conn.cursor() as cur:
await cur.execute(
"DELETE FROM scada.scada_data WHERE device_id = %s AND time >= %s AND time <= %s",
"DELETE FROM scada.measurements WHERE device_id = %s AND time >= %s AND time <= %s",
(device_id, start_time, end_time),
)
app/infra/db/timescaledb/repositories/scheme.py
-710
View File
@@ -1,710 +0,0 @@
from typing import List, Any, Dict
from datetime import datetime, timedelta
from collections import defaultdict
from psycopg import AsyncConnection, Connection, sql
import app.services.globals as globals
from app.services.time_api import parse_clock_duration_seconds, parse_utc_time
class SchemeRepository:
@staticmethod
def _get_result_timestep(result_timestep_seconds: int | None) -> timedelta:
if result_timestep_seconds is not None:
if result_timestep_seconds <= 0:
raise ValueError("result_timestep_seconds must be greater than 0.")
return timedelta(seconds=result_timestep_seconds)
timestep_seconds = parse_clock_duration_seconds(
globals.hydraulic_timestep,
field_name="HYDRAULIC TIMESTEP",
)
if timestep_seconds <= 0:
raise ValueError("HYDRAULIC TIMESTEP must be greater than 0.")
return timedelta(seconds=timestep_seconds)
# --- Link Simulation ---
@staticmethod
async def insert_links_batch(conn: AsyncConnection, data: List[dict]):
"""Batch insert for scheme.link_simulation using DELETE then COPY for performance."""
if not data:
return
# 获取批次中所有不同的时间点
all_times = list(set(item["time"] for item in data))
target_scheme_type = data[0]["scheme_type"]
target_scheme_name = data[0]["scheme_name"]
# 使用事务确保原子性
async with conn.transaction():
async with conn.cursor() as cur:
# 1. 删除该批次涉及的所有时间点、scheme_type、scheme_name 的旧数据
await cur.execute(
"DELETE FROM scheme.link_simulation WHERE time = ANY(%s) AND scheme_type = %s AND scheme_name = %s",
(all_times, target_scheme_type, target_scheme_name),
)
# 2. 使用 COPY 快速写入新数据
async with cur.copy(
"COPY scheme.link_simulation (time, scheme_type, scheme_name, 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_type"],
item["scheme_name"],
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
def insert_links_batch_sync(conn: Connection, data: List[dict]):
"""Batch insert for scheme.link_simulation using DELETE then COPY for performance (sync version)."""
if not data:
return
# 获取批次中所有不同的时间点
all_times = list(set(item["time"] for item in data))
target_scheme_type = data[0]["scheme_type"]
target_scheme_name = data[0]["scheme_name"]
# 使用事务确保原子性
with conn.transaction():
with conn.cursor() as cur:
# 1. 删除该批次涉及的所有时间点、scheme_type、scheme_name 的旧数据
cur.execute(
"DELETE FROM scheme.link_simulation WHERE time = ANY(%s) AND scheme_type = %s AND scheme_name = %s",
(all_times, target_scheme_type, target_scheme_name),
)
# 2. 使用 COPY 快速写入新数据
with cur.copy(
"COPY scheme.link_simulation (time, scheme_type, scheme_name, id, flow, friction, headloss, quality, reaction, setting, status, velocity) FROM STDIN"
) as copy:
for item in data:
copy.write_row(
(
item["time"],
item["scheme_type"],
item["scheme_name"],
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_link_by_scheme_and_time_range(
conn: AsyncConnection,
scheme_type: str,
scheme_name: 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_type = %s AND scheme_name = %s AND time >= %s AND time <= %s AND id = %s",
(scheme_type, scheme_name, start_time, end_time, link_id),
)
return await cur.fetchall()
@staticmethod
async def get_links_by_scheme_and_time_range(
conn: AsyncConnection,
scheme_type: str,
scheme_name: 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_type = %s AND scheme_name = %s AND time >= %s AND time <= %s",
(scheme_type, scheme_name, start_time, end_time),
)
return await cur.fetchall()
@staticmethod
async def get_link_field_by_scheme_and_time_range(
conn: AsyncConnection,
scheme_type: str,
scheme_name: str,
start_time: datetime,
end_time: datetime,
link_id: str,
field: str,
) -> List[Dict[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 time, {} FROM scheme.link_simulation WHERE scheme_type = %s AND scheme_name = %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_type, scheme_name, start_time, end_time, link_id)
)
rows = await cur.fetchall()
return [
{"time": row["time"].isoformat(), "value": row[field]} for row in rows
]
@staticmethod
async def get_links_field_by_scheme_and_time_range(
conn: AsyncConnection,
scheme_type: str,
scheme_name: str,
start_time: datetime,
end_time: datetime,
field: str,
) -> dict:
# 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 id, time, {} FROM scheme.link_simulation WHERE scheme_type = %s AND scheme_name = %s AND time >= %s AND time <= %s"
).format(sql.Identifier(field))
async with conn.cursor() as cur:
await cur.execute(query, (scheme_type, scheme_name, start_time, end_time))
rows = await cur.fetchall()
result = defaultdict(list)
for row in rows:
result[row["id"]].append(
{"time": row["time"].isoformat(), "value": row[field]}
)
return dict(result)
@staticmethod
async def update_link_field(
conn: AsyncConnection,
time: datetime,
scheme_type: str,
scheme_name: 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_type = %s AND scheme_name = %s AND id = %s"
).format(sql.Identifier(field))
async with conn.cursor() as cur:
await cur.execute(query, (value, time, scheme_type, scheme_name, link_id))
@staticmethod
async def delete_links_by_scheme_and_time_range(
conn: AsyncConnection,
scheme_type: str,
scheme_name: str,
start_time: datetime,
end_time: datetime,
):
async with conn.cursor() as cur:
await cur.execute(
"DELETE FROM scheme.link_simulation WHERE scheme_type = %s AND scheme_name = %s AND time >= %s AND time <= %s",
(scheme_type, scheme_name, start_time, end_time),
)
# --- Node Simulation ---
@staticmethod
async def insert_nodes_batch(conn: AsyncConnection, data: List[dict]):
if not data:
return
# 获取批次中所有不同的时间点
all_times = list(set(item["time"] for item in data))
target_scheme_type = data[0]["scheme_type"]
target_scheme_name = data[0]["scheme_name"]
# 使用事务确保原子性
async with conn.transaction():
async with conn.cursor() as cur:
# 1. 删除该批次涉及的所有时间点、scheme_type、scheme_name 的旧数据
await cur.execute(
"DELETE FROM scheme.node_simulation WHERE time = ANY(%s) AND scheme_type = %s AND scheme_name = %s",
(all_times, target_scheme_type, target_scheme_name),
)
# 2. 使用 COPY 快速写入新数据
async with cur.copy(
"COPY scheme.node_simulation (time, scheme_type, scheme_name, id, actual_demand, total_head, pressure, quality) FROM STDIN"
) as copy:
for item in data:
await copy.write_row(
(
item["time"],
item["scheme_type"],
item["scheme_name"],
item["id"],
item.get("actual_demand"),
item.get("total_head"),
item.get("pressure"),
item.get("quality"),
)
)
@staticmethod
def insert_nodes_batch_sync(conn: Connection, data: List[dict]):
if not data:
return
# 获取批次中所有不同的时间点
all_times = list(set(item["time"] for item in data))
target_scheme_type = data[0]["scheme_type"]
target_scheme_name = data[0]["scheme_name"]
# 使用事务确保原子性
with conn.transaction():
with conn.cursor() as cur:
# 1. 删除该批次涉及的所有时间点、scheme_type、scheme_name 的旧数据
cur.execute(
"DELETE FROM scheme.node_simulation WHERE time = ANY(%s) AND scheme_type = %s AND scheme_name = %s",
(all_times, target_scheme_type, target_scheme_name),
)
# 2. 使用 COPY 快速写入新数据
with cur.copy(
"COPY scheme.node_simulation (time, scheme_type, scheme_name, id, actual_demand, total_head, pressure, quality) FROM STDIN"
) as copy:
for item in data:
copy.write_row(
(
item["time"],
item["scheme_type"],
item["scheme_name"],
item["id"],
item.get("actual_demand"),
item.get("total_head"),
item.get("pressure"),
item.get("quality"),
)
)
@staticmethod
async def get_node_by_scheme_and_time_range(
conn: AsyncConnection,
scheme_type: str,
scheme_name: 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_type = %s AND scheme_name = %s AND time >= %s AND time <= %s AND id = %s",
(scheme_type, scheme_name, start_time, end_time, node_id),
)
return await cur.fetchall()
@staticmethod
async def get_nodes_by_scheme_and_time_range(
conn: AsyncConnection,
scheme_type: str,
scheme_name: 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_type = %s AND scheme_name = %s AND time >= %s AND time <= %s",
(scheme_type, scheme_name, start_time, end_time),
)
return await cur.fetchall()
@staticmethod
async def get_node_field_by_scheme_and_time_range(
conn: AsyncConnection,
scheme_type: str,
scheme_name: str,
start_time: datetime,
end_time: datetime,
node_id: str,
field: str,
) -> List[Dict[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 time, {} FROM scheme.node_simulation WHERE scheme_type = %s AND scheme_name = %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_type, scheme_name, start_time, end_time, node_id)
)
rows = await cur.fetchall()
return [
{"time": row["time"].isoformat(), "value": row[field]} for row in rows
]
@staticmethod
async def get_nodes_field_by_scheme_and_time_range(
conn: AsyncConnection,
scheme_type: str,
scheme_name: str,
start_time: datetime,
end_time: datetime,
field: str,
) -> dict:
# 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 id, time, {} FROM scheme.node_simulation WHERE scheme_type = %s AND scheme_name = %s AND time >= %s AND time <= %s"
).format(sql.Identifier(field))
async with conn.cursor() as cur:
await cur.execute(query, (scheme_type, scheme_name, start_time, end_time))
rows = await cur.fetchall()
result = defaultdict(list)
for row in rows:
result[row["id"]].append(
{"time": row["time"].isoformat(), "value": row[field]}
)
return dict(result)
@staticmethod
async def update_node_field(
conn: AsyncConnection,
time: datetime,
scheme_type: str,
scheme_name: 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_type = %s AND scheme_name = %s AND id = %s"
).format(sql.Identifier(field))
async with conn.cursor() as cur:
await cur.execute(query, (value, time, scheme_type, scheme_name, node_id))
@staticmethod
async def delete_nodes_by_scheme_and_time_range(
conn: AsyncConnection,
scheme_type: str,
scheme_name: str,
start_time: datetime,
end_time: datetime,
):
async with conn.cursor() as cur:
await cur.execute(
"DELETE FROM scheme.node_simulation WHERE scheme_type = %s AND scheme_name = %s AND time >= %s AND time <= %s",
(scheme_type, scheme_name, start_time, end_time),
)
# --- 复合查询 ---
@staticmethod
async def store_scheme_simulation_result(
conn: AsyncConnection,
scheme_type: str,
scheme_name: str,
node_result_list: List[Dict[str, any]],
link_result_list: List[Dict[str, any]],
result_start_time: str,
num_periods: int = 1,
result_timestep_seconds: int | None = None,
):
"""
Store scheme simulation results to TimescaleDB.
Args:
conn: Database connection
scheme_type: Scheme type
scheme_name: 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)
"""
simulation_time = parse_utc_time(
result_start_time, field_name="result_start_time"
)
timestep = SchemeRepository._get_result_timestep(result_timestep_seconds)
# Prepare node data for batch insert
node_data = []
for node_result in node_result_list:
node_id = node_result.get("node")
result_rows = node_result.get("result", [])
for period_index in range(min(num_periods, len(result_rows))):
current_time = simulation_time + (timestep * period_index)
data = result_rows[period_index]
node_data.append(
{
"time": current_time,
"scheme_type": scheme_type,
"scheme_name": scheme_name,
"id": node_id,
"actual_demand": data.get("demand"),
"total_head": data.get("head"),
"pressure": data.get("pressure"),
"quality": data.get("quality"),
}
)
# Prepare link data for batch insert
link_data = []
for link_result in link_result_list:
link_id = link_result.get("link")
result_rows = link_result.get("result", [])
for period_index in range(min(num_periods, len(result_rows))):
current_time = simulation_time + (timestep * period_index)
data = result_rows[period_index]
link_data.append(
{
"time": current_time,
"scheme_type": scheme_type,
"scheme_name": scheme_name,
"id": link_id,
"flow": data.get("flow"),
"friction": data.get("friction"),
"headloss": data.get("headloss"),
"quality": data.get("quality"),
"reaction": data.get("reaction"),
"setting": data.get("setting"),
"status": data.get("status"),
"velocity": data.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
def store_scheme_simulation_result_sync(
conn: Connection,
scheme_type: str,
scheme_name: str,
node_result_list: List[Dict[str, any]],
link_result_list: List[Dict[str, any]],
result_start_time: str,
num_periods: int = 1,
result_timestep_seconds: int | None = None,
):
"""
Store scheme simulation results to TimescaleDB (sync version).
Args:
conn: Database connection
scheme_type: Scheme type
scheme_name: 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)
"""
simulation_time = parse_utc_time(
result_start_time, field_name="result_start_time"
)
timestep = SchemeRepository._get_result_timestep(result_timestep_seconds)
# Prepare node data for batch insert
node_data = []
for node_result in node_result_list:
node_id = node_result.get("node")
result_rows = node_result.get("result", [])
for period_index in range(min(num_periods, len(result_rows))):
current_time = simulation_time + (timestep * period_index)
data = result_rows[period_index]
node_data.append(
{
"time": current_time,
"scheme_type": scheme_type,
"scheme_name": scheme_name,
"id": node_id,
"actual_demand": data.get("demand"),
"total_head": data.get("head"),
"pressure": data.get("pressure"),
"quality": data.get("quality"),
}
)
# Prepare link data for batch insert
link_data = []
for link_result in link_result_list:
link_id = link_result.get("link")
result_rows = link_result.get("result", [])
for period_index in range(min(num_periods, len(result_rows))):
current_time = simulation_time + (timestep * period_index)
data = result_rows[period_index]
link_data.append(
{
"time": current_time,
"scheme_type": scheme_type,
"scheme_name": scheme_name,
"id": link_id,
"flow": data.get("flow"),
"friction": data.get("friction"),
"headloss": data.get("headloss"),
"quality": data.get("quality"),
"reaction": data.get("reaction"),
"setting": data.get("setting"),
"status": data.get("status"),
"velocity": data.get("velocity"),
}
)
# Insert data using batch methods
if node_data:
SchemeRepository.insert_nodes_batch_sync(conn, node_data)
if link_data:
SchemeRepository.insert_links_batch_sync(conn, link_data)
@staticmethod
async def query_all_record_by_scheme_time_property(
conn: AsyncConnection,
scheme_type: str,
scheme_name: str,
query_time: str,
type: str,
property: str,
) -> list:
"""
Query all records by scheme, time and property from TimescaleDB.
Args:
conn: Database connection
scheme_type: Scheme type
scheme_name: 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
"""
target_time = parse_utc_time(query_time, field_name="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":
data = await SchemeRepository.get_nodes_field_by_scheme_and_time_range(
conn, scheme_type, scheme_name, start_time, end_time, property
)
elif type.lower() == "link":
data = await SchemeRepository.get_links_field_by_scheme_and_time_range(
conn, scheme_type, scheme_name, start_time, end_time, property
)
else:
raise ValueError(f"Invalid type: {type}. Must be 'node' or 'link'")
# Format the results
# Format the results
result = []
for id, items in data.items():
for item in items:
result.append({"ID": id, "value": item["value"]})
return result
@staticmethod
async def query_scheme_simulation_result_by_id_time(
conn: AsyncConnection,
scheme_type: str,
scheme_name: 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_type: Scheme type
scheme_name: 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
"""
target_time = parse_utc_time(query_time, field_name="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_type, scheme_name, start_time, end_time, id
)
elif type.lower() == "link":
return await SchemeRepository.get_link_by_scheme_and_time_range(
conn, scheme_type, scheme_name, start_time, end_time, id
)
else:
raise ValueError(f"Invalid type: {type}. Must be 'node' or 'link'")