refactor(db): isolate hydraulic playback data access

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

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

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:
+    nodes, pipes = await HydraulicGeometryRepository.get_rows(conn)
-        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()
     units = await NetworkSettingsRepository.get_result_units(conn)
     return await asyncio.to_thread(build_geometry, nodes, pipes, units)
 
@@ -154,22 +144,13 @@ 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:
+    async for row in HydraulicPlaybackRepository.iter_node_pressures(conn, times):
-        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("""
+    async for row in HydraulicPlaybackRepository.iter_pipe_metrics(conn, times):
-            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"]]