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.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_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"),
|
|
)
|
|
)
|