refactor(db)!: adopt project-routed pooled databases
Reorganize WNDB by responsibility and remove legacy scheme endpoints.\n\nRoute analysis and time-series access through project pools, preserve transactional realtime replacement, and refresh GIS materialized views after writes.\n\nAdd database architecture documentation, live pooling coverage, API contract updates, and executable container verification.\n\nBREAKING CHANGE: legacy scheme APIs and flat app.native.wndb module imports are removed.
This commit is contained in:
@@ -0,0 +1 @@
|
||||
"""WNDB connection, transaction, and project lifecycle infrastructure."""
|
||||
@@ -0,0 +1,224 @@
|
||||
from collections import OrderedDict
|
||||
from collections.abc import Iterator
|
||||
from contextlib import contextmanager
|
||||
from contextvars import ContextVar
|
||||
from threading import RLock
|
||||
|
||||
from psycopg import Connection
|
||||
from psycopg.rows import dict_row
|
||||
from psycopg_pool import ConnectionPool
|
||||
|
||||
from app.core.config import settings
|
||||
from app.infra.db.project_routing import get_project_pgconn_string
|
||||
|
||||
_pools: OrderedDict[str, ConnectionPool] = OrderedDict()
|
||||
_pool_conninfo: dict[str, str] = {}
|
||||
_pool_borrows: dict[str, int] = {}
|
||||
_admin_pools: OrderedDict[str, ConnectionPool] = OrderedDict()
|
||||
_admin_pool_borrows: dict[str, int] = {}
|
||||
_registry_lock = RLock()
|
||||
_active_project_connection: ContextVar[tuple[str, Connection] | None] = ContextVar(
|
||||
"wndb_active_project_connection",
|
||||
default=None,
|
||||
)
|
||||
|
||||
|
||||
def _close_pool(pool: ConnectionPool) -> None:
|
||||
if not pool.closed:
|
||||
pool.close()
|
||||
|
||||
|
||||
def _evict_idle_project_pools(*, protected: str | None = None) -> None:
|
||||
limit = max(1, settings.PROJECT_PG_CACHE_SIZE)
|
||||
while len(_pools) > limit:
|
||||
candidate = next(
|
||||
(key for key in _pools if key != protected and _pool_borrows.get(key, 0) == 0),
|
||||
None,
|
||||
)
|
||||
if candidate is None:
|
||||
return
|
||||
pool = _pools.pop(candidate)
|
||||
_pool_conninfo.pop(candidate, None)
|
||||
_pool_borrows.pop(candidate, None)
|
||||
_close_pool(pool)
|
||||
|
||||
|
||||
def _evict_idle_admin_pools(*, protected: str | None = None) -> None:
|
||||
limit = max(1, settings.PROJECT_PG_CACHE_SIZE)
|
||||
while len(_admin_pools) > limit:
|
||||
candidate = next(
|
||||
(
|
||||
key
|
||||
for key in _admin_pools
|
||||
if key != protected and _admin_pool_borrows.get(key, 0) == 0
|
||||
),
|
||||
None,
|
||||
)
|
||||
if candidate is None:
|
||||
return
|
||||
pool = _admin_pools.pop(candidate)
|
||||
_admin_pool_borrows.pop(candidate, None)
|
||||
_close_pool(pool)
|
||||
|
||||
|
||||
def get_project_pool(name: str) -> ConnectionPool:
|
||||
"""Return the routed synchronous pool used by native WNDB operations."""
|
||||
conninfo = get_project_pgconn_string(db_name=name)
|
||||
with _registry_lock:
|
||||
pool = _pools.get(name)
|
||||
if pool is not None and _pool_conninfo.get(name) == conninfo and not pool.closed:
|
||||
_pools.move_to_end(name)
|
||||
return pool
|
||||
if pool is not None:
|
||||
if _pool_borrows.get(name, 0):
|
||||
raise RuntimeError(f"Cannot replace active project pool {name!r}")
|
||||
_close_pool(pool)
|
||||
pool = ConnectionPool(
|
||||
conninfo=conninfo,
|
||||
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},
|
||||
open=True,
|
||||
)
|
||||
_pools[name] = pool
|
||||
_pool_conninfo[name] = conninfo
|
||||
_pool_borrows.setdefault(name, 0)
|
||||
_evict_idle_project_pools(protected=name)
|
||||
return pool
|
||||
|
||||
|
||||
def is_project_pool_open(name: str) -> bool:
|
||||
with _registry_lock:
|
||||
pool = _pools.get(name)
|
||||
if pool is None or pool.closed:
|
||||
return False
|
||||
if _pool_conninfo.get(name) != get_project_pgconn_string(db_name=name):
|
||||
if _pool_borrows.get(name, 0):
|
||||
return False
|
||||
_close_pool(pool)
|
||||
_pools.pop(name, None)
|
||||
_pool_conninfo.pop(name, None)
|
||||
_pool_borrows.pop(name, None)
|
||||
return False
|
||||
return True
|
||||
|
||||
|
||||
def close_project_pool(name: str) -> None:
|
||||
with _registry_lock:
|
||||
if _pool_borrows.get(name, 0):
|
||||
raise RuntimeError(f"Cannot close active project pool {name!r}")
|
||||
pool = _pools.pop(name, None)
|
||||
_pool_conninfo.pop(name, None)
|
||||
_pool_borrows.pop(name, None)
|
||||
if pool is not None:
|
||||
_close_pool(pool)
|
||||
|
||||
|
||||
def close_all_project_pools() -> None:
|
||||
"""Close every WNDB project pool and the administration pool."""
|
||||
with _registry_lock:
|
||||
pools = [*_pools.values(), *_admin_pools.values()]
|
||||
_pools.clear()
|
||||
_pool_conninfo.clear()
|
||||
_pool_borrows.clear()
|
||||
_admin_pools.clear()
|
||||
_admin_pool_borrows.clear()
|
||||
for pool in pools:
|
||||
_close_pool(pool)
|
||||
|
||||
|
||||
def get_admin_pool() -> ConnectionPool:
|
||||
"""Return the administration pool for the current routed PostgreSQL host."""
|
||||
conninfo = get_project_pgconn_string(db_name="postgres")
|
||||
with _registry_lock:
|
||||
pool = _admin_pools.get(conninfo)
|
||||
if pool is not None and not pool.closed:
|
||||
_admin_pools.move_to_end(conninfo)
|
||||
return pool
|
||||
pool = ConnectionPool(
|
||||
conninfo=conninfo,
|
||||
min_size=settings.PROJECT_PG_POOL_MIN_SIZE,
|
||||
max_size=settings.PROJECT_PG_POOL_SIZE,
|
||||
kwargs={"autocommit": True, "row_factory": dict_row},
|
||||
open=True,
|
||||
)
|
||||
_admin_pools[conninfo] = pool
|
||||
_admin_pool_borrows.setdefault(conninfo, 0)
|
||||
_evict_idle_admin_pools(protected=conninfo)
|
||||
return pool
|
||||
|
||||
|
||||
@contextmanager
|
||||
def project_connection(name: str) -> Iterator[Connection]:
|
||||
"""Borrow one routed WNDB connection and return it to its pool on exit."""
|
||||
active = _active_project_connection.get()
|
||||
if active is not None:
|
||||
active_name, conn = active
|
||||
if active_name != name:
|
||||
raise RuntimeError(
|
||||
f"Cannot access project {name!r} inside transaction for {active_name!r}"
|
||||
)
|
||||
yield conn
|
||||
return
|
||||
with _registry_lock:
|
||||
pool = get_project_pool(name)
|
||||
_pool_borrows[name] = _pool_borrows.get(name, 0) + 1
|
||||
try:
|
||||
with pool.connection() as conn:
|
||||
yield conn
|
||||
finally:
|
||||
with _registry_lock:
|
||||
_pool_borrows[name] -= 1
|
||||
_evict_idle_project_pools()
|
||||
|
||||
|
||||
@contextmanager
|
||||
def project_transaction(name: str) -> Iterator[Connection]:
|
||||
"""Run all nested WNDB operations on one pooled connection and transaction."""
|
||||
active = _active_project_connection.get()
|
||||
if active is not None:
|
||||
active_name, conn = active
|
||||
if active_name != name:
|
||||
raise RuntimeError(
|
||||
f"Cannot nest project {name!r} inside transaction for {active_name!r}"
|
||||
)
|
||||
with conn.transaction():
|
||||
yield conn
|
||||
return
|
||||
|
||||
with _registry_lock:
|
||||
pool = get_project_pool(name)
|
||||
_pool_borrows[name] = _pool_borrows.get(name, 0) + 1
|
||||
try:
|
||||
with pool.connection() as conn:
|
||||
token = _active_project_connection.set((name, conn))
|
||||
try:
|
||||
with conn.transaction():
|
||||
yield conn
|
||||
finally:
|
||||
_active_project_connection.reset(token)
|
||||
finally:
|
||||
with _registry_lock:
|
||||
_pool_borrows[name] -= 1
|
||||
_evict_idle_project_pools()
|
||||
|
||||
|
||||
def is_project_transaction_active(name: str) -> bool:
|
||||
active = _active_project_connection.get()
|
||||
return active is not None and active[0] == name
|
||||
|
||||
|
||||
@contextmanager
|
||||
def admin_connection() -> Iterator[Connection]:
|
||||
"""Borrow a PostgreSQL administration connection from its pool."""
|
||||
with _registry_lock:
|
||||
pool = get_admin_pool()
|
||||
conninfo = pool.conninfo
|
||||
_admin_pool_borrows[conninfo] = _admin_pool_borrows.get(conninfo, 0) + 1
|
||||
try:
|
||||
with pool.connection() as conn:
|
||||
yield conn
|
||||
finally:
|
||||
with _registry_lock:
|
||||
_admin_pool_borrows[conninfo] -= 1
|
||||
_evict_idle_admin_pools()
|
||||
@@ -0,0 +1,143 @@
|
||||
from collections.abc import Mapping, Sequence
|
||||
from typing import Any
|
||||
|
||||
from psycopg import sql
|
||||
from psycopg.rows import Row, dict_row
|
||||
|
||||
from .connection import is_project_transaction_active, project_connection
|
||||
|
||||
API_ADD = "add"
|
||||
API_UPDATE = "update"
|
||||
API_DELETE = "delete"
|
||||
|
||||
g_add_prefix = {"operation": API_ADD}
|
||||
g_update_prefix = {"operation": API_UPDATE}
|
||||
g_delete_prefix = {"operation": API_DELETE}
|
||||
|
||||
|
||||
class ChangeSet:
|
||||
def __init__(self, ps: dict[str, Any] | None = None):
|
||||
self.operations: list[dict[str, Any]] = []
|
||||
if ps is not None:
|
||||
self.append(ps)
|
||||
|
||||
@staticmethod
|
||||
def from_list(ps: list[dict[str, Any]]):
|
||||
change_set = ChangeSet()
|
||||
for item in ps:
|
||||
change_set.append(item)
|
||||
return change_set
|
||||
|
||||
def add(self, ps: dict[str, Any]):
|
||||
self.operations.append(g_add_prefix | ps)
|
||||
return self
|
||||
|
||||
def update(self, ps: dict[str, Any]):
|
||||
self.operations.append(g_update_prefix | ps)
|
||||
return self
|
||||
|
||||
def delete(self, ps: dict[str, Any]):
|
||||
self.operations.append(g_delete_prefix | ps)
|
||||
return self
|
||||
|
||||
def append(self, ps: dict[str, Any]):
|
||||
self.operations.append(ps)
|
||||
return self
|
||||
|
||||
def merge(self, change_set):
|
||||
self.operations.extend(change_set.operations)
|
||||
return self
|
||||
|
||||
def dump(self):
|
||||
for operation in self.operations:
|
||||
print(operation)
|
||||
|
||||
def compress(self):
|
||||
return self
|
||||
|
||||
|
||||
class DatabaseCommand:
|
||||
def __init__(self, statement: str, changes: list[dict[str, Any]]) -> None:
|
||||
self.sql = statement
|
||||
self.changes = changes
|
||||
|
||||
QueryParams = Sequence[Any] | Mapping[str, Any]
|
||||
|
||||
|
||||
def sql_literal(value: Any) -> str:
|
||||
"""Render one PostgreSQL literal for legacy WNDB SQL batch builders.
|
||||
|
||||
WNDB still assembles multi-statement model changes before executing them as
|
||||
one transaction. Every interpolated value must pass through this helper;
|
||||
identifiers remain static strings owned by the backend.
|
||||
"""
|
||||
return sql.Literal(value).as_string()
|
||||
|
||||
|
||||
def _execute(cur, query: str, params: QueryParams | None = None):
|
||||
return cur.execute(query, params) if params is not None else cur.execute(query)
|
||||
|
||||
|
||||
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)
|
||||
row = cur.fetchone()
|
||||
if row is None:
|
||||
raise LookupError(query)
|
||||
return row
|
||||
|
||||
|
||||
def read_all(
|
||||
name: str, query: str, params: QueryParams | None = None
|
||||
) -> list[Row]:
|
||||
with project_connection(name) as conn, conn.cursor(row_factory=dict_row) as cur:
|
||||
_execute(cur, query, params)
|
||||
return cur.fetchall()
|
||||
|
||||
|
||||
def try_read(
|
||||
name: str, query: str, params: QueryParams | None = None
|
||||
) -> Row | None:
|
||||
with project_connection(name) as conn, conn.cursor(row_factory=dict_row) as cur:
|
||||
_execute(cur, query, params)
|
||||
return cur.fetchone()
|
||||
|
||||
|
||||
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)
|
||||
|
||||
|
||||
def refresh_materialized_views(name: str, *, concurrently: bool = True) -> None:
|
||||
"""Refresh the GIS query layer after committed model or asset changes."""
|
||||
with project_connection(name) as conn, conn.cursor() as cur:
|
||||
cur.execute("CALL gis.refresh_all_materialized_views(%s)", (concurrently,))
|
||||
|
||||
|
||||
_MATERIALIZED_VIEW_SOURCES = (
|
||||
"network.nodes",
|
||||
"network.junctions",
|
||||
"network.reservoirs",
|
||||
"network.tanks",
|
||||
"network.links",
|
||||
"network.pipes",
|
||||
"network.pumps",
|
||||
"network.valves",
|
||||
"network.demands",
|
||||
"gis.node_geometries",
|
||||
"gis.link_vertices",
|
||||
"asset.scada_devices",
|
||||
)
|
||||
|
||||
|
||||
def _affects_materialized_views(command: DatabaseCommand) -> bool:
|
||||
statement = command.sql.lower()
|
||||
return any(source in statement for source in _MATERIALIZED_VIEW_SOURCES)
|
||||
|
||||
|
||||
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)
|
||||
return ChangeSet.from_list(command.changes)
|
||||
@@ -0,0 +1,107 @@
|
||||
from psycopg import sql
|
||||
from psycopg.rows import dict_row
|
||||
from .connection import (
|
||||
admin_connection,
|
||||
close_project_pool,
|
||||
get_project_pool,
|
||||
is_project_pool_open,
|
||||
)
|
||||
|
||||
_server_databases = ["template0", "template1", "postgres", "project"]
|
||||
|
||||
|
||||
def list_project() -> list[str]:
|
||||
ps = []
|
||||
with admin_connection() as conn:
|
||||
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,),
|
||||
):
|
||||
ps.append(p["datname"])
|
||||
return ps
|
||||
|
||||
|
||||
def have_project(name: str) -> bool:
|
||||
with admin_connection() as conn:
|
||||
with conn.cursor() as cur:
|
||||
cur.execute("select 1 from pg_database where datname = %s", (name,))
|
||||
return cur.fetchone() is not None
|
||||
|
||||
|
||||
def copy_project(source: str, new: str) -> None:
|
||||
close_project_pool(source)
|
||||
|
||||
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,),
|
||||
)
|
||||
|
||||
|
||||
def create_project(name: str) -> None:
|
||||
return copy_project("project", name)
|
||||
|
||||
|
||||
def delete_project(name: str) -> None:
|
||||
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))
|
||||
)
|
||||
|
||||
|
||||
def clean_project(excluded: list[str] = []) -> None:
|
||||
projects = list_project()
|
||||
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))
|
||||
)
|
||||
|
||||
|
||||
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)
|
||||
Reference in New Issue
Block a user