jiang
fa188af0b1
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.
225 lines
7.4 KiB
Python
225 lines
7.4 KiB
Python
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()
|