jiang
5966d039de
Reorganize algorithm packages by business responsibility, move orchestration into services, and keep database access behind pooled repositories. Harden analysis API validation, remove unsafe legacy simulation endpoints, and add regression and architecture boundary coverage. BREAKING CHANGE: legacy algorithm module paths and obsolete simulation endpoints are removed.
316 lines
12 KiB
Python
316 lines
12 KiB
Python
from datetime import datetime, timedelta
|
|
from typing import Any
|
|
from uuid import UUID
|
|
|
|
from psycopg import AsyncConnection, Connection, sql
|
|
|
|
from app.domain.time 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_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,
|
|
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"),
|
|
)
|
|
)
|