From 4c441fcbf2c4ac662677bdbf50e93a2b2dab38db Mon Sep 17 00:00:00 2001 From: Huarch Date: Sun, 15 Nov 2026 23:49:03 +0800 Subject: [PATCH] refactor(db): isolate hydraulic playback data access --- app/infra/db/postgresql/hydraulic_playback.py | 24 ++++++++ .../repositories/hydraulic_playback.py | 29 ++++++++++ app/services/hydraulic_playback.py | 55 ++++++------------- 3 files changed, 71 insertions(+), 37 deletions(-) create mode 100644 app/infra/db/postgresql/hydraulic_playback.py create mode 100644 app/infra/db/timescaledb/repositories/hydraulic_playback.py diff --git a/app/infra/db/postgresql/hydraulic_playback.py b/app/infra/db/postgresql/hydraulic_playback.py new file mode 100644 index 0000000..9f08480 --- /dev/null +++ b/app/infra/db/postgresql/hydraulic_playback.py @@ -0,0 +1,24 @@ +"""Published network geometry used to build indexed hydraulic playback packets.""" + +from typing import Any + +from psycopg import AsyncConnection + + +class HydraulicGeometryRepository: + @staticmethod + async def get_rows(conn: AsyncConnection) -> tuple[list[dict[str, Any]], list[dict[str, Any]]]: + async with conn.cursor() as cursor: + await cursor.execute(""" + SELECT n.id, ST_X(ST_Transform(g.geom,4326)) AS longitude, + ST_Y(ST_Transform(g.geom,4326)) AS latitude + FROM gis.network_nodes n JOIN gis.node_geometries g ON g.node_id=n.id + WHERE g.geom IS NOT NULL AND NOT ST_IsEmpty(g.geom) ORDER BY n.id + """) + nodes = await cursor.fetchall() + await cursor.execute(""" + SELECT id, start_node_id, diameter, ST_AsGeoJSON(ST_Transform(ST_LineMerge(geom),4326))::json AS geometry + FROM gis.pipes WHERE geom IS NOT NULL AND NOT ST_IsEmpty(geom) ORDER BY id + """) + pipes = await cursor.fetchall() + return nodes, pipes diff --git a/app/infra/db/timescaledb/repositories/hydraulic_playback.py b/app/infra/db/timescaledb/repositories/hydraulic_playback.py new file mode 100644 index 0000000..96400a7 --- /dev/null +++ b/app/infra/db/timescaledb/repositories/hydraulic_playback.py @@ -0,0 +1,29 @@ +"""Stream exact playback samples without materializing full-network JSON results.""" + +from collections.abc import AsyncIterator +from datetime import datetime +from typing import Any + +from psycopg import AsyncConnection + + +class HydraulicPlaybackRepository: + @staticmethod + async def iter_node_pressures(conn: AsyncConnection, times: list[datetime]) -> AsyncIterator[dict[str, Any]]: + async with conn.cursor() as cursor: + await cursor.execute(""" + SELECT time, node_id, pressure FROM realtime.node_results + WHERE time = ANY(%s) ORDER BY time, node_id + """, (times,)) + async for row in cursor: + yield row + + @staticmethod + async def iter_pipe_metrics(conn: AsyncConnection, times: list[datetime]) -> AsyncIterator[dict[str, Any]]: + async with conn.cursor() as cursor: + await cursor.execute(""" + SELECT time, link_id, flow, velocity FROM realtime.link_results + WHERE time = ANY(%s) ORDER BY time, link_id + """, (times,)) + async for row in cursor: + yield row diff --git a/app/services/hydraulic_playback.py b/app/services/hydraulic_playback.py index 4135f3a..68f2b94 100644 --- a/app/services/hydraulic_playback.py +++ b/app/services/hydraulic_playback.py @@ -14,6 +14,8 @@ from typing import Any from psycopg import AsyncConnection from app.infra.db.postgresql.network_settings import NetworkSettingsRepository +from app.infra.db.postgresql.hydraulic_playback import HydraulicGeometryRepository +from app.infra.db.timescaledb.repositories.hydraulic_playback import HydraulicPlaybackRepository WORLD_METERS = 40075016.68557849 MAX_FRAME_COUNT = 8 @@ -122,19 +124,7 @@ def build_geometry(nodes: list[dict], pipes: list[dict], units: dict[str, str]) async def read_geometry(conn: AsyncConnection) -> PlaybackGeometry: - async with conn.cursor() as cursor: - await cursor.execute(""" - SELECT n.id, ST_X(ST_Transform(g.geom,4326)) AS longitude, - ST_Y(ST_Transform(g.geom,4326)) AS latitude - FROM gis.network_nodes n JOIN gis.node_geometries g ON g.node_id=n.id - WHERE g.geom IS NOT NULL AND NOT ST_IsEmpty(g.geom) ORDER BY n.id - """) - nodes = await cursor.fetchall() - await cursor.execute(""" - SELECT id, start_node_id, diameter, ST_AsGeoJSON(ST_Transform(ST_LineMerge(geom),4326))::json AS geometry - FROM gis.pipes WHERE geom IS NOT NULL AND NOT ST_IsEmpty(geom) ORDER BY id - """) - pipes = await cursor.fetchall() + nodes, pipes = await HydraulicGeometryRepository.get_rows(conn) units = await NetworkSettingsRepository.get_result_units(conn) return await asyncio.to_thread(build_geometry, nodes, pipes, units) @@ -154,30 +144,21 @@ async def read_frames(conn: AsyncConnection, geometry: PlaybackGeometry, start: values = array("f", [math.nan]) * (stride * frame_count) time_index = {value: i for i, value in enumerate(times)} valid_counts = {metric: [0] * frame_count for metric in ("pressure", "flow", "velocity")} - async with conn.cursor() as cursor: - await cursor.execute(""" - SELECT time, node_id, pressure FROM realtime.node_results - WHERE time = ANY(%s) ORDER BY time, node_id - """, (times,)) - async for row in cursor: - index = geometry.node_index.get(str(row["node_id"])) - if index is not None: - frame = time_index[row["time"]] - values[frame * stride + index] = finite(row["pressure"], geometry.pressure_factor) - valid_counts["pressure"][frame] += math.isfinite(values[frame * stride + index]) - await cursor.execute(""" - SELECT time, link_id, flow, velocity FROM realtime.link_results - WHERE time = ANY(%s) ORDER BY time, link_id - """, (times,)) - async for row in cursor: - index = geometry.pipe_index.get(str(row["link_id"])) - if index is not None: - frame = time_index[row["time"]] - offset = frame * stride + node_count - values[offset + index] = finite(row["flow"], geometry.flow_factor) - values[offset + pipe_count + index] = finite(row["velocity"], geometry.velocity_factor) - valid_counts["flow"][frame] += math.isfinite(values[offset + index]) - valid_counts["velocity"][frame] += math.isfinite(values[offset + pipe_count + index]) + async for row in HydraulicPlaybackRepository.iter_node_pressures(conn, times): + index = geometry.node_index.get(str(row["node_id"])) + if index is not None: + frame = time_index[row["time"]] + values[frame * stride + index] = finite(row["pressure"], geometry.pressure_factor) + valid_counts["pressure"][frame] += math.isfinite(values[frame * stride + index]) + async for row in HydraulicPlaybackRepository.iter_pipe_metrics(conn, times): + index = geometry.pipe_index.get(str(row["link_id"])) + if index is not None: + frame = time_index[row["time"]] + offset = frame * stride + node_count + values[offset + index] = finite(row["flow"], geometry.flow_factor) + values[offset + pipe_count + index] = finite(row["velocity"], geometry.velocity_factor) + valid_counts["flow"][frame] += math.isfinite(values[offset + index]) + valid_counts["velocity"][frame] += math.isfinite(values[offset + pipe_count + index]) return packet({"kind": "frames", "protocol": 1, "modelVersion": geometry.version, "nodeCount": node_count, "pipeCount": pipe_count, "times": [int(t.timestamp() * 1000) for t in times], "validCounts": valid_counts,