381 lines
14 KiB
Python
381 lines
14 KiB
Python
import asyncio
|
|
import os
|
|
from concurrent.futures import ThreadPoolExecutor
|
|
from uuid import uuid4
|
|
|
|
import pytest
|
|
from psycopg import connect
|
|
|
|
from app.core.config import get_pgconn_string
|
|
from app.infra.db.dynamic_manager import ProjectConnectionManager
|
|
from app.infra.db.postgresql.scada import ScadaInfoRepository
|
|
from app.infra.db.timescaledb.sync_pool import timescale_connection
|
|
from app.native.wndb.commands.api import delete_pattern_cascade
|
|
from app.native.wndb.core.connection import project_connection, project_transaction
|
|
from app.native.wndb.core.projects import have_project, temporary_project_database
|
|
from app.native.wndb.core.database import ChangeSet, g_delete_prefix, write
|
|
from app.native.wndb.model import demands, junctions, patterns
|
|
from app.services.scheme_management import create_analysis_run, update_analysis_run
|
|
|
|
|
|
pytestmark = pytest.mark.skipif(
|
|
os.getenv("RUN_DB_INTEGRATION") != "1",
|
|
reason="set RUN_DB_INTEGRATION=1 to test configured PostgreSQL databases",
|
|
)
|
|
|
|
PROJECT = os.getenv("DB_INTEGRATION_PROJECT", "tjwater_v2")
|
|
|
|
|
|
def _read_business_database(_: int) -> str:
|
|
with project_connection(PROJECT) as conn, conn.cursor() as cur:
|
|
cur.execute("select current_database()")
|
|
return str(cur.fetchone()["current_database"])
|
|
|
|
|
|
def _read_timeseries_database(_: int) -> str:
|
|
with timescale_connection(PROJECT) as conn, conn.cursor() as cur:
|
|
cur.execute("select current_database()")
|
|
return str(cur.fetchone()["current_database"])
|
|
|
|
|
|
def test_business_pool_handles_concurrent_borrows() -> None:
|
|
with ThreadPoolExecutor(max_workers=16) as executor:
|
|
names = list(executor.map(_read_business_database, range(64)))
|
|
|
|
assert names == [PROJECT] * 64
|
|
|
|
|
|
def test_timeseries_pool_handles_concurrent_borrows() -> None:
|
|
with ThreadPoolExecutor(max_workers=16) as executor:
|
|
names = list(executor.map(_read_timeseries_database, range(64)))
|
|
|
|
assert names == [PROJECT] * 64
|
|
|
|
|
|
def test_temporary_project_clone_copies_model_scada_and_views() -> None:
|
|
def counts(project: str) -> dict:
|
|
with project_connection(project) as conn, conn.cursor() as cur:
|
|
cur.execute(
|
|
"""
|
|
select
|
|
(select count(*) from network.nodes) as nodes,
|
|
(select count(*) from network.links) as links,
|
|
(select count(*) from asset.scada_devices) as scada,
|
|
(select count(*) from gis.junctions) as mv_junctions,
|
|
(select count(*) from gis.pipes) as mv_pipes
|
|
"""
|
|
)
|
|
return dict(cur.fetchone())
|
|
|
|
source_counts = counts(PROJECT)
|
|
temporary = None
|
|
with temporary_project_database(PROJECT, "clone_validation") as temporary:
|
|
assert counts(temporary) == source_counts
|
|
|
|
assert temporary is not None
|
|
assert have_project(temporary) is False
|
|
|
|
|
|
def test_dynamic_pool_replaces_terminated_idle_connection() -> None:
|
|
async def exercise_pool() -> None:
|
|
manager = ProjectConnectionManager()
|
|
dsn = get_pgconn_string(db_name=PROJECT)
|
|
try:
|
|
project_id = uuid4()
|
|
async with manager.pg_connection(
|
|
project_id, "biz_data", dsn, 1, 1
|
|
) as conn:
|
|
async with conn.cursor() as cur:
|
|
await cur.execute("select pg_backend_pid()")
|
|
backend_pid = int((await cur.fetchone())["pg_backend_pid"])
|
|
|
|
with connect(dsn, autocommit=True) as admin_conn:
|
|
with admin_conn.cursor() as cur:
|
|
cur.execute("select pg_terminate_backend(%s)", (backend_pid,))
|
|
assert cur.fetchone()[0] is True
|
|
|
|
async with manager.pg_connection(
|
|
project_id, "biz_data", dsn, 1, 1
|
|
) as conn:
|
|
async with conn.cursor() as cur:
|
|
await cur.execute("select current_database()")
|
|
assert (await cur.fetchone())["current_database"] == PROJECT
|
|
finally:
|
|
await manager.close_all()
|
|
|
|
asyncio.run(exercise_pool())
|
|
|
|
|
|
def test_scada_api_repository_reads_device_ids_through_dynamic_pool() -> None:
|
|
async def read_scada_devices() -> list[dict]:
|
|
manager = ProjectConnectionManager()
|
|
try:
|
|
async with manager.pg_connection(
|
|
uuid4(), "biz_data", get_pgconn_string(db_name=PROJECT), 1, 4
|
|
) as conn:
|
|
return await ScadaInfoRepository.get_scadas(conn)
|
|
finally:
|
|
await manager.close_all()
|
|
|
|
devices = asyncio.run(read_scada_devices())
|
|
|
|
assert devices
|
|
device_ids = [device["device_id"] for device in devices]
|
|
assert len(device_ids) == len(set(device_ids))
|
|
assert all(device_ids)
|
|
assert all(
|
|
device["longitude"] is not None and device["latitude"] is not None
|
|
for device in devices
|
|
)
|
|
|
|
|
|
def test_nested_wndb_writes_roll_back_as_one_transaction() -> None:
|
|
with pytest.raises(RuntimeError, match="force rollback"):
|
|
with project_transaction(PROJECT) as conn:
|
|
with conn.cursor() as cur:
|
|
cur.execute(
|
|
"create temporary table wndb_pool_rollback_probe (value integer) on commit drop"
|
|
)
|
|
cur.execute("insert into wndb_pool_rollback_probe values (1)")
|
|
with project_connection(PROJECT) as nested:
|
|
assert nested is conn
|
|
raise RuntimeError("force rollback")
|
|
|
|
with project_connection(PROJECT) as conn, conn.cursor() as cur:
|
|
cur.execute("select to_regclass('pg_temp.wndb_pool_rollback_probe')")
|
|
assert cur.fetchone()["to_regclass"] is None
|
|
|
|
|
|
def test_analysis_run_lifecycle_uses_one_execution_id() -> None:
|
|
run_id = None
|
|
with pytest.raises(RuntimeError, match="force rollback"):
|
|
with project_transaction(PROJECT) as conn:
|
|
run_id = create_analysis_run(
|
|
PROJECT,
|
|
"integration-lifecycle-probe",
|
|
"integration_test",
|
|
"pytest",
|
|
"2026-08-24T00:00:00Z",
|
|
{"temporary": True},
|
|
)
|
|
update_analysis_run(
|
|
PROJECT,
|
|
run_id,
|
|
status="completed",
|
|
username="pytest",
|
|
scheme_detail={"temporary": True},
|
|
)
|
|
with conn.cursor() as cur:
|
|
cur.execute(
|
|
"select status from analysis.runs where run_id = %s", (run_id,)
|
|
)
|
|
assert cur.fetchone()["status"] == "completed"
|
|
raise RuntimeError("force rollback")
|
|
|
|
with project_connection(PROJECT) as conn, conn.cursor() as cur:
|
|
cur.execute("select count(*) as count from analysis.runs where run_id = %s", (run_id,))
|
|
assert cur.fetchone()["count"] == 0
|
|
|
|
|
|
def test_legacy_wndb_batch_treats_malicious_id_as_data() -> None:
|
|
malicious = "integration'); DROP SCHEMA network CASCADE; --"
|
|
command = patterns._add_pattern(
|
|
PROJECT,
|
|
ChangeSet({"id": malicious, "factors": [1.0]}),
|
|
)
|
|
|
|
with pytest.raises(RuntimeError, match="force rollback"):
|
|
with project_transaction(PROJECT) as conn:
|
|
write(PROJECT, command.sql)
|
|
with conn.cursor() as cur:
|
|
cur.execute("select count(*) as count from network.patterns where id = %s", (malicious,))
|
|
assert cur.fetchone()["count"] == 1
|
|
cur.execute("select to_regnamespace('network') as namespace")
|
|
assert cur.fetchone()["namespace"] is not None
|
|
raise RuntimeError("force rollback")
|
|
|
|
|
|
def test_wndb_database_command_pattern_crud_uses_pooled_transaction() -> None:
|
|
pattern_id = f"integration-command-{uuid4()}"
|
|
|
|
with pytest.raises(RuntimeError, match="force rollback"):
|
|
with project_transaction(PROJECT) as conn:
|
|
added = patterns.add_pattern(
|
|
PROJECT,
|
|
ChangeSet({"id": pattern_id, "factors": [1.0, 1.1]}),
|
|
)
|
|
updated = patterns.set_pattern(
|
|
PROJECT,
|
|
ChangeSet({"id": pattern_id, "factors": [0.8, 1.2, 1.0]}),
|
|
)
|
|
|
|
with conn.cursor() as cur:
|
|
cur.execute(
|
|
"select factor from network.pattern_values "
|
|
"where pattern_id = %s order by sequence_no",
|
|
(pattern_id,),
|
|
)
|
|
assert [float(row["factor"]) for row in cur.fetchall()] == [
|
|
0.8,
|
|
1.2,
|
|
1.0,
|
|
]
|
|
|
|
deleted = patterns.delete_pattern(
|
|
PROJECT,
|
|
ChangeSet({"id": pattern_id}),
|
|
)
|
|
with conn.cursor() as cur:
|
|
cur.execute(
|
|
"select count(*) as count from network.patterns where id = %s",
|
|
(pattern_id,),
|
|
)
|
|
assert cur.fetchone()["count"] == 0
|
|
|
|
assert added.operations[0]["operation"] == "add"
|
|
assert updated.operations[0]["operation"] == "update"
|
|
assert deleted.operations == [
|
|
{"operation": "delete", "type": "pattern", "id": pattern_id}
|
|
]
|
|
raise RuntimeError("force rollback")
|
|
|
|
with project_connection(PROJECT) as conn, conn.cursor() as cur:
|
|
cur.execute(
|
|
"select count(*) as count from network.patterns where id = %s",
|
|
(pattern_id,),
|
|
)
|
|
assert cur.fetchone()["count"] == 0
|
|
|
|
|
|
def test_wndb_ordered_detail_tables_use_parent_scoped_primary_keys() -> None:
|
|
expected = {
|
|
"gis.link_vertices": "PRIMARY KEY (link_id, sequence_no)",
|
|
"network.curve_points": "PRIMARY KEY (curve_id, sequence_no)",
|
|
"network.demands": "PRIMARY KEY (junction_id, sequence_no)",
|
|
"network.pattern_flow_samples": "PRIMARY KEY (pattern_id, sequence_no)",
|
|
"network.pattern_values": "PRIMARY KEY (pattern_id, sequence_no)",
|
|
}
|
|
|
|
with project_connection(PROJECT) as conn, conn.cursor() as cur:
|
|
cur.execute(
|
|
"select conrelid::regclass::text as table_name, "
|
|
"pg_get_constraintdef(oid) as definition "
|
|
"from pg_constraint "
|
|
"where conname = any(%s) order by 1",
|
|
([f"{table.rsplit('.', 1)[1]}_pkey" for table in expected],),
|
|
)
|
|
actual = {row["table_name"]: row["definition"] for row in cur.fetchall()}
|
|
|
|
assert actual == expected
|
|
|
|
|
|
def test_gis_unified_views_cover_all_materialized_network_layers() -> None:
|
|
with project_connection(PROJECT) as conn, conn.cursor() as cur:
|
|
cur.execute(
|
|
"""
|
|
SELECT
|
|
(SELECT COUNT(*) FROM gis.network_nodes) AS nodes,
|
|
(SELECT COUNT(*) FROM gis.junctions)
|
|
+ (SELECT COUNT(*) FROM gis.reservoirs)
|
|
+ (SELECT COUNT(*) FROM gis.tanks) AS source_nodes,
|
|
(SELECT COUNT(*) FROM gis.network_links) AS links,
|
|
(SELECT COUNT(*) FROM gis.pipes)
|
|
+ (SELECT COUNT(*) FROM gis.pumps)
|
|
+ (SELECT COUNT(*) FROM gis.valves) AS source_links
|
|
"""
|
|
)
|
|
counts = cur.fetchone()
|
|
cur.execute(
|
|
"""
|
|
SELECT obj_description('gis.network_nodes'::regclass) AS node_comment,
|
|
obj_description('gis.network_links'::regclass) AS link_comment
|
|
"""
|
|
)
|
|
comments = cur.fetchone()
|
|
cur.execute(
|
|
"""
|
|
SELECT COUNT(*) AS undocumented_columns
|
|
FROM pg_attribute
|
|
WHERE attrelid = ANY(
|
|
ARRAY['gis.network_nodes'::regclass, 'gis.network_links'::regclass]
|
|
)
|
|
AND attnum > 0
|
|
AND NOT attisdropped
|
|
AND col_description(attrelid, attnum) IS NULL
|
|
"""
|
|
)
|
|
undocumented_columns = cur.fetchone()["undocumented_columns"]
|
|
|
|
assert counts["nodes"] == counts["source_nodes"]
|
|
assert counts["links"] == counts["source_links"]
|
|
assert comments["node_comment"]
|
|
assert comments["link_comment"]
|
|
assert undocumented_columns == 0
|
|
|
|
|
|
def test_wndb_pattern_cascade_unsets_dependent_demand_atomically() -> None:
|
|
suffix = uuid4()
|
|
junction_id = f"integration-junction-{suffix}"
|
|
pattern_id = f"integration-cascade-{suffix}"
|
|
retained_pattern_id = f"integration-retained-{suffix}"
|
|
|
|
with pytest.raises(RuntimeError, match="force rollback"):
|
|
with project_transaction(PROJECT):
|
|
junctions.add_junction(
|
|
PROJECT,
|
|
ChangeSet(
|
|
{"id": junction_id, "x": 0.0, "y": 0.0, "elevation": 1.0}
|
|
),
|
|
)
|
|
patterns.add_pattern(
|
|
PROJECT,
|
|
ChangeSet({"id": pattern_id, "factors": [1.0]}),
|
|
)
|
|
patterns.add_pattern(
|
|
PROJECT,
|
|
ChangeSet({"id": retained_pattern_id, "factors": [1.0]}),
|
|
)
|
|
demands.set_demand(
|
|
PROJECT,
|
|
ChangeSet(
|
|
{
|
|
"junction": junction_id,
|
|
"demands": [
|
|
{
|
|
"demand": 1.0,
|
|
"pattern": pattern_id,
|
|
"category": "integration",
|
|
},
|
|
{
|
|
"demand": 2.0,
|
|
"pattern": retained_pattern_id,
|
|
"category": "retained",
|
|
},
|
|
],
|
|
}
|
|
),
|
|
)
|
|
|
|
result = delete_pattern_cascade(
|
|
PROJECT,
|
|
ChangeSet(g_delete_prefix | {"id": pattern_id}),
|
|
)
|
|
|
|
assert patterns.get_pattern(PROJECT, pattern_id) == {}
|
|
assert demands.get_demand(PROJECT, junction_id)["demands"] == [
|
|
{"demand": 1.0, "pattern": None, "category": "integration"},
|
|
{
|
|
"demand": 2.0,
|
|
"pattern": retained_pattern_id,
|
|
"category": "retained",
|
|
},
|
|
]
|
|
assert result.operations[-1] == {
|
|
"operation": "delete",
|
|
"type": "pattern",
|
|
"id": pattern_id,
|
|
}
|
|
raise RuntimeError("force rollback")
|
|
|
|
assert junctions.get_junction(PROJECT, junction_id) == {}
|