"""Bounded read primitives over existing project views and time-series tables.""" from datetime import datetime from typing import Any from uuid import UUID from psycopg import AsyncConnection from psycopg import sql from psycopg.rows import dict_row from fastapi import HTTPException from app.infra.db.postgresql.scada import _SCADA_VIEW_SELECT, _device from app.infra.db.timescaledb.repositories.realtime import RealtimeRepository from app.infra.db.timescaledb.repositories.analysis import AnalysisResultsRepository ELEMENT_VIEWS = { v: v + "s" for v in ("junction", "pipe", "reservoir", "tank", "pump", "valve") } MAX_QUERY_BYTES = 8 * 1024 * 1024 async def read_rows(conn, statement, params=(), *, timeout_ms=30000): async with conn.transaction(), conn.cursor(row_factory=dict_row) as cur: await cur.execute( "SELECT set_config('statement_timeout',%s,true)", (str(timeout_ms),) ) await cur.execute(statement, params) return await cur.fetchall() async def asset_rows(conn, element_type, *, element_id=None, limit=100, offset=0): view = ELEMENT_VIEWS[element_type] where = sql.SQL("WHERE v.id=%s") if element_id else sql.SQL("") params = (element_id,) if element_id else () statement = sql.SQL( "SELECT to_jsonb(v)-'geom' AS record FROM gis.{} v {} ORDER BY v.id LIMIT %s OFFSET %s" ).format(sql.Identifier(view), where) rows = await read_rows(conn, statement, params + (limit, offset)) return [row["record"] for row in rows] async def scada_devices(conn, *, device_id=None, limit=100, offset=0): where = " WHERE id=%s" if device_id else "" params = (device_id,) if device_id else () rows = await read_rows( conn, _SCADA_VIEW_SELECT + where + " ORDER BY device_id LIMIT %s OFFSET %s", params + (limit, offset), ) return [_device(row) for row in rows] async def series_rows( conn, *, source, element_type, ids, start=None, end=None, time=None, field=None, run_id=None, limit=1000, offset=0, end_exclusive=False, ): node = element_type == "junction" if source == "scada": table, id_column, fields = ( "measurements", "device_id", {"monitored_value", "cleaned_value"}, ) else: table, id_column = ( ("node_results", "node_id") if node else ("link_results", "link_id") ) fields = ( ( AnalysisResultsRepository.NODE_FIELDS if node else AnalysisResultsRepository.LINK_FIELDS ) if source == "analysis" else ( RealtimeRepository.NODE_FIELDS if node else RealtimeRepository.LINK_FIELDS ) ) if field is not None and field not in fields: raise ValueError("unsupported field for this data source and element type") columns = ( sql.SQL("{},time,{} AS value").format( sql.Identifier(id_column), sql.Identifier(field) ) if field else sql.SQL("*") ) filters = [] params = [] if ids: filters.append(sql.SQL("{}=ANY(%s)").format(sql.Identifier(id_column))) params.append(ids) if run_id: filters.append(sql.SQL("run_id=%s")) params.append(run_id) if start is not None: filters.append( sql.SQL( "time >= %s AND time < %s" if end_exclusive else "time >= %s AND time <= %s" ) ) params.extend((start, end)) else: # Existing point-query semantics use +/- one second. filters.append( sql.SQL("time BETWEEN %s-interval '1 second' AND %s+interval '1 second'") ) params.extend((time, time)) statement = sql.SQL( "SELECT {} FROM {}.{} WHERE {} ORDER BY time,{} LIMIT %s OFFSET %s" ).format( columns, sql.Identifier(source), sql.Identifier(table), sql.SQL(" AND ").join(filters), sql.Identifier(id_column), ) return await read_rows(conn, statement, tuple(params) + (limit, offset)) async def sensor_runs(conn, *, run_id=None, limit=100, offset=0): where = " AND r.run_id=%s" if run_id else "" params = (run_id,) if run_id else () return await read_rows( conn, """SELECT r.run_id,r.name,r.created_by,r.created_at,r.status, result.payload FROM analysis.runs r JOIN LATERAL (SELECT payload FROM analysis.results WHERE run_id=r.run_id AND result_type='sensor_placement' ORDER BY created_at DESC,result_id DESC LIMIT 1) result ON true WHERE r.run_type='sensor_placement'""" + where + " ORDER BY r.created_at DESC,r.run_id LIMIT %s OFFSET %s", params + (limit, offset), ) async def history_rows( conn: AsyncConnection, descriptors: list[dict[str, Any]], *, start: datetime, end: datetime, run_id: UUID | None, limit: int, offset: int, ) -> list[dict[str, Any]]: """Page points across raw, cleaned and model series in one SQL statement.""" parts, params = [], [] for index, item in enumerate(descriptors): source, field = item["source"], item["field"] if source == "scada": table, id_column, fields = "measurements", "device_id", {"monitored_value", "cleaned_value"} elif source in {"realtime", "analysis"}: node = item["element_type"] == "junction" table, id_column = ("node_results", "node_id") if node else ("link_results", "link_id") fields = {"pressure"} if node else {"flow"} else: raise ValueError("unsupported history source") if field not in fields or (source == "analysis" and run_id is None): raise ValueError("unsupported history field or missing run") parts.append(sql.SQL( "SELECT %s::integer AS series_index,time,{} AS value FROM {}.{} " "WHERE {}=%s AND time >= %s AND time <= %s{}" ).format( sql.Identifier(field), sql.Identifier(source), sql.Identifier(table), sql.Identifier(id_column), sql.SQL(" AND run_id=%s" if source == "analysis" else ""), )) params.extend((index, item["id"], start, end)) if source == "analysis": params.append(run_id) if not parts: return [] statement = sql.SQL( "SELECT series_index,time,value FROM ({}) points " "ORDER BY series_index, time LIMIT %s OFFSET %s" ).format(sql.SQL(" UNION ALL ").join(parts)) return await read_rows(conn, statement, tuple(params) + (limit, offset)) async def evidence_page( conn: AsyncConnection, *, project_id: UUID, task_id: UUID, evidence_id: str, data_path: list[str], limit: int, offset: int, mcp_sources: list[str], full_snapshot: bool = False, ) -> dict[str, Any]: """Project the immutable snapshot in PostgreSQL before transferring JSON.""" return await _snapshot_page(conn, stored_sql=""" SELECT document->'data_snapshot' AS snapshot FROM workbench.operational_task_evidence WHERE project_id=%s AND task_id=%s AND evidence_id=%s AND jsonb_typeof(document->'data_snapshot')='object' """, stored_params=(project_id, task_id, evidence_id), data_path=data_path, limit=limit, offset=offset, mcp_sources=mcp_sources, full_snapshot=full_snapshot) async def _snapshot_page( conn: AsyncConnection, *, stored_sql: str, stored_params: tuple, data_path: list[str], limit: int, offset: int, mcp_sources: list[str], full_snapshot: bool = False, max_bytes: int = MAX_QUERY_BYTES, ) -> dict[str, Any]: """Shared JSONB projection; callers supply only fixed, scoped SQL sources.""" def descriptor(value, path): return f"""jsonb_build_object( 'collection_type',jsonb_typeof({value}), 'count',CASE jsonb_typeof({value}) WHEN 'array' THEN jsonb_array_length({value}) WHEN 'string' THEN char_length({value} #>> '{{}}') ELSE (SELECT count(*) FROM jsonb_object_keys({value})) END, 'data_path',to_jsonb({path}))""" collection = descriptor("p.value", "p.path") child = descriptor("child.item", "p.path || ARRAY[child.key]") statement = """ WITH stored AS MATERIALIZED ( STORED_QUERY ), captured AS ( SELECT snapshot, snapshot->>'source'=ANY(%s::text[]) AND jsonb_typeof(snapshot->'data')='object' AS is_mcp FROM stored ), selected AS MATERIALIZED ( SELECT snapshot, CASE WHEN is_mcp THEN (snapshot->'data')-'data' END AS query_metadata, (CASE WHEN is_mcp THEN snapshot->'data'->'data' ELSE snapshot->'data' END) #> %s::text[] AS value FROM captured ), entries AS ( SELECT ordinality AS position, NULL::text AS key, item AS value FROM selected, jsonb_array_elements( CASE WHEN jsonb_typeof(value)='array' THEN value ELSE '[]'::jsonb END ) WITH ORDINALITY AS a(item, ordinality) UNION ALL SELECT row_number() OVER (ORDER BY e.key COLLATE "C"), e.key, e.item FROM selected, jsonb_each( CASE WHEN jsonb_typeof(value)='object' THEN value ELSE '{}'::jsonb END ) AS e(key,item) ), paged AS MATERIALIZED ( SELECT * FROM entries ORDER BY position LIMIT %s OFFSET %s ), paths AS ( SELECT paged.*, jsonb_typeof(selected.value) AS parent_type, %s::text[] || ARRAY[CASE WHEN jsonb_typeof(selected.value)='array' THEN (position-1)::text ELSE key END] AS path FROM paged,selected ), summaries AS ( SELECT p.position,p.key,p.value AS original, CASE WHEN jsonb_typeof(p.value) IN ('array','object') OR (jsonb_typeof(p.value)='string' AND char_length(p.value #>> '{}')>256) THEN CASE WHEN p.parent_type='object' OR jsonb_typeof(p.value) IN ('array','string') THEN COLLECTION_PREVIEW ELSE COALESCE((SELECT jsonb_object_agg(child.key, CASE WHEN jsonb_typeof(child.item) IN ('array','object') OR (jsonb_typeof(child.item)='string' AND char_length(child.item #>> '{}')>256) THEN CHILD_PREVIEW ELSE child.item END) FROM jsonb_each(p.value) AS child(key,item)), '{}'::jsonb) END ELSE p.value END AS value FROM paths p ), projected AS ( SELECT CASE WHEN %s THEN snapshot ELSE jsonb_build_object( 'snapshot_metadata',snapshot-'data', 'query_metadata',query_metadata, 'nested',EXISTS(SELECT 1 FROM summaries WHERE original<>value), 'value_type',jsonb_typeof(value), 'total',CASE jsonb_typeof(value) WHEN 'array' THEN jsonb_array_length(value) WHEN 'object' THEN (SELECT count(*) FROM jsonb_object_keys(value)) WHEN 'string' THEN char_length(value #>> '{}') END, 'value',CASE jsonb_typeof(value) WHEN 'array' THEN COALESCE((SELECT jsonb_agg(value ORDER BY position) FROM summaries),'[]'::jsonb) WHEN 'object' THEN COALESCE((SELECT jsonb_object_agg(key,value) FROM summaries),'{}'::jsonb) WHEN 'string' THEN to_jsonb(substring(value #>> '{}' FROM %s FOR %s)) ELSE value END ) END AS payload, value IS NULL AS missing_path FROM selected ), sized AS MATERIALIZED ( SELECT payload,missing_path,octet_length(payload::text)>%s AS too_large FROM projected ) SELECT CASE WHEN NOT too_large THEN payload END AS payload,too_large,missing_path FROM sized """ statement = statement.replace("STORED_QUERY", stored_sql).replace("COLLECTION_PREVIEW", collection).replace("CHILD_PREVIEW", child) rows = await read_rows(conn, statement, ( *stored_params, mcp_sources, data_path, limit, offset, data_path, full_snapshot, offset + 1, limit, max_bytes, )) if not rows: raise HTTPException(404, "Stored query document was not found") row = rows[0] if row["too_large"]: raise HTTPException(413, "Query exceeds the byte budget; select a deeper data_path or smaller page") if row.get("missing_path") and not full_snapshot: raise HTTPException(422, "Evidence data_path does not name a stored value") return row["payload"] async def task_context_summary(conn: AsyncConnection, *, project_id: UUID, task_id: UUID, report_revision: int | None) -> dict: """Read fixed-size metadata and explicit previews, never a complete report.""" report = """jsonb_build_object('revision',r.revision, 'execution_generation',r.document->'execution_generation', 'stage',r.document#>'{report,stage}','outcome',r.document#>'{report,outcome}', 'current_conclusion_preview',left(r.document#>>'{report,current_conclusion}',200), 'preview_only',true,'detail',jsonb_build_object('section','report','report_revision',r.revision))""" rows = await read_rows(conn, """ SELECT jsonb_build_object('task',jsonb_build_object( 'task_id',t.task_id,'task_type',t.task_type,'workflow',t.document->>'workflow','status',t.status, 'lifecycle_status',t.lifecycle_status,'priority',t.priority,'revision',t.revision, 'execution_generation',t.execution_generation, 'current_report_revision',t.current_report_revision, 'title_preview',left(t.document->>'title',160),'preview_only',true, 'detail',jsonb_build_object('section','task')), 'latest_report',(SELECT REPORT_SUMMARY FROM workbench.operational_task_reports r WHERE r.project_id=t.project_id AND r.task_id=t.task_id ORDER BY r.revision DESC LIMIT 1), 'selected_report',CASE WHEN %s::int IS NULL THEN NULL ELSE (SELECT REPORT_SUMMARY FROM workbench.operational_task_reports r WHERE r.project_id=t.project_id AND r.task_id=t.task_id AND r.revision=%s) END ) AS payload FROM workbench.operational_tasks t WHERE t.project_id=%s AND t.task_id=%s """.replace("REPORT_SUMMARY", report), (report_revision, report_revision, project_id, task_id)) if not rows: raise HTTPException(404, "Task was not found in this project") result = rows[0]["payload"] if report_revision is not None and result["selected_report"] is None: raise HTTPException(404, "Selected report revision was not found") return result async def task_context_page( conn: AsyncConnection, *, project_id: UUID, task_id: UUID, section: str, report_revision: int | None, data_path: list[str], limit: int, offset: int, ) -> dict: """Page only the requested task section inside PostgreSQL, before transfer.""" collections = { "evidence": ("operational_task_evidence", "evidence_id"), "actions": ("operational_task_actions", "action_id"), "access_requests": ("operational_task_access_requests", "access_request_id"), } documents = {"task": "operational_tasks", "report": "operational_task_reports", "plan": "operational_task_plans"} if section in {"access_requests", "progress"}: table = "operational_task_access_requests" if section == "access_requests" else "operational_task_progress" available = await read_rows(conn, "SELECT to_regclass(%s) AS relation", ("workbench." + table,)) if not available[0]["relation"]: return {"snapshot_metadata": {"section": section, "available": False}, "value": None, "value_type": "null", "total": None, "nested": False} projection_offset = offset if section in collections: table, key = collections[section] # These identifiers come exclusively from the fixed table above. relation = sql.Identifier("workbench", table).as_string() identifier = sql.Identifier(key).as_string() if data_path: stored = f"""SELECT jsonb_build_object('section',%s::text,'data',jsonb_build_object({identifier}::text,document)) AS snapshot FROM {relation} WHERE project_id=%s AND task_id=%s AND {identifier}::text=%s""" params = (section, project_id, task_id, data_path[0]) else: # Fetch just this page of rows. Collection keys remain stable across # paging, so every nested data_path uses the real evidence/action ID. stored = f"""WITH items AS MATERIALIZED ( SELECT {identifier}::text AS key,document FROM {relation} WHERE project_id=%s AND task_id=%s ORDER BY {identifier}::text COLLATE "C" LIMIT %s OFFSET %s) SELECT jsonb_build_object('section',%s::text, 'total',(SELECT count(*) FROM {relation} WHERE project_id=%s AND task_id=%s), 'data',COALESCE(jsonb_object_agg(key,document),'{{}}'::jsonb)) AS snapshot FROM items""" params = (project_id, task_id, limit + 1, offset, section, project_id, task_id) projection_offset = 0 elif section in documents: relation = sql.Identifier("workbench", documents[section]).as_string() where = "" params = (section, project_id, task_id) if section == "report": where = " AND (%s::int IS NULL OR revision=%s) ORDER BY revision DESC LIMIT 1" params += (report_revision, report_revision) elif section == "plan": where = " ORDER BY revision DESC LIMIT 1" stored = f"SELECT jsonb_build_object('section',%s::text,'data',document) AS snapshot FROM {relation} WHERE project_id=%s AND task_id=%s{where}" elif section == "run": stored = """SELECT jsonb_build_object('section','run','data',to_jsonb(r)-'project_id') AS snapshot FROM workbench.operational_task_runs r WHERE project_id=%s AND task_id=%s ORDER BY generation DESC LIMIT 1""" params = (project_id, task_id) elif section == "progress": stored = """SELECT jsonb_build_object('section','progress','data',p.document) AS snapshot FROM workbench.operational_task_progress p JOIN workbench.operational_task_runs r ON r.project_id=p.project_id AND r.task_id=p.task_id AND r.run_id=p.run_id JOIN workbench.operational_tasks t ON t.project_id=p.project_id AND t.task_id=p.task_id AND t.execution_generation=r.generation WHERE p.project_id=%s AND p.task_id=%s""" params = (project_id, task_id) else: raise HTTPException(422, "Unknown task context section") return await _snapshot_page(conn, stored_sql=stored, stored_params=params, data_path=data_path, limit=limit, offset=projection_offset, mcp_sources=[], max_bytes=8000) async def task_data_scopes(conn: AsyncConnection, *, project_id: UUID, task_id: UUID, permissions: set[str]) -> list[str]: """Bound scope discovery to seven known sources, without loading task detail.""" known = ["workbench_runtime", "peripheral_points", "peripheral_work_orders", "peripheral_observations", "peripheral_test_data", "scada_devices", "scada_readings"] scopes = {name for name in known if ("scada.view" if name.startswith("scada_") else "simulation.view") in permissions} available = await read_rows(conn, "SELECT to_regclass('workbench.operational_task_access_requests') AS relation") if available[0]["relation"]: rows = await read_rows(conn, """SELECT DISTINCT scope FROM workbench.operational_task_access_requests, LATERAL jsonb_array_elements_text(COALESCE(document->'approved_scopes','[]'::jsonb)) AS granted(scope) WHERE project_id=%s AND task_id=%s AND status='approved' AND scope=ANY(%s::text[])""", (project_id, task_id, known)) scopes.update(row["scope"] for row in rows) return [name for name in known if name in scopes]