refactor(db)!: finalize pooled WNDB v2 migration
This commit is contained in:
@@ -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."""
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
"""
|
||||
)
|
||||
@@ -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)
|
||||
|
||||
Reference in New Issue
Block a user