refactor(db)!: finalize pooled WNDB v2 migration

This commit is contained in:
2026-08-27 17:26:22 +08:00
parent fa188af0b1
commit b74799a39d
105 changed files with 4988 additions and 5565 deletions
+121 -3
View File
@@ -1,12 +1,17 @@
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.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
@@ -17,7 +22,7 @@ pytestmark = pytest.mark.skipif(
reason="set RUN_DB_INTEGRATION=1 to test configured PostgreSQL databases",
)
PROJECT = os.getenv("DB_INTEGRATION_PROJECT", "tjwater_next")
PROJECT = os.getenv("DB_INTEGRATION_PROJECT", "tjwater_v2")
def _read_business_database(_: int) -> str:
@@ -46,6 +51,60 @@ def test_timeseries_pool_handles_concurrent_borrows() -> None:
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_nested_wndb_writes_roll_back_as_one_transaction() -> None:
with pytest.raises(RuntimeError, match="force rollback"):
with project_transaction(PROJECT) as conn:
@@ -186,10 +245,55 @@ def test_wndb_ordered_detail_tables_use_parent_scoped_primary_keys() -> None:
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):
@@ -203,6 +307,10 @@ def test_wndb_pattern_cascade_unsets_dependent_demand_atomically() -> None:
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(
@@ -213,7 +321,12 @@ def test_wndb_pattern_cascade_unsets_dependent_demand_atomically() -> None:
"demand": 1.0,
"pattern": pattern_id,
"category": "integration",
}
},
{
"demand": 2.0,
"pattern": retained_pattern_id,
"category": "retained",
},
],
}
),
@@ -226,7 +339,12 @@ def test_wndb_pattern_cascade_unsets_dependent_demand_atomically() -> None:
assert patterns.get_pattern(PROJECT, pattern_id) == {}
assert demands.get_demand(PROJECT, junction_id)["demands"] == [
{"demand": 1.0, "pattern": None, "category": "integration"}
{"demand": 1.0, "pattern": None, "category": "integration"},
{
"demand": 2.0,
"pattern": retained_pattern_id,
"category": "retained",
},
]
assert result.operations[-1] == {
"operation": "delete",