import os from concurrent.futures import ThreadPoolExecutor from uuid import uuid4 import pytest 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.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_next") 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_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_wndb_pattern_cascade_unsets_dependent_demand_atomically() -> None: suffix = uuid4() junction_id = f"integration-junction-{suffix}" pattern_id = f"integration-cascade-{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]}), ) demands.set_demand( PROJECT, ChangeSet( { "junction": junction_id, "demands": [ { "demand": 1.0, "pattern": pattern_id, "category": "integration", } ], } ), ) 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"} ] assert result.operations[-1] == { "operation": "delete", "type": "pattern", "id": pattern_id, } raise RuntimeError("force rollback") assert junctions.get_junction(PROJECT, junction_id) == {}