jiang
be6317eb48
Add project-bound OAuth MCP queries and controlled analysis with scoped evidence, limits, timeouts and a nearby-model spatial query using EPSG:4326 inputs. Default nearby searches to 200 meters and five results. Integrate investigation messaging/presentation, health/readiness and typed history; remove duplicate web search and fix online simulation and compressed baseline writes. Include project migrations 023-025; keep local credentials and demo seeding excluded. Validation: 319 focused pytest checks; real PostGIS temporary-table regression; read-only Jinqian bridge radius comparison; cross-language MCP discovery/calls.
407 lines
20 KiB
Python
407 lines
20 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 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]
|