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
+20
View File
@@ -11,6 +11,7 @@ from psycopg_pool import ConnectionPool
from app.core.config import settings
from app.infra.db.project_routing import get_project_pgconn_string
_check_connection = ConnectionPool.check_connection
_pools: OrderedDict[str, ConnectionPool] = OrderedDict()
_pool_conninfo: dict[str, str] = {}
_pool_borrows: dict[str, int] = {}
@@ -21,6 +22,10 @@ _active_project_connection: ContextVar[tuple[str, Connection] | None] = ContextV
"wndb_active_project_connection",
default=None,
)
_active_model_mutation_locks: ContextVar[frozenset[str]] = ContextVar(
"wndb_active_model_mutation_locks",
default=frozenset(),
)
def _close_pool(pool: ConnectionPool) -> None:
@@ -78,6 +83,7 @@ def get_project_pool(name: str) -> ConnectionPool:
min_size=settings.PROJECT_PG_POOL_MIN_SIZE,
max_size=settings.PROJECT_PG_POOL_SIZE + settings.PROJECT_PG_MAX_OVERFLOW,
kwargs={"autocommit": True, "row_factory": dict_row},
check=_check_connection,
open=True,
)
_pools[name] = pool
@@ -140,6 +146,7 @@ def get_admin_pool() -> ConnectionPool:
min_size=settings.PROJECT_PG_POOL_MIN_SIZE,
max_size=settings.PROJECT_PG_POOL_SIZE,
kwargs={"autocommit": True, "row_factory": dict_row},
check=_check_connection,
open=True,
)
_admin_pools[conninfo] = pool
@@ -192,10 +199,12 @@ def project_transaction(name: str) -> Iterator[Connection]:
try:
with pool.connection() as conn:
token = _active_project_connection.set((name, conn))
lock_token = _active_model_mutation_locks.set(frozenset())
try:
with conn.transaction():
yield conn
finally:
_active_model_mutation_locks.reset(lock_token)
_active_project_connection.reset(token)
finally:
with _registry_lock:
@@ -208,6 +217,17 @@ def is_project_transaction_active(name: str) -> bool:
return active is not None and active[0] == name
def is_model_mutation_lock_active(name: str) -> bool:
"""Return whether the current project transaction already owns its model lock."""
return name in _active_model_mutation_locks.get()
def mark_model_mutation_lock_active(name: str) -> None:
"""Record a transaction-scoped advisory lock to avoid duplicate round trips."""
locks = _active_model_mutation_locks.get()
_active_model_mutation_locks.set(locks | {name})
@contextmanager
def admin_connection() -> Iterator[Connection]:
"""Borrow a PostgreSQL administration connection from its pool."""
+103 -6
View File
@@ -1,10 +1,19 @@
from collections.abc import Mapping, Sequence
from typing import Any
from collections.abc import Iterator, Mapping, Sequence
from contextlib import contextmanager
from typing import Any, Callable
from psycopg import sql
from psycopg.rows import Row, dict_row
from .connection import is_project_transaction_active, project_connection
from app.infra.db.project_routing import get_project_database_name
from .connection import (
is_model_mutation_lock_active,
is_project_transaction_active,
mark_model_mutation_lock_active,
project_connection,
project_transaction,
)
API_ADD = "add"
API_UPDATE = "update"
@@ -61,6 +70,20 @@ class DatabaseCommand:
self.sql = statement
self.changes = changes
class MaterializedViewRefreshAfterCommitError(RuntimeError):
"""Report a failed view refresh without implying that the write rolled back."""
changes_committed = True
def __init__(self, project: str) -> None:
self.project = project
super().__init__(
f"Project {project!r} changes were committed, but materialized view "
"refresh failed"
)
QueryParams = Sequence[Any] | Mapping[str, Any]
@@ -78,6 +101,27 @@ def _execute(cur, query: str, params: QueryParams | None = None):
return cur.execute(query, params) if params is not None else cur.execute(query)
def acquire_model_mutation_lock(conn, name: str) -> None:
"""Serialize model replacement and ordinary WNDB mutations per database."""
if is_model_mutation_lock_active(name):
return
physical_name = get_project_database_name(name)
with conn.cursor() as cur:
cur.execute(
"select pg_advisory_xact_lock(hashtextextended(%s, 0))",
(f"tjwater:wndb:model:{physical_name}",),
)
mark_model_mutation_lock_active(name)
@contextmanager
def model_mutation_transaction(name: str) -> Iterator[Any]:
"""Open a project transaction and acquire its model lock before reading."""
with project_transaction(name) as conn:
acquire_model_mutation_lock(conn, name)
yield conn
def read(name: str, query: str, params: QueryParams | None = None) -> Row:
with project_connection(name) as conn, conn.cursor(row_factory=dict_row) as cur:
_execute(cur, query, params)
@@ -104,8 +148,15 @@ def try_read(
def write(name: str, query: str, params: QueryParams | None = None) -> None:
with project_connection(name) as conn, conn.cursor() as cur:
_execute(cur, query, params)
connection_context = (
project_connection(name)
if is_project_transaction_active(name)
else model_mutation_transaction(name)
)
with connection_context as conn:
acquire_model_mutation_lock(conn, name)
with conn.cursor() as cur:
_execute(cur, query, params)
def refresh_materialized_views(name: str, *, concurrently: bool = True) -> None:
@@ -114,6 +165,13 @@ def refresh_materialized_views(name: str, *, concurrently: bool = True) -> None:
cur.execute("CALL gis.refresh_all_materialized_views(%s)", (concurrently,))
def refresh_materialized_views_after_commit(name: str) -> None:
try:
refresh_materialized_views(name)
except Exception as exc:
raise MaterializedViewRefreshAfterCommitError(name) from exc
_MATERIALIZED_VIEW_SOURCES = (
"network.nodes",
"network.junctions",
@@ -135,9 +193,48 @@ def _affects_materialized_views(command: DatabaseCommand) -> bool:
return any(source in statement for source in _MATERIALIZED_VIEW_SOURCES)
_MATERIALIZED_VIEW_ELEMENT_TYPES = frozenset(
{
"junction",
"reservoir",
"tank",
"pipe",
"pump",
"valve",
"demand",
"vertex",
}
)
def changes_affect_materialized_views(change_set: ChangeSet) -> bool:
"""Return whether a dispatched WNDB batch changes a published GIS source."""
return any(
operation.get("type") in _MATERIALIZED_VIEW_ELEMENT_TYPES
for operation in change_set.operations
)
def execute_command(name: str, command: DatabaseCommand) -> ChangeSet:
"""Apply a model mutation without the removed database undo/redo journal."""
write(name, command.sql)
if _affects_materialized_views(command) and not is_project_transaction_active(name):
refresh_materialized_views(name)
refresh_materialized_views_after_commit(name)
return ChangeSet.from_list(command.changes)
def execute_locked_command(
name: str,
builder: Callable[[], DatabaseCommand | None],
) -> ChangeSet:
"""Build a read-modify-write command only after acquiring the model lock."""
nested_transaction = is_project_transaction_active(name)
command: DatabaseCommand | None = None
with model_mutation_transaction(name):
command = builder()
if command is None:
return ChangeSet()
result = execute_command(name, command)
if not nested_transaction and _affects_materialized_views(command):
refresh_materialized_views_after_commit(name)
return result
+284
View File
@@ -0,0 +1,284 @@
from psycopg import Connection, sql
from .connection import project_connection, project_transaction
from .database import acquire_model_mutation_lock
_MODEL_SCHEMAS = ("network", "gis")
_ALLOWED_EXTERNAL_REFERENCES = {
("analysis", "results", "network", "nodes"),
("analysis", "results", "network", "links"),
("asset", "scada_devices", "network", "nodes"),
("asset", "scada_devices", "network", "links"),
}
def _model_tables(conn: Connection) -> list[tuple[str, str]]:
with conn.cursor() as cur:
cur.execute(
"""
select n.nspname as schema_name, c.relname as table_name
from pg_class c
join pg_namespace n on n.oid = c.relnamespace
where n.nspname = any(%s) and c.relkind in ('r', 'p')
order by n.nspname, c.relname
""",
(list(_MODEL_SCHEMAS),),
)
return [(row["schema_name"], row["table_name"]) for row in cur.fetchall()]
def _copy_order(
conn: Connection, tables: list[tuple[str, str]]
) -> list[tuple[str, str]]:
table_set = set(tables)
dependencies: dict[tuple[str, str], set[tuple[str, str]]] = {
table: set() for table in tables
}
with conn.cursor() as cur:
cur.execute(
"""
select source_ns.nspname as source_schema,
source.relname as source_table,
target_ns.nspname as target_schema,
target.relname as target_table
from pg_constraint constraint_row
join pg_class source on source.oid = constraint_row.conrelid
join pg_namespace source_ns on source_ns.oid = source.relnamespace
join pg_class target on target.oid = constraint_row.confrelid
join pg_namespace target_ns on target_ns.oid = target.relnamespace
where constraint_row.contype = 'f'
and source_ns.nspname = any(%s)
and target_ns.nspname = any(%s)
""",
(list(_MODEL_SCHEMAS), list(_MODEL_SCHEMAS)),
)
for row in cur.fetchall():
source = (row["source_schema"], row["source_table"])
target = (row["target_schema"], row["target_table"])
if source in table_set and target in table_set and source != target:
dependencies[source].add(target)
ordered: list[tuple[str, str]] = []
remaining = set(tables)
while remaining:
ready = sorted(
table for table in remaining if not (dependencies[table] & remaining)
)
if not ready:
cycle = ", ".join(f"{schema}.{table}" for schema, table in sorted(remaining))
raise RuntimeError(f"Model table foreign-key cycle detected: {cycle}")
ordered.extend(ready)
remaining.difference_update(ready)
return ordered
def _table_columns(
conn: Connection, schema_name: str, table_name: str
) -> list[str]:
with conn.cursor() as cur:
cur.execute(
"""
select column_name
from information_schema.columns
where table_schema = %s and table_name = %s
and is_generated = 'NEVER'
order by ordinal_position
""",
(schema_name, table_name),
)
return [row["column_name"] for row in cur.fetchall()]
def _external_references(conn: Connection) -> set[tuple[str, str, str, str]]:
with conn.cursor() as cur:
cur.execute(
"""
select source_ns.nspname as source_schema,
source.relname as source_table,
target_ns.nspname as target_schema,
target.relname as target_table
from pg_constraint constraint_row
join pg_class source on source.oid = constraint_row.conrelid
join pg_namespace source_ns on source_ns.oid = source.relnamespace
join pg_class target on target.oid = constraint_row.confrelid
join pg_namespace target_ns on target_ns.oid = target.relnamespace
where constraint_row.contype = 'f'
and target_ns.nspname = any(%s)
and source_ns.nspname <> all(%s)
""",
(list(_MODEL_SCHEMAS), list(_MODEL_SCHEMAS)),
)
return {
(
row["source_schema"],
row["source_table"],
row["target_schema"],
row["target_table"],
)
for row in cur.fetchall()
}
def _copy_table(
source_conn: Connection,
target_conn: Connection,
schema_name: str,
table_name: str,
columns: list[str],
) -> None:
relation = sql.Identifier(schema_name, table_name)
column_list = sql.SQL(", ").join(map(sql.Identifier, columns))
copy_out = sql.SQL("copy {} ({}) to stdout").format(relation, column_list)
copy_in = sql.SQL("copy {} ({}) from stdin").format(relation, column_list)
with source_conn.cursor().copy(copy_out) as source_copy:
with target_conn.cursor().copy(copy_in) as target_copy:
for chunk in source_copy:
target_copy.write(chunk)
def replace_project_model(
target_project: str,
source_project: str,
*,
copy_source_scada: bool = False,
) -> None:
"""Atomically replace WNDB/GIS model tables from a validated staging DB.
Business and analysis tables remain in the target database. Historical
analysis rows keep references that still exist and become element-neutral
when an element disappeared. SCADA devices are retained only when their
bound node or link still exists in the replacement model.
``copy_source_scada`` is used only for temporary project clones. Normal INP
replacement keeps the target project's existing device mappings and drops
mappings whose model element no longer exists.
"""
with project_connection(source_project) as source_conn, source_conn.transaction():
with source_conn.cursor() as cur:
cur.execute("set transaction isolation level repeatable read, read only")
source_tables = _model_tables(source_conn)
source_columns = {
table: _table_columns(source_conn, *table) for table in source_tables
}
source_scada_columns = (
_table_columns(source_conn, "asset", "scada_devices")
if copy_source_scada
else []
)
copy_order = _copy_order(source_conn, source_tables)
with project_transaction(target_project) as target_conn:
acquire_model_mutation_lock(target_conn, target_project)
target_tables = _model_tables(target_conn)
if set(target_tables) != set(source_tables):
missing = sorted(set(source_tables) - set(target_tables))
extra = sorted(set(target_tables) - set(source_tables))
raise RuntimeError(
f"Staging/target model schema mismatch; missing={missing}, extra={extra}"
)
for table, columns in source_columns.items():
if _table_columns(target_conn, *table) != columns:
raise RuntimeError(
f"Staging/target columns differ for {table[0]}.{table[1]}"
)
if copy_source_scada and _table_columns(
target_conn, "asset", "scada_devices"
) != source_scada_columns:
raise RuntimeError(
"Source/target columns differ for asset.scada_devices"
)
unexpected = _external_references(target_conn) - _ALLOWED_EXTERNAL_REFERENCES
if unexpected:
formatted = ", ".join(
f"{source_schema}.{source_table}->{target_schema}.{target_table}"
for source_schema, source_table, target_schema, target_table in sorted(
unexpected
)
)
raise RuntimeError(
f"Model replacement has unsupported external references: {formatted}"
)
with target_conn.cursor() as cur:
if not copy_source_scada:
cur.execute(
"create temporary table model_scada_snapshot on commit drop "
"as table asset.scada_devices"
)
cur.execute("delete from asset.scada_devices")
cur.execute(
"create temporary table model_result_refs on commit drop as "
"select result_id, node_id, link_id from analysis.results"
)
cur.execute("update analysis.results set node_id = null, link_id = null")
for schema_name, table_name in reversed(copy_order):
cur.execute(
sql.SQL("delete from {}").format(
sql.Identifier(schema_name, table_name)
)
)
for schema_name, table_name in copy_order:
_copy_table(
source_conn,
target_conn,
schema_name,
table_name,
source_columns[(schema_name, table_name)],
)
if copy_source_scada:
_copy_table(
source_conn,
target_conn,
"asset",
"scada_devices",
source_scada_columns,
)
with target_conn.cursor() as cur:
if not copy_source_scada:
cur.execute(
"""
insert into asset.scada_devices
select snapshot.*
from model_scada_snapshot snapshot
where (
snapshot.node_id is not null
and exists (
select 1 from network.nodes node
where node.id = snapshot.node_id
)
) or (
snapshot.link_id is not null
and exists (
select 1 from network.links link
where link.id = snapshot.link_id
)
)
"""
)
cur.execute(
"""
update analysis.results result
set node_id = case
when exists (
select 1 from network.nodes node
where node.id = refs.node_id
) then refs.node_id
else null
end,
link_id = case
when exists (
select 1 from network.links link
where link.id = refs.link_id
) then refs.link_id
else null
end
from model_result_refs refs
where refs.result_id = result.result_id
"""
)
+222 -62
View File
@@ -1,13 +1,46 @@
from collections.abc import Iterable
from contextlib import contextmanager
import re
from uuid import uuid4
from psycopg import sql
from psycopg.rows import dict_row
from app.core.config import settings
from app.infra.db.project_routing import (
get_project_database_name,
get_project_template_database_name,
)
from .connection import (
admin_connection,
close_project_pool,
get_project_pool,
is_project_pool_open,
)
_server_databases = ["template0", "template1", "postgres", "project"]
_SERVER_DATABASES = frozenset({"template0", "template1", "postgres", "project"})
_TEMPORARY_DATABASE_PREFIX = "tjw_tmp_"
def _protected_databases() -> frozenset[str]:
return _SERVER_DATABASES | {
settings.METADATA_DB_NAME,
settings.WNDB_TEMPLATE_DB_NAME,
}
def _validate_project_database(name: str, *, allow_template_source: bool = False) -> None:
if not name:
raise ValueError("Project database name must not be empty")
protected = {database.casefold() for database in _protected_databases()}
is_template = name.casefold().endswith("_template")
if (
allow_template_source
and name.casefold() == settings.WNDB_TEMPLATE_DB_NAME.casefold()
):
return
if name.casefold() in protected or is_template:
raise ValueError(f"Database {name!r} is protected and cannot be managed as a project")
def list_project() -> list[str]:
@@ -16,92 +49,219 @@ def list_project() -> list[str]:
with conn.cursor(row_factory=dict_row) as cur:
for p in cur.execute(
"select datname from pg_database where datname <> all(%s) order by datname",
(_server_databases,),
(list(_protected_databases()),),
):
ps.append(p["datname"])
if not str(p["datname"]).casefold().endswith("_template"):
ps.append(p["datname"])
return ps
@contextmanager
def _database_locks(cur, *database_names: str):
"""Serialize physical database lifecycle operations across server workers."""
lock_names = sorted(set(database_names), key=str.casefold)
for database_name in lock_names:
cur.execute(
"select pg_advisory_lock(hashtextextended(%s, 0))",
(f"tjwater:wndb:{database_name}",),
)
try:
yield
finally:
for database_name in reversed(lock_names):
cur.execute(
"select pg_advisory_unlock(hashtextextended(%s, 0))",
(f"tjwater:wndb:{database_name}",),
)
@contextmanager
def _temporary_database_capacity(cur, database_name: str):
"""Serialize temporary creation and enforce a server-wide hard limit."""
if not database_name.startswith(_TEMPORARY_DATABASE_PREFIX):
yield
return
lock_name = "tjwater:wndb:temporary-database-capacity"
cur.execute("select pg_advisory_lock(hashtextextended(%s, 0))", (lock_name,))
try:
cur.execute(
"select count(*) as count from pg_database where datname like %s",
(f"{_TEMPORARY_DATABASE_PREFIX}%",),
)
row = cur.fetchone()
count = int(row["count"] if row is not None else 0)
limit = max(1, settings.WNDB_TEMP_DB_MAX_COUNT)
if count >= limit:
raise RuntimeError(
f"Temporary database limit reached ({count}/{limit}); "
"retry after an active analysis completes"
)
yield
finally:
cur.execute("select pg_advisory_unlock(hashtextextended(%s, 0))", (lock_name,))
def _database_allows_connections(cur, database_name: str) -> bool:
cur.execute(
"select datallowconn from pg_database where datname = %s",
(database_name,),
)
row = cur.fetchone()
if row is None:
raise ValueError(f"Database {database_name!r} does not exist")
return bool(row["datallowconn"])
def _set_database_connections(cur, database_name: str, *, allowed: bool) -> None:
cur.execute(
"update pg_database set datallowconn = %s where datname = %s",
(allowed, database_name),
)
def temporary_project_name(project: str, purpose: str) -> str:
"""Return a collision-resistant physical database name for one run."""
physical_name = get_project_database_name(project)
safe_purpose = re.sub(r"[^a-z0-9_]+", "_", purpose.casefold()).strip("_")
safe_project = re.sub(r"[^a-z0-9_]+", "_", physical_name.casefold()).strip("_")
prefix = (
f"{_TEMPORARY_DATABASE_PREFIX}{safe_purpose or 'run'}_"
f"{safe_project or 'project'}"
)[:29]
return f"{prefix}_{uuid4().hex}"
@contextmanager
def temporary_project_database(project: str, purpose: str):
"""Clone one project's runnable model into an isolated temporary database."""
temporary_name = temporary_project_name(project, purpose)
try:
copy_project(get_project_template_database_name(project), temporary_name)
# Import lazily to keep physical database lifecycle independent from
# model-copy implementation details at module import time.
from .database import refresh_materialized_views_after_commit
from .model_replace import replace_project_model
replace_project_model(
temporary_name,
project,
copy_source_scada=True,
)
refresh_materialized_views_after_commit(temporary_name)
yield temporary_name
finally:
if have_project(temporary_name):
delete_project(temporary_name)
@contextmanager
def temporary_template_database(name_hint: str, purpose: str):
"""Create an empty schema-only temporary database from the fixed template."""
temporary_name = temporary_project_name(name_hint, purpose)
try:
copy_project(get_project_template_database_name(name_hint), temporary_name)
yield temporary_name
finally:
if have_project(temporary_name):
delete_project(temporary_name)
def have_project(name: str) -> bool:
database_name = get_project_database_name(name)
with admin_connection() as conn:
with conn.cursor() as cur:
cur.execute("select 1 from pg_database where datname = %s", (name,))
cur.execute("select 1 from pg_database where datname = %s", (database_name,))
return cur.fetchone() is not None
def copy_project(source: str, new: str) -> None:
physical_source = get_project_database_name(source)
physical_new = get_project_database_name(new)
_validate_project_database(physical_source, allow_template_source=True)
_validate_project_database(physical_new)
close_project_pool(source)
close_project_pool(new)
with admin_connection() as admin_conn:
with admin_conn.cursor() as cur:
cur.execute(
"update pg_database set datallowconn = false where datname = %s",
(source,),
)
try:
cur.execute(
"select pg_terminate_backend(pid) from pg_stat_activity where datname = %s and pid <> pg_backend_pid()",
(source,),
)
cur.execute(
sql.SQL("create database {} with template = {}").format(
sql.Identifier(new), sql.Identifier(source)
)
)
finally:
cur.execute(
"update pg_database set datallowconn = true where datname = %s",
(source,),
)
with _temporary_database_capacity(cur, physical_new):
with _database_locks(cur, physical_source, physical_new):
source_allowed = _database_allows_connections(cur, physical_source)
if source_allowed:
_set_database_connections(
cur,
physical_source,
allowed=False,
)
try:
cur.execute(
"select pg_terminate_backend(pid) from pg_stat_activity "
"where datname = %s and pid <> pg_backend_pid()",
(physical_source,),
)
cur.execute(
sql.SQL("create database {} with template = {}").format(
sql.Identifier(physical_new),
sql.Identifier(physical_source),
)
)
finally:
if source_allowed:
_set_database_connections(
cur,
physical_source,
allowed=True,
)
def create_project(name: str) -> None:
return copy_project("project", name)
return copy_project(get_project_template_database_name(name), name)
def delete_project(name: str) -> None:
database_name = get_project_database_name(name)
_validate_project_database(database_name)
close_project_pool(name)
with admin_connection() as conn:
with conn.cursor() as cur:
cur.execute(
"select pg_terminate_backend(pid) from pg_stat_activity "
"where datname = %s and pid <> pg_backend_pid()",
(name,),
)
cur.execute(
sql.SQL("drop database {}").format(sql.Identifier(name))
)
with _database_locks(cur, database_name):
was_allowed = _database_allows_connections(cur, database_name)
if was_allowed:
_set_database_connections(cur, database_name, allowed=False)
try:
cur.execute(
"select pg_terminate_backend(pid) from pg_stat_activity "
"where datname = %s and pid <> pg_backend_pid()",
(database_name,),
)
cur.execute(
sql.SQL("drop database {}").format(
sql.Identifier(database_name)
)
)
except Exception:
if was_allowed:
_set_database_connections(cur, database_name, allowed=True)
raise
def clean_project(excluded: list[str] = []) -> None:
projects = list_project()
def clean_project(projects: Iterable[str]) -> None:
"""Delete only the explicitly supplied project databases."""
targets = list(dict.fromkeys(projects))
physical_targets = [get_project_database_name(project) for project in targets]
for database_name in physical_targets:
_validate_project_database(database_name)
if not targets:
return
with admin_connection() as conn:
with conn.cursor(row_factory=dict_row) as cur:
row = cur.execute("select current_database()").fetchone()
if row != None:
current_db = row["current_database"]
if current_db in projects:
projects.remove(current_db)
for project in projects:
if project in _server_databases or project in excluded:
continue
cur.execute(
"select pg_terminate_backend(pid) from pg_stat_activity "
"where datname = %s and pid <> pg_backend_pid()",
(project,),
)
cur.execute(
sql.SQL("drop database {}").format(sql.Identifier(project))
)
current_db = row["current_database"] if row is not None else None
if current_db in physical_targets:
raise ValueError(f"Cannot delete the current database {current_db!r}")
def open_project(name: str) -> None:
get_project_pool(name)
def is_project_open(name: str) -> bool:
return is_project_pool_open(name)
def close_project(name: str) -> None:
close_project_pool(name)
for project in targets:
delete_project(project)