feat(server): integrate MCP queries and investigation messaging
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.
This commit is contained in:
@@ -0,0 +1,216 @@
|
||||
"""Opt-in SQL checks against a disposable PostgreSQL database; no business data."""
|
||||
|
||||
import asyncio
|
||||
import os
|
||||
from datetime import datetime, timedelta, timezone
|
||||
from types import SimpleNamespace
|
||||
from uuid import uuid4
|
||||
|
||||
import pytest
|
||||
from fastapi import HTTPException
|
||||
from psycopg import AsyncConnection
|
||||
from psycopg.rows import dict_row
|
||||
from psycopg.types.json import Jsonb
|
||||
|
||||
from app.infra.db.postgresql import read_queries
|
||||
from app.mcp.models import PeripheralInput
|
||||
from app.mcp.peripheral import read_peripheral_query
|
||||
from app.services.content_digest import digest
|
||||
|
||||
|
||||
def test_history_and_evidence_pages_with_real_postgresql():
|
||||
dsn = os.environ.get("MCP_QUERY_TEST_DSN")
|
||||
if not dsn:
|
||||
pytest.skip("MCP_QUERY_TEST_DSN must point to a disposable PostgreSQL database")
|
||||
|
||||
async def check():
|
||||
conn = await AsyncConnection.connect(dsn, row_factory=dict_row)
|
||||
async with conn, conn.transaction(force_rollback=True):
|
||||
await conn.execute("CREATE SCHEMA workbench")
|
||||
await conn.execute("CREATE TABLE workbench.operational_task_evidence (project_id uuid, task_id uuid, evidence_id text, document jsonb)")
|
||||
await conn.execute("CREATE SCHEMA scada")
|
||||
await conn.execute("CREATE TABLE scada.measurements (device_id text, time timestamptz, monitored_value double precision, cleaned_value double precision)")
|
||||
await conn.execute("CREATE SCHEMA analysis")
|
||||
await conn.execute("CREATE TABLE analysis.node_results (run_id uuid, node_id text, time timestamptz, pressure double precision)")
|
||||
await conn.execute("CREATE SCHEMA peripheral_sample")
|
||||
await conn.execute("CREATE TABLE peripheral_sample.demo_datasets (project_id uuid, dataset_id text, document jsonb, digest text)")
|
||||
project, task, run = uuid4(), uuid4(), uuid4()
|
||||
start = datetime(2026, 10, 1, tzinfo=timezone.utc)
|
||||
for i in range(2):
|
||||
await conn.execute("INSERT INTO scada.measurements VALUES (%s,%s,%s,%s)", ("D1", start+timedelta(minutes=i), i, i+10))
|
||||
await conn.execute("INSERT INTO analysis.node_results VALUES (%s,%s,%s,%s)", (run, "J1", start+timedelta(minutes=i), i+20))
|
||||
await conn.execute("INSERT INTO analysis.node_results VALUES (%s,%s,%s,%s)", (uuid4(), "J1", start+timedelta(minutes=i), 999))
|
||||
descriptors = [
|
||||
{"source": "scada", "element_type": "junction", "id": "D1", "field": field}
|
||||
for field in ("monitored_value", "cleaned_value")
|
||||
] + [{"source": "analysis", "element_type": "junction", "id": "J1", "field": "pressure"}]
|
||||
seen = []
|
||||
for offset in (0, 2, 4, 6):
|
||||
rows = await read_queries.history_rows(conn, descriptors, start=start, end=start+timedelta(hours=1), run_id=run, limit=2, offset=offset)
|
||||
seen.extend(row["value"] for row in rows)
|
||||
assert seen == [0, 1, 10, 11, 20, 21]
|
||||
|
||||
snapshot = {
|
||||
"source": "query_scada_readings", "source_ref": "saved:1", "queried_at": start.isoformat(),
|
||||
"origin": "synthetic", "truncated": False,
|
||||
"data": {"semantics": "observed", "data": {
|
||||
"items": [{"value": 1}, {"value": 2, "points": [3, 4]}],
|
||||
"text": "甲😀乙丙", "none": None,
|
||||
}},
|
||||
}
|
||||
await conn.execute("INSERT INTO workbench.operational_task_evidence VALUES (%s,%s,%s,%s)", (project, task, "E1", Jsonb({"data_snapshot": snapshot})))
|
||||
|
||||
async def evidence(path, **overrides):
|
||||
values = dict(project_id=project, task_id=task, evidence_id="E1", data_path=path, limit=1, offset=0, mcp_sources=["query_scada_readings"])
|
||||
return await read_queries.evidence_page(conn, **(values | overrides))
|
||||
|
||||
page = await evidence(["items"], offset=1)
|
||||
assert page["total"] == 2 and page["nested"]
|
||||
assert page["snapshot_metadata"]["origin"] == "synthetic"
|
||||
assert page["query_metadata"] == {"semantics": "observed"}
|
||||
assert page["value"] == [{"value": 2, "points": {"collection_type": "array", "count": 2, "data_path": ["items", "1", "points"]}}]
|
||||
assert (await evidence(["items", "1", "points"], offset=1))["value"] == [4]
|
||||
assert (await evidence(["text"], offset=1, limit=2))["value"] == "😀乙"
|
||||
assert (await evidence(["none"]))["value"] is None
|
||||
assert (await evidence([], full_snapshot=True)) == snapshot
|
||||
assert (await evidence([]))["value"]["items"]["data_path"] == ["items"]
|
||||
for overrides, status in (({"project_id": uuid4()}, 404), ({"task_id": uuid4()}, 404), ({"data_path": ["absent"]}, 422)):
|
||||
with pytest.raises(HTTPException) as caught:
|
||||
await evidence([], **overrides)
|
||||
assert caught.value.status_code == status
|
||||
|
||||
legacy = {**snapshot, "source": "scada_readings", "data": {"items": [5, 6]}}
|
||||
await conn.execute("INSERT INTO workbench.operational_task_evidence VALUES (%s,%s,%s,%s)", (project, task, "legacy", Jsonb({"data_snapshot": legacy})))
|
||||
old_page = await evidence(["items"], evidence_id="legacy", offset=1)
|
||||
assert old_page["value"] == [6] and old_page["query_metadata"] is None
|
||||
|
||||
# A large nested collection is represented by a path, not transferred in full.
|
||||
snapshot["data"]["data"]["huge"] = ["x" * (read_queries.MAX_QUERY_BYTES + 1)]
|
||||
await conn.execute("UPDATE workbench.operational_task_evidence SET document=%s", (Jsonb({"data_snapshot": snapshot}),))
|
||||
assert (await evidence(["huge"], offset=1))["value"] == []
|
||||
assert (await evidence([], limit=1))["value"]["huge"]["count"] == 1
|
||||
descriptor = (await evidence(["huge"]))["value"][0]
|
||||
assert descriptor["collection_type"] == "string" and descriptor["data_path"] == ["huge", "0"]
|
||||
assert (await evidence(["huge", "0"], offset=7, limit=3))["value"] == "xxx"
|
||||
with pytest.raises(HTTPException) as caught:
|
||||
await evidence([], full_snapshot=True)
|
||||
assert caught.value.status_code == 413
|
||||
|
||||
payload = {"project_id": str(project), "points": [{"point_id": "DEMO-1"}], "observations": [
|
||||
{"point_id": "DEMO-1", "observed_at": (start+timedelta(minutes=i)).isoformat(), "value": i} for i in range(3)
|
||||
]}
|
||||
await conn.execute("INSERT INTO peripheral_sample.demo_datasets VALUES (%s,'water-demo-v1',%s,%s)", (project, Jsonb(payload), digest(payload)))
|
||||
args = PeripheralInput(context_handle="a"*43, source="peripheral_observations", point_id="DEMO-1", start=start, end=start+timedelta(hours=1), limit=1, offset=1)
|
||||
scope = SimpleNamespace(connection=conn, project=SimpleNamespace(project_id=project))
|
||||
demo = await read_peripheral_query(scope, args)
|
||||
assert demo.origin == "synthetic" and demo.truncated and demo.data["items"][0]["value"] == 1
|
||||
await conn.execute("UPDATE peripheral_sample.demo_datasets SET digest='corrupted'")
|
||||
with pytest.raises(HTTPException) as caught:
|
||||
await read_peripheral_query(scope, args)
|
||||
assert caught.value.status_code == 503
|
||||
payload["padding"] = "x" * (read_queries.MAX_QUERY_BYTES + 1)
|
||||
await conn.execute("UPDATE peripheral_sample.demo_datasets SET document=%s", (Jsonb(payload),))
|
||||
with pytest.raises(HTTPException) as caught:
|
||||
await read_peripheral_query(scope, args)
|
||||
assert caught.value.status_code == 413
|
||||
|
||||
asyncio.run(check())
|
||||
|
||||
|
||||
def test_task_context_summary_and_sections_with_real_postgresql():
|
||||
dsn = os.environ.get("MCP_QUERY_TEST_DSN")
|
||||
if not dsn:
|
||||
pytest.skip("MCP_QUERY_TEST_DSN must point to a disposable PostgreSQL database")
|
||||
|
||||
async def check():
|
||||
conn = await AsyncConnection.connect(dsn, row_factory=dict_row)
|
||||
async with conn, conn.transaction(force_rollback=True):
|
||||
await conn.execute("CREATE SCHEMA workbench")
|
||||
await conn.execute("""CREATE TABLE workbench.operational_tasks (
|
||||
project_id uuid,task_id uuid,task_type text,status text,lifecycle_status text,priority text,
|
||||
revision bigint,execution_generation bigint,current_report_revision bigint,document jsonb)""")
|
||||
for name, key in (("reports", "revision bigint"), ("plans", "revision bigint"),
|
||||
("evidence", "evidence_id text"), ("actions", "action_id uuid")):
|
||||
await conn.execute(f"CREATE TABLE workbench.operational_task_{name} (project_id uuid,task_id uuid,{key},document jsonb)")
|
||||
await conn.execute("CREATE TABLE workbench.operational_task_runs (project_id uuid,task_id uuid,run_id uuid,generation bigint,status text)")
|
||||
project, task, run = uuid4(), uuid4(), uuid4()
|
||||
long_text = "甲😀" * 10000
|
||||
await conn.execute("INSERT INTO workbench.operational_tasks VALUES (%s,%s,'event_investigation','running','active','normal',1,2,4,%s)",
|
||||
(project, task, Jsonb({"title": long_text, "source": {"description": long_text}})))
|
||||
for revision in (3, 4):
|
||||
await conn.execute("INSERT INTO workbench.operational_task_reports VALUES (%s,%s,%s,%s)",
|
||||
(project, task, revision, Jsonb({"revision": revision, "report": {"stage": "investigation", "outcome": "pending", "current_conclusion": long_text}})))
|
||||
await conn.execute("INSERT INTO workbench.operational_task_plans VALUES (%s,%s,1,%s)",
|
||||
(project, task, Jsonb({"nodes": [{"id": "N1"}]})))
|
||||
for evidence_id in ("E1", "E2", "E3"):
|
||||
await conn.execute("INSERT INTO workbench.operational_task_evidence VALUES (%s,%s,%s,%s)",
|
||||
(project, task, evidence_id, Jsonb({"data_snapshot": {"text": long_text}, "none": None})))
|
||||
action_id = uuid4()
|
||||
await conn.execute("INSERT INTO workbench.operational_task_actions VALUES (%s,%s,%s,%s)",
|
||||
(project, task, action_id, Jsonb({"reason": long_text})))
|
||||
await conn.execute("INSERT INTO workbench.operational_task_runs VALUES (%s,%s,%s,2,'running')", (project, task, run))
|
||||
|
||||
summary = await read_queries.task_context_summary(conn, project_id=project, task_id=task, report_revision=3)
|
||||
assert summary["selected_report"]["revision"] == 3 and summary["latest_report"]["revision"] == 4
|
||||
assert summary["task"]["preview_only"] and len(summary["task"]["title_preview"]) == 160
|
||||
assert len(summary["selected_report"]["current_conclusion_preview"]) == 200
|
||||
assert "source" not in summary["task"]
|
||||
for overrides in ({"report_revision": 99}, {"project_id": uuid4()}, {"task_id": uuid4()}):
|
||||
with pytest.raises(HTTPException) as caught:
|
||||
await read_queries.task_context_summary(conn, **(dict(project_id=project, task_id=task, report_revision=3) | overrides))
|
||||
assert caught.value.status_code == 404
|
||||
|
||||
async def page(section, path=None, **overrides):
|
||||
values = dict(project_id=project, task_id=task, section=section, report_revision=None, data_path=path or [], limit=1, offset=0)
|
||||
return await read_queries.task_context_page(conn, **(values | overrides))
|
||||
|
||||
second = await page("evidence", offset=1)
|
||||
assert second["snapshot_metadata"]["total"] == 3
|
||||
assert list(second["value"]) == ["E2"]
|
||||
assert second["value"]["E2"]["data_path"] == ["E2"]
|
||||
nested = await page("evidence", ["E2"], limit=20)
|
||||
assert nested["value"]["data_snapshot"]["data_path"] == ["E2", "data_snapshot"]
|
||||
assert nested["value"]["none"] is None
|
||||
text = await page("evidence", ["E2", "data_snapshot", "text"], offset=1, limit=2)
|
||||
assert text["value"] == "😀甲" and text["total"] == len(long_text)
|
||||
assert (await page("evidence", offset=3))["value"] == {}
|
||||
assert (await page("evidence", project_id=uuid4()))["value"] == {}
|
||||
with pytest.raises(HTTPException) as caught:
|
||||
await page("evidence", ["E2"], task_id=uuid4())
|
||||
assert caught.value.status_code == 404
|
||||
|
||||
report = await page("report", ["report", "current_conclusion"], report_revision=3, offset=2, limit=2)
|
||||
assert report["value"] == "甲😀"
|
||||
assert (await page("task"))["nested"]
|
||||
assert (await page("plan"))["value"]["nodes"]["collection_type"] == "array"
|
||||
assert (await page("actions"))["value"][str(action_id)]["data_path"] == [str(action_id)]
|
||||
assert "project_id" not in (await page("run", limit=20))["value"]
|
||||
|
||||
for section in ("progress", "access_requests"):
|
||||
assert (await page(section))["snapshot_metadata"]["available"] is False
|
||||
assert await read_queries.task_data_scopes(conn, project_id=project, task_id=task, permissions={"scada.view"}) == ["scada_devices", "scada_readings"]
|
||||
await conn.execute("CREATE TABLE workbench.operational_task_progress (project_id uuid,task_id uuid,run_id uuid,document jsonb)")
|
||||
await conn.execute("INSERT INTO workbench.operational_task_progress VALUES (%s,%s,%s,%s)", (project, task, run, Jsonb({"content": long_text})))
|
||||
assert (await page("progress"))["value"]["content"]["collection_type"] == "string"
|
||||
await conn.execute("CREATE TABLE workbench.operational_task_access_requests (project_id uuid,task_id uuid,access_request_id uuid,status text,document jsonb)")
|
||||
for grant_project, grant_task, status, scopes in (
|
||||
(project, task, "approved", ["peripheral_points", "unknown_source"]),
|
||||
(project, task, "pending", ["scada_readings"]),
|
||||
(uuid4(), task, "approved", ["peripheral_observations"]),
|
||||
(project, uuid4(), "approved", ["peripheral_test_data"]),
|
||||
):
|
||||
await conn.execute("INSERT INTO workbench.operational_task_access_requests VALUES (%s,%s,%s,%s,%s)",
|
||||
(grant_project, grant_task, uuid4(), status, Jsonb({"approved_scopes": scopes})))
|
||||
assert await read_queries.task_data_scopes(conn, project_id=project, task_id=task, permissions=set()) == ["peripheral_points"]
|
||||
assert (await page("access_requests", limit=20))["snapshot_metadata"]["total"] == 2
|
||||
|
||||
# A very wide document must fail before it crosses the DB boundary;
|
||||
# a smaller page remains readable.
|
||||
wide = {str(i): "x" * 250 for i in range(200)}
|
||||
await conn.execute("UPDATE workbench.operational_tasks SET document=%s", (Jsonb(wide),))
|
||||
with pytest.raises(HTTPException) as caught:
|
||||
await page("task", limit=200)
|
||||
assert caught.value.status_code == 413
|
||||
assert len((await page("task", limit=1))["value"]) == 1
|
||||
|
||||
asyncio.run(check())
|
||||
@@ -0,0 +1,58 @@
|
||||
"""Opt-in PostGIS checks using transaction-local temporary model tables only."""
|
||||
|
||||
import asyncio
|
||||
import os
|
||||
|
||||
import pytest
|
||||
from psycopg import AsyncConnection, sql
|
||||
from psycopg.rows import dict_row
|
||||
|
||||
from app.infra.db.postgresql import spatial_queries
|
||||
|
||||
|
||||
def test_nearby_model_queries_with_real_postgis(monkeypatch):
|
||||
dsn = os.environ.get("MCP_SPATIAL_TEST_DSN")
|
||||
if not dsn:
|
||||
pytest.skip("MCP_SPATIAL_TEST_DSN must point to a PostGIS database; only TEMP tables are written")
|
||||
monkeypatch.setitem(spatial_queries._MODEL_RELATIONS, "pipe", sql.Identifier("pg_temp", "mcp_model_pipes"))
|
||||
monkeypatch.setitem(spatial_queries._MODEL_RELATIONS, "junction", sql.Identifier("pg_temp", "mcp_model_junctions"))
|
||||
|
||||
async def check():
|
||||
conn = await AsyncConnection.connect(dsn, row_factory=dict_row)
|
||||
async with conn, conn.transaction(force_rollback=True):
|
||||
await conn.execute("CREATE TEMP TABLE mcp_model_pipes (id text PRIMARY KEY, diameter float, geom geometry(LineString,3857)) ON COMMIT DROP")
|
||||
await conn.execute("CREATE TEMP TABLE mcp_model_junctions (id text PRIMARY KEY, geom geometry(Point,3857)) ON COMMIT DROP")
|
||||
await conn.execute("CREATE INDEX ON mcp_model_pipes USING gist (geom)")
|
||||
# The crossing pipe has distant endpoints, while its middle is at the center.
|
||||
for pipe_id, wkt in [
|
||||
("crossing", "LINESTRING(121.49 30.9,121.51 30.9)"),
|
||||
("near-a", "LINESTRING(121.4999 30.9005,121.5001 30.9005)"),
|
||||
("near-b", "LINESTRING(121.4999 30.9005,121.5001 30.9005)"),
|
||||
("outside", "LINESTRING(121.4999 30.91,121.5001 30.91)"),
|
||||
]:
|
||||
await conn.execute("INSERT INTO mcp_model_pipes VALUES (%s,300,ST_Transform(ST_GeomFromText(%s,4326),3857))", (pipe_id, wkt))
|
||||
await conn.execute("INSERT INTO mcp_model_pipes VALUES ('missing',300,NULL),('empty',300,ST_GeomFromText('LINESTRING EMPTY',3857))")
|
||||
await conn.execute("INSERT INTO mcp_model_junctions VALUES ('junction',ST_Transform(ST_SetSRID(ST_MakePoint(121.5,30.9002),4326),3857))")
|
||||
|
||||
async def nearby(kind="pipe", **overrides):
|
||||
values = dict(center=(121.5, 30.9), radius_m=100, limit=10, offset=0)
|
||||
return await spatial_queries.nearby_elements(conn, kind, **(values | overrides))
|
||||
|
||||
rows = await nearby()
|
||||
assert [row["element_id"] for row in rows] == ["crossing", "near-a", "near-b"]
|
||||
assert rows[0]["distance_m"] < 1
|
||||
assert 54 < rows[1]["distance_m"] < 57
|
||||
assert rows[1]["properties"]["diameter"] == 300 and "geom" not in rows[1]["properties"]
|
||||
# 55 m geodesic is ~65 m in Web Mercator here: a 60 m radius must include it.
|
||||
assert len(await nearby(radius_m=60)) == 3
|
||||
assert len(await nearby(radius_m=50)) == 1
|
||||
assert [row["element_id"] for row in await nearby(limit=1, offset=2)] == ["near-b"]
|
||||
assert await nearby(center=(0, 0)) == []
|
||||
assert [row["element_id"] for row in await nearby("junction")] == ["junction"]
|
||||
assert await nearby(center=(0, 90)) == []
|
||||
|
||||
# A point west of the dateline must find a model east of it.
|
||||
await conn.execute("INSERT INTO mcp_model_pipes VALUES ('dateline',300,ST_Transform(ST_GeomFromText('LINESTRING(179.9999 30,179.9999 30.0001)',4326),3857))")
|
||||
assert [row["element_id"] for row in await nearby(center=(-179.9999, 30), radius_m=100)] == ["dateline"]
|
||||
|
||||
asyncio.run(check())
|
||||
@@ -66,6 +66,8 @@ async def test_postgres_task_artifacts_approval_and_project_isolation_roll_back(
|
||||
).fetchall()
|
||||
assert len(comment_rows) >= 17
|
||||
assert all(row["description"] for row in comment_rows)
|
||||
messages_migration = (Path(__file__).parents[2] / "resources/sql/024_agent_messages.sql").read_text()
|
||||
await conn.execute(messages_migration.replace("\nBEGIN;", "").replace("\nCOMMIT;", ""))
|
||||
repository = OperationalTaskRepository(conn, context.project_id)
|
||||
service = OperationalTaskService(repository, context, permissions_for_context(context))
|
||||
now = datetime.now(timezone.utc)
|
||||
|
||||
@@ -0,0 +1,101 @@
|
||||
"""Opt-in TimescaleDB regression; only accepts an independently disposable DB."""
|
||||
|
||||
import asyncio
|
||||
import os
|
||||
|
||||
import psycopg
|
||||
import pytest
|
||||
from psycopg.conninfo import conninfo_to_dict
|
||||
|
||||
from app.core.config import settings
|
||||
from app.infra.db.timescaledb.repositories.realtime import RealtimeRepository
|
||||
|
||||
|
||||
DSN = os.environ.get("REALTIME_COMPRESSION_TEST_DSN", "")
|
||||
pytestmark = pytest.mark.skipif(not DSN, reason="Needs an isolated TimescaleDB test database")
|
||||
START = "2026-10-08T04:00:00Z"
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def compressed_snapshot():
|
||||
assert conninfo_to_dict(DSN).get("dbname", "").startswith("tjw_test_realtime_"), "Refusing a non-test database"
|
||||
with psycopg.connect(DSN, autocommit=True) as conn:
|
||||
conn.execute("CREATE SCHEMA realtime")
|
||||
conn.execute("""CREATE TABLE realtime.node_results (
|
||||
time timestamptz NOT NULL, node_id text NOT NULL,
|
||||
actual_demand double precision, total_head double precision,
|
||||
pressure double precision, quality double precision,
|
||||
PRIMARY KEY(time,node_id))""")
|
||||
conn.execute("""CREATE TABLE realtime.link_results (
|
||||
time timestamptz NOT NULL, link_id text NOT NULL, flow double precision,
|
||||
friction double precision, headloss double precision, quality double precision,
|
||||
reaction double precision, setting double precision, status double precision,
|
||||
velocity double precision, PRIMARY KEY(time,link_id))""")
|
||||
for table, element, value in [("node_results", "node_id", "pressure"), ("link_results", "link_id", "flow")]:
|
||||
conn.execute("SELECT create_hypertable(%s,'time',chunk_time_interval=>interval '1 day')", (f"realtime.{table}",))
|
||||
conn.execute(f"""INSERT INTO realtime.{table}(time,{element},{value})
|
||||
SELECT %s::timestamptz + period * interval '15 minutes', id::text, 0
|
||||
FROM generate_series(0,3) period CROSS JOIN generate_series(1,100) id""", (START,))
|
||||
conn.execute(f"ALTER TABLE realtime.{table} SET (timescaledb.compress, timescaledb.compress_orderby='{element},time')")
|
||||
conn.execute("SELECT compress_chunk(show_chunks(%s))", (f"realtime.{table}",))
|
||||
# Scale down the production 100000-row guard: each compressed table has
|
||||
# 400 rows, and deleting one snapshot needs batch decompression.
|
||||
conn.execute("SET timescaledb.max_tuples_decompressed_per_dml_transaction=100")
|
||||
try:
|
||||
yield conn
|
||||
finally:
|
||||
conn.execute("DROP SCHEMA realtime CASCADE")
|
||||
|
||||
|
||||
def snapshot():
|
||||
return (
|
||||
[{"node": str(i), "result": [{"pressure": 999}]} for i in range(1,101)],
|
||||
[{"link": str(i), "result": [{"flow": 888}]} for i in range(1,101)],
|
||||
)
|
||||
|
||||
|
||||
def assert_snapshot(conn, expected_pressure, expected_flow):
|
||||
for table, value, expected in [("node_results", "pressure", expected_pressure), ("link_results", "flow", expected_flow)]:
|
||||
assert conn.execute(f"SELECT count(*),min({value}),max({value}) FROM realtime.{table} WHERE time=%s", (START,)).fetchone() == (100,expected,expected)
|
||||
assert conn.execute(f"SELECT count(*) FROM realtime.{table}").fetchone() == (400,)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("asynchronous", [False, True])
|
||||
def test_compressed_snapshot_replacement_and_transaction_local_budget(monkeypatch, compressed_snapshot, asynchronous):
|
||||
conn = compressed_snapshot
|
||||
nodes, links = snapshot()
|
||||
monkeypatch.setattr(settings, "TIMESCALEDB_REALTIME_MAX_DECOMPRESSED_TUPLES", 1000)
|
||||
# Execute the old DELETE in the actual compressed database, then roll it
|
||||
# back. This proves the guard really fires rather than trusting a SQL mock.
|
||||
with pytest.raises(psycopg.Error, match="tuple decompression limit exceeded"):
|
||||
with conn.transaction():
|
||||
conn.execute("DELETE FROM realtime.node_results WHERE time=%s", (START,))
|
||||
assert_snapshot(conn, 0, 0)
|
||||
|
||||
if asynchronous:
|
||||
async def run():
|
||||
async with await psycopg.AsyncConnection.connect(DSN, autocommit=True) as async_conn:
|
||||
await async_conn.execute("SET timescaledb.max_tuples_decompressed_per_dml_transaction=100")
|
||||
await RealtimeRepository.store_realtime_simulation_result(async_conn, nodes, links, START)
|
||||
cursor = await async_conn.execute("SHOW timescaledb.max_tuples_decompressed_per_dml_transaction")
|
||||
assert await cursor.fetchone() == ("100",)
|
||||
asyncio.run(run())
|
||||
else:
|
||||
RealtimeRepository.store_realtime_simulation_result_sync(conn, nodes, links, START)
|
||||
assert_snapshot(conn, 999, 888)
|
||||
assert conn.execute("SHOW timescaledb.max_tuples_decompressed_per_dml_transaction").fetchone() == ("100",)
|
||||
|
||||
|
||||
def test_copy_failure_rolls_back_both_tables_and_local_budget(monkeypatch, compressed_snapshot):
|
||||
conn = compressed_snapshot
|
||||
monkeypatch.setattr(settings, "TIMESCALEDB_REALTIME_MAX_DECOMPRESSED_TUPLES", 1000)
|
||||
|
||||
def fail_copy(*_):
|
||||
raise RuntimeError("fixture link copy failure")
|
||||
|
||||
monkeypatch.setattr(RealtimeRepository, "_copy_links_sync", fail_copy)
|
||||
nodes, links = snapshot()
|
||||
with pytest.raises(RuntimeError, match="fixture link copy failure"):
|
||||
RealtimeRepository.store_realtime_simulation_result_sync(conn, nodes, links, START)
|
||||
assert_snapshot(conn, 0, 0)
|
||||
assert conn.execute("SHOW timescaledb.max_tuples_decompressed_per_dml_transaction").fetchone() == ("100",)
|
||||
Reference in New Issue
Block a user