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

refactor(db)!: finalize pooled WNDB v2 migration

Browse Source
This commit is contained in:
jiang
2026-08-27 17:26:22 +08:00
parent fa188af0b1
commit b74799a39d
105 changed files with 4988 additions and 5565 deletions
Download Patch File Download Diff File Expand all files Collapse all files
app/infra/db/timescaledb/composite_queries.py
+111 -83
View File
@@ -56,38 +56,40 @@ class CompositeQueries:
Raises:
ValueError: 当 SCADA 设备未找到或字段无效时
"""
result = {}
scada_by_id = await CompositeQueries._get_project_scada_index(postgres_conn)
link_devices: dict[str, str] = {}
node_devices: dict[str, str] = {}
for device_id in device_ids:
target_scada = scada_by_id.get(device_id)
if not target_scada:
raise ValueError(f"SCADA device {device_id} not found")
scada_type = target_scada["device_type"]
element_id = (
target_scada["link_id"]
if scada_type in {"pipe_flow", "flow"}
else target_scada["node_id"]
)
if scada_type == "pipe_flow":
# 查询 link 模拟数据
res = await RealtimeRepository.get_link_field_by_time_range(
timescale_conn, start_time, end_time, element_id, "flow"
)
if scada_type in {"pipe_flow", "flow"}:
link_devices[device_id] = target_scada["link_id"]
elif scada_type == "pressure":
# 查询 node 模拟数据
res = await RealtimeRepository.get_node_field_by_time_range(
timescale_conn, start_time, end_time, element_id, "pressure"
)
node_devices[device_id] = target_scada["node_id"]
else:
raise ValueError(f"Unknown SCADA type: {scada_type}")
# 添加 scada_id 到每个数据项
for item in res:
item["scada_id"] = device_id
result[device_id] = res
return result
link_series = await RealtimeRepository.get_link_fields_by_ids_time_range(
timescale_conn, start_time, end_time,
list(dict.fromkeys(link_devices.values())), "flow",
)
node_series = await RealtimeRepository.get_node_fields_by_ids_time_range(
timescale_conn, start_time, end_time,
list(dict.fromkeys(node_devices.values())), "pressure",
)
return {
device_id: [
{**item, "scada_id": device_id}
for item in (
link_series.get(element_id, [])
if device_id in link_devices
else node_series.get(element_id, [])
)
]
for device_id, element_id in (link_devices | node_devices).items()
}
@staticmethod
async def get_scada_associated_analysis_simulation_data(
@@ -117,38 +119,41 @@ class CompositeQueries:
Raises:
ValueError: 当 SCADA 设备未找到或字段无效时
"""
result = {}
scada_by_id = await CompositeQueries._get_project_scada_index(postgres_conn)
link_devices: dict[str, str] = {}
node_devices: dict[str, str] = {}
for device_id in device_ids:
target_scada = scada_by_id.get(device_id)
if not target_scada:
raise ValueError(f"SCADA device {device_id} not found")
scada_type = target_scada["device_type"]
element_id = (
target_scada["link_id"]
if scada_type in {"pipe_flow", "flow"}
else target_scada["node_id"]
)
if scada_type == "pipe_flow":
# 查询 link 模拟数据
res = await AnalysisResultsRepository.get_link_series(
timescale_conn, run_id, element_id, start_time, end_time, "flow"
)
if scada_type in {"pipe_flow", "flow"}:
link_devices[device_id] = target_scada["link_id"]
elif scada_type == "pressure":
# 查询 node 模拟数据
res = await AnalysisResultsRepository.get_node_series(
timescale_conn, run_id, element_id, start_time, end_time, "pressure"
)
node_devices[device_id] = target_scada["node_id"]
else:
raise ValueError(f"Unknown SCADA type: {scada_type}")
# 添加 scada_id 到每个数据项
for item in res:
item["scada_id"] = device_id
result[device_id] = res
return result
link_series = await AnalysisResultsRepository.get_series_by_ids(
timescale_conn, run_id, "link",
list(dict.fromkeys(link_devices.values())), start_time, end_time, "flow",
)
node_series = await AnalysisResultsRepository.get_series_by_ids(
timescale_conn, run_id, "node",
list(dict.fromkeys(node_devices.values())), start_time, end_time, "pressure",
)
return {
device_id: [
{**item, "scada_id": device_id}
for item in (
link_series.get(element_id, [])
if device_id in link_devices
else node_series.get(element_id, [])
)
]
for device_id, element_id in (link_devices | node_devices).items()
}
@staticmethod
async def get_realtime_simulation_data(
@@ -175,26 +180,33 @@ class CompositeQueries:
Raises:
ValueError: 当 SCADA 设备未找到或字段无效时
"""
result = {}
pipe_ids: list[str] = []
junction_ids: list[str] = []
for feature_id, feature_type in feature_infos:
if feature_type.lower() == "pipe":
# 查询 link 模拟数据
res = await RealtimeRepository.get_link_field_by_time_range(
timescale_conn, start_time, end_time, feature_id, "flow"
)
pipe_ids.append(feature_id)
elif feature_type.lower() == "junction":
# 查询 node 模拟数据
res = await RealtimeRepository.get_node_field_by_time_range(
timescale_conn, start_time, end_time, feature_id, "pressure"
)
junction_ids.append(feature_id)
else:
raise ValueError(f"Unknown type: {feature_type}")
# 添加 scada_id 到每个数据项
for item in res:
item["feature_id"] = feature_id
result[feature_id] = res
return result
link_series = await RealtimeRepository.get_link_fields_by_ids_time_range(
timescale_conn, start_time, end_time, list(dict.fromkeys(pipe_ids)), "flow"
)
node_series = await RealtimeRepository.get_node_fields_by_ids_time_range(
timescale_conn, start_time, end_time,
list(dict.fromkeys(junction_ids)), "pressure",
)
return {
feature_id: [
{**item, "feature_id": feature_id}
for item in (
link_series.get(feature_id, [])
if feature_type.lower() == "pipe"
else node_series.get(feature_id, [])
)
]
for feature_id, feature_type in feature_infos
}
@staticmethod
async def get_analysis_simulation_data(
@@ -223,25 +235,34 @@ class CompositeQueries:
Raises:
ValueError: 当类型无效时
"""
result = {}
pipe_ids: list[str] = []
junction_ids: list[str] = []
for feature_id, feature_type in feature_infos:
if feature_type.lower() == "pipe":
# 查询 link 模拟数据
res = await AnalysisResultsRepository.get_link_series(
timescale_conn, run_id, feature_id, start_time, end_time, "flow"
)
pipe_ids.append(feature_id)
elif feature_type.lower() == "junction":
# 查询 node 模拟数据
res = await AnalysisResultsRepository.get_node_series(
timescale_conn, run_id, feature_id, start_time, end_time, "pressure"
)
junction_ids.append(feature_id)
else:
raise ValueError(f"Unknown type: {feature_type}")
# 添加 feature_id 到每个数据项
for item in res:
item["feature_id"] = feature_id
result[feature_id] = res
return result
link_series = await AnalysisResultsRepository.get_series_by_ids(
timescale_conn, run_id, "link", list(dict.fromkeys(pipe_ids)),
start_time, end_time, "flow",
)
node_series = await AnalysisResultsRepository.get_series_by_ids(
timescale_conn, run_id, "node", list(dict.fromkeys(junction_ids)),
start_time, end_time, "pressure",
)
return {
feature_id: [
{**item, "feature_id": feature_id}
for item in (
link_series.get(feature_id, [])
if feature_type.lower() == "pipe"
else node_series.get(feature_id, [])
)
]
for feature_id, feature_type in feature_infos
}
@staticmethod
async def get_element_associated_scada_data(
@@ -399,7 +420,7 @@ class CompositeQueries:
if scada_by_id[device_id]["device_type"] in {"pipe_flow", "flow"}
]
updated_rows = 0
cleaned_rows: list[tuple[datetime, str, float | None]] = []
for grouped_ids, cleaning_function in (
(pressure_ids, clean_pressure_data_df_km),
(flow_ids, clean_flow_data_df_kf),
@@ -422,18 +443,25 @@ class CompositeQueries:
if isinstance(time_value, datetime)
else datetime.fromisoformat(str(time_value))
)
await ScadaRepository.update_scada_field(
timescale_conn,
time_dt,
device_id,
"cleaned_value",
value,
cleaned_rows.append(
(
time_dt,
device_id,
None if pd.isna(value) else float(value),
)
)
updated_rows += 1
if updated_rows == 0:
if not cleaned_rows:
raise ValueError("SCADA 数据清洗未产生任何数据库更新")
updated_rows = await ScadaRepository.update_scada_field_batch(
timescale_conn,
cleaned_rows,
"cleaned_value",
)
if updated_rows == 0:
raise ValueError("SCADA 清洗结果未匹配任何已有监测数据")
return "success"
@staticmethod
app/infra/db/timescaledb/internal_queries.py
+3 -2
View File
@@ -259,7 +259,7 @@ class InternalQueries:
query = sql.SQL(
"SELECT btrim({}::text) AS id, time, {} FROM {}.{} "
"WHERE run_id = %s AND time >= %s AND time <= %s "
"AND btrim({}::text) = ANY(%s)"
"AND btrim({}::text) = ANY(%s) ORDER BY id, time"
).format(
sql.Identifier(id_column),
sql.Identifier(field),
@@ -274,7 +274,8 @@ class InternalQueries:
else:
query = sql.SQL(
"SELECT btrim({}::text) AS id, time, {} FROM {}.{} "
"WHERE time >= %s AND time <= %s AND btrim({}::text) = ANY(%s)"
"WHERE time >= %s AND time <= %s "
"AND btrim({}::text) = ANY(%s) ORDER BY id, time"
).format(
sql.Identifier(id_column),
sql.Identifier(field),
app/infra/db/timescaledb/repositories/analysis.py
+47
View File
@@ -175,6 +175,53 @@ class AnalysisResultsRepository:
await cur.execute(query, (run_id, link_id, start_time, end_time))
return await cur.fetchall()
@staticmethod
async def get_series_by_ids(
conn: AsyncConnection,
run_id: UUID,
element_type: str,
element_ids: list[str],
start_time: datetime,
end_time: datetime,
field: str,
) -> dict[str, list[dict[str, Any]]]:
if element_type == "node":
table_name, id_column, valid_fields = (
"node_results", "node_id", AnalysisResultsRepository.NODE_FIELDS
)
elif element_type == "link":
table_name, id_column, valid_fields = (
"link_results", "link_id", AnalysisResultsRepository.LINK_FIELDS
)
else:
raise ValueError(f"invalid analysis element type: {element_type}")
if field not in valid_fields:
raise ValueError(f"invalid {element_type} result field: {field}")
result: dict[str, list[dict[str, Any]]] = {
element_id: [] for element_id in element_ids
}
if not element_ids:
return result
query = sql.SQL(
"SELECT {} AS element_id, time, {} AS value FROM analysis.{} "
"WHERE run_id = %s AND {} = ANY(%s) AND time BETWEEN %s AND %s "
"ORDER BY {}, time"
).format(
sql.Identifier(id_column),
sql.Identifier(field),
sql.Identifier(table_name),
sql.Identifier(id_column),
sql.Identifier(id_column),
)
async with conn.cursor() as cur:
await cur.execute(query, (run_id, element_ids, start_time, end_time))
for row in await cur.fetchall():
result.setdefault(str(row["element_id"]), []).append(
{"time": row["time"], "value": row["value"]}
)
return result
@staticmethod
async def get_values_at_time(
conn: AsyncConnection,
app/infra/db/timescaledb/repositories/realtime.py
+120 -20
View File
@@ -7,6 +7,19 @@ from app.services.time_api import parse_utc_time
class RealtimeRepository:
@staticmethod
def _batch_time(data: List[dict]) -> datetime:
"""Return one normalized timestamp shared by every row in a snapshot."""
if not data:
raise ValueError("Realtime batch must not be empty")
times = {
parse_utc_time(item["time"], field_name="time")
for item in data
}
if len(times) != 1:
raise ValueError("Realtime batch must contain exactly one timestamp")
return times.pop()
# --- Link Simulation ---
@staticmethod
@@ -15,8 +28,7 @@ class RealtimeRepository:
if not data:
return
# 假设同一批次的数据时间是相同的
target_time = data[0]["time"]
target_time = RealtimeRepository._batch_time(data)
# 使用事务确保原子性
async with conn.transaction():
@@ -38,7 +50,7 @@ class RealtimeRepository:
for item in data:
await copy.write_row(
(
item["time"],
target_time,
item["id"],
item.get("flow"),
item.get("friction"),
@@ -57,8 +69,7 @@ class RealtimeRepository:
if not data:
return
# 假设同一批次的数据时间是相同的
target_time = data[0]["time"]
target_time = RealtimeRepository._batch_time(data)
# 使用事务确保原子性
with conn.transaction():
@@ -80,7 +91,7 @@ class RealtimeRepository:
for item in data:
copy.write_row(
(
item["time"],
target_time,
item["id"],
item.get("flow"),
item.get("friction"),
@@ -99,7 +110,8 @@ class RealtimeRepository:
) -> List[dict]:
async with conn.cursor() as cur:
await cur.execute(
"SELECT * FROM realtime.link_results WHERE time >= %s AND time <= %s AND link_id = %s",
"SELECT * FROM realtime.link_results WHERE time >= %s AND time <= %s "
"AND link_id = %s ORDER BY time",
(start_time, end_time, link_id),
)
return await cur.fetchall()
@@ -112,7 +124,8 @@ 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_results WHERE time >= %s AND time <= %s",
"SELECT * FROM realtime.link_results WHERE time >= %s AND time <= %s "
"ORDER BY time, link_id",
(normalized_start_time, normalized_end_time),
)
return await cur.fetchall()
@@ -140,7 +153,8 @@ class RealtimeRepository:
raise ValueError(f"Invalid field: {field}")
query = sql.SQL(
"SELECT time, {} FROM realtime.link_results WHERE time >= %s AND time <= %s AND link_id = %s"
"SELECT time, {} FROM realtime.link_results WHERE time >= %s "
"AND time <= %s AND link_id = %s ORDER BY time"
).format(sql.Identifier(field))
async with conn.cursor() as cur:
@@ -150,6 +164,36 @@ class RealtimeRepository:
{"time": row["time"].isoformat(), "value": row[field]} for row in rows
]
@staticmethod
async def get_link_fields_by_ids_time_range(
conn: AsyncConnection,
start_time: datetime,
end_time: datetime,
link_ids: list[str],
field: str,
) -> dict[str, list[dict[str, Any]]]:
valid_fields = {
"flow", "friction", "headloss", "quality", "reaction",
"setting", "status", "velocity",
}
if field not in valid_fields:
raise ValueError(f"Invalid field: {field}")
result = {link_id: [] for link_id in link_ids}
if not link_ids:
return result
query = sql.SQL(
"SELECT link_id, time, {} FROM realtime.link_results "
"WHERE time BETWEEN %s AND %s AND link_id = ANY(%s) "
"ORDER BY link_id, time"
).format(sql.Identifier(field))
async with conn.cursor() as cur:
await cur.execute(query, (start_time, end_time, link_ids))
for row in await cur.fetchall():
result.setdefault(str(row["link_id"]), []).append(
{"time": row["time"].isoformat(), "value": row[field]}
)
return result
@staticmethod
async def get_links_field_by_time_range(
conn: AsyncConnection,
@@ -172,7 +216,8 @@ class RealtimeRepository:
raise ValueError(f"Invalid field: {field}")
query = sql.SQL(
"SELECT link_id, time, {} FROM realtime.link_results WHERE time >= %s AND time <= %s"
"SELECT link_id, time, {} FROM realtime.link_results "
"WHERE time >= %s AND time <= %s ORDER BY link_id, time"
).format(sql.Identifier(field))
async with conn.cursor() as cur:
@@ -230,8 +275,7 @@ class RealtimeRepository:
if not data:
return
# 假设同一批次的数据时间是相同的
target_time = data[0]["time"]
target_time = RealtimeRepository._batch_time(data)
# 使用事务确保原子性
async with conn.transaction():
@@ -253,7 +297,7 @@ class RealtimeRepository:
for item in data:
await copy.write_row(
(
item["time"],
target_time,
item["id"],
item.get("actual_demand"),
item.get("total_head"),
@@ -267,8 +311,7 @@ class RealtimeRepository:
if not data:
return
# 假设同一批次的数据时间是相同的
target_time = data[0]["time"]
target_time = RealtimeRepository._batch_time(data)
# 使用事务确保原子性
with conn.transaction():
@@ -290,7 +333,7 @@ class RealtimeRepository:
for item in data:
copy.write_row(
(
item["time"],
target_time,
item["id"],
item.get("actual_demand"),
item.get("total_head"),
@@ -305,7 +348,8 @@ class RealtimeRepository:
) -> List[dict]:
async with conn.cursor() as cur:
await cur.execute(
"SELECT * FROM realtime.node_results WHERE time >= %s AND time <= %s AND node_id = %s",
"SELECT * FROM realtime.node_results WHERE time >= %s AND time <= %s "
"AND node_id = %s ORDER BY time",
(start_time, end_time, node_id),
)
return await cur.fetchall()
@@ -318,7 +362,8 @@ 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_results WHERE time >= %s AND time <= %s",
"SELECT * FROM realtime.node_results WHERE time >= %s AND time <= %s "
"ORDER BY time, node_id",
(normalized_start_time, normalized_end_time),
)
return await cur.fetchall()
@@ -336,7 +381,8 @@ class RealtimeRepository:
raise ValueError(f"Invalid field: {field}")
query = sql.SQL(
"SELECT time, {} FROM realtime.node_results WHERE time >= %s AND time <= %s AND node_id = %s"
"SELECT time, {} FROM realtime.node_results WHERE time >= %s "
"AND time <= %s AND node_id = %s ORDER BY time"
).format(sql.Identifier(field))
async with conn.cursor() as cur:
@@ -346,6 +392,33 @@ class RealtimeRepository:
{"time": row["time"].isoformat(), "value": row[field]} for row in rows
]
@staticmethod
async def get_node_fields_by_ids_time_range(
conn: AsyncConnection,
start_time: datetime,
end_time: datetime,
node_ids: list[str],
field: str,
) -> dict[str, list[dict[str, Any]]]:
valid_fields = {"actual_demand", "total_head", "pressure", "quality"}
if field not in valid_fields:
raise ValueError(f"Invalid field: {field}")
result = {node_id: [] for node_id in node_ids}
if not node_ids:
return result
query = sql.SQL(
"SELECT node_id, time, {} FROM realtime.node_results "
"WHERE time BETWEEN %s AND %s AND node_id = ANY(%s) "
"ORDER BY node_id, time"
).format(sql.Identifier(field))
async with conn.cursor() as cur:
await cur.execute(query, (start_time, end_time, node_ids))
for row in await cur.fetchall():
result.setdefault(str(row["node_id"]), []).append(
{"time": row["time"].isoformat(), "value": row[field]}
)
return result
@staticmethod
async def get_nodes_field_by_time_range(
conn: AsyncConnection, start_time: datetime, end_time: datetime, field: str
@@ -355,7 +428,8 @@ class RealtimeRepository:
raise ValueError(f"Invalid field: {field}")
query = sql.SQL(
"SELECT node_id, time, {} FROM realtime.node_results WHERE time >= %s AND time <= %s"
"SELECT node_id, time, {} FROM realtime.node_results "
"WHERE time >= %s AND time <= %s ORDER BY node_id, time"
).format(sql.Identifier(field))
async with conn.cursor() as cur:
@@ -459,6 +533,19 @@ class RealtimeRepository:
# transactions (savepoints), while this outer transaction guarantees
# that a link write failure also rolls back the node replacement.
async with conn.transaction():
async with conn.cursor() as cur:
await cur.execute(
"SELECT pg_advisory_xact_lock(hashtextextended(%s::text, 1))",
(simulation_time,),
)
await cur.execute(
"DELETE FROM realtime.node_results WHERE time = %s",
(simulation_time,),
)
await cur.execute(
"DELETE FROM realtime.link_results WHERE time = %s",
(simulation_time,),
)
if node_data:
await RealtimeRepository.insert_nodes_batch(conn, node_data)
@@ -525,6 +612,19 @@ class RealtimeRepository:
# transactions (savepoints), while this outer transaction guarantees
# that a link write failure also rolls back the node replacement.
with conn.transaction():
with conn.cursor() as cur:
cur.execute(
"SELECT pg_advisory_xact_lock(hashtextextended(%s::text, 1))",
(simulation_time,),
)
cur.execute(
"DELETE FROM realtime.node_results WHERE time = %s",
(simulation_time,),
)
cur.execute(
"DELETE FROM realtime.link_results WHERE time = %s",
(simulation_time,),
)
if node_data:
RealtimeRepository.insert_nodes_batch_sync(conn, node_data)
app/infra/db/timescaledb/repositories/scada.py
+32 -3
View File
@@ -35,7 +35,8 @@ class ScadaRepository:
) -> List[dict]:
async with conn.cursor() as cur:
await cur.execute(
"SELECT * FROM scada.measurements 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 ORDER BY device_id, time",
(device_ids, start_time, end_time),
)
return await cur.fetchall()
@@ -49,7 +50,8 @@ class ScadaRepository:
) -> List[dict]:
with conn.cursor(row_factory=dict_row) as cur:
cur.execute(
"SELECT * FROM scada.measurements 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 ORDER BY device_id, time",
(device_ids, start_time, end_time),
)
return cur.fetchall()
@@ -88,7 +90,9 @@ class ScadaRepository:
raise ValueError(f"Invalid field: {field}")
query = sql.SQL(
"SELECT device_id, time, {} FROM scada.measurements 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) "
"ORDER BY device_id, time"
).format(sql.Identifier(field))
async with conn.cursor() as cur:
@@ -122,6 +126,31 @@ class ScadaRepository:
if cur.rowcount == 0:
await cur.execute(insert_query, (time, device_id, value))
@staticmethod
async def update_scada_field_batch(
conn: AsyncConnection,
rows: list[tuple[datetime, str, float | None]],
field: str,
) -> int:
"""Update existing SCADA samples in one set-based statement."""
valid_fields = {"monitored_value", "cleaned_value"}
if field not in valid_fields:
raise ValueError(f"Invalid field: {field}")
if not rows:
return 0
query = sql.SQL(
"UPDATE scada.measurements AS measurement SET {} = batch.value "
"FROM unnest(%s::timestamptz[], %s::text[], %s::double precision[]) "
"AS batch(time, device_id, value) "
"WHERE measurement.time = batch.time "
"AND measurement.device_id = batch.device_id"
).format(sql.Identifier(field))
times, device_ids, values = zip(*rows)
async with conn.cursor() as cur:
await cur.execute(query, (list(times), list(device_ids), list(values)))
return cur.rowcount
@staticmethod
async def delete_scada_by_id_time_range(
conn: AsyncConnection, device_id: str, start_time: datetime, end_time: datetime
app/infra/db/timescaledb/sync_pool.py
+2
View File
@@ -11,6 +11,7 @@ from app.core.config import settings
from app.infra.db.project_routing import get_project_timescale_pgconn_string
_check_connection = ConnectionPool.check_connection
_pools: OrderedDict[str, ConnectionPool] = OrderedDict()
_pool_conninfo: dict[str, str] = {}
_pool_borrows: dict[str, int] = {}
@@ -49,6 +50,7 @@ def get_timescale_pool(db_name: str) -> ConnectionPool:
min_size=settings.PROJECT_TS_POOL_MIN_SIZE,
max_size=settings.PROJECT_TS_POOL_MAX_SIZE,
kwargs={"row_factory": dict_row},
check=_check_connection,
open=True,
)
_pools[db_name] = pool