Accumulate investigation paths and preserve failed attempts across rounds. Keep graph annotations separate from execution state, synchronize the API contract, and consolidate documentation.
439 lines
22 KiB
Python
439 lines
22 KiB
Python
"""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 psycopg.types.json import Jsonb
|
|
from fastapi import HTTPException
|
|
from app.domain.schemas.operational_task import LEGACY_INVESTIGATION_NODE_ANNOTATIONS
|
|
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')),
|
|
'handling_context',COALESCE(t.document->'handling_context','{}'::jsonb),
|
|
'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
|
|
|
|
|
|
def _reference_diagram_projection(stored_sql: str, stored_params: tuple, section: str) -> tuple[str, tuple]:
|
|
"""Normalize only saved diagram annotations before any data_path or page projection."""
|
|
if section not in {"progress", "report"}:
|
|
raise ValueError("reference diagrams only belong to progress or report sections")
|
|
path = "{data," + section + ",diagram,node_refs}"
|
|
statement = f"""
|
|
WITH reference_documents AS (
|
|
{stored_sql}
|
|
), reference_aliases AS (
|
|
SELECT %s::jsonb AS annotations
|
|
)
|
|
SELECT CASE WHEN jsonb_typeof(snapshot#>'{path}')='array' THEN
|
|
jsonb_set(snapshot,'{path}',COALESCE((
|
|
SELECT jsonb_agg(CASE WHEN annotations ? (node->>'status')
|
|
THEN jsonb_set(node,'{{status}}',annotations->(node->>'status'))
|
|
ELSE node END ORDER BY position)
|
|
FROM jsonb_array_elements(snapshot#>'{path}') WITH ORDINALITY AS nodes(node,position)
|
|
),'[]'::jsonb)) ELSE snapshot END AS snapshot
|
|
FROM reference_documents CROSS JOIN reference_aliases
|
|
"""
|
|
return statement, (*stored_params, Jsonb(LEGACY_INVESTIGATION_NODE_ANNOTATIONS))
|
|
|
|
|
|
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"),
|
|
"questions": ("operational_task_questions", "question_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
|
|
WHERE p.project_id=%s AND p.task_id=%s AND r.generation<=t.execution_generation
|
|
AND (t.execution_generation=r.generation
|
|
OR jsonb_typeof(p.document#>'{progress,diagram,edges}')='array')
|
|
ORDER BY (t.execution_generation=r.generation) DESC,r.generation DESC,p.revision DESC LIMIT 1"""
|
|
params = (project_id, task_id)
|
|
else:
|
|
raise HTTPException(422, "Unknown task context section")
|
|
if section in {"progress", "report"}:
|
|
stored, params = _reference_diagram_projection(stored, params, 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]
|