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.
102 lines
5.3 KiB
Python
102 lines
5.3 KiB
Python
"""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",)
|