Files
TJWaterServerBinary/app/infra/db/project_routing.py
T

94 lines
3.1 KiB
Python

from __future__ import annotations
from contextlib import contextmanager
from contextvars import ContextVar, Token
from dataclasses import dataclass
from typing import Iterator
from psycopg.conninfo import conninfo_to_dict, make_conninfo
from app.core.config import get_pgconn_string, get_timescaledb_pgconn_string, settings
@dataclass(frozen=True)
class ActiveProjectRouting:
project_code: str
business_dsn: str
timescale_dsn: str | None = None
@property
def business_database_name(self) -> str:
"""Return the physical BizDB name selected by metadata routing."""
database_name = conninfo_to_dict(self.business_dsn).get("dbname")
if not database_name:
raise RuntimeError(
f"Business database routing for project {self.project_code!r} "
"does not contain a database name"
)
return database_name
_active_project_routing: ContextVar[ActiveProjectRouting | None] = ContextVar(
"active_project_routing",
default=None,
)
def get_active_project_routing() -> ActiveProjectRouting | None:
return _active_project_routing.get()
@contextmanager
def activate_project_routing(
routing: ActiveProjectRouting,
) -> Iterator[ActiveProjectRouting]:
token: Token[ActiveProjectRouting | None] = _active_project_routing.set(routing)
try:
yield routing
finally:
_active_project_routing.reset(token)
def _dsn_for_database(dsn: str, database_name: str) -> str:
return make_conninfo(dsn, dbname=database_name)
def get_project_database_name(name: str) -> str:
"""Resolve a logical project code to its routed physical BizDB name."""
routing = get_active_project_routing()
if routing is not None and name == routing.project_code:
return routing.business_database_name
return name
def get_project_template_database_name(name: str | None = None) -> str:
"""Return the configured immutable template for the WNDB schema version.
The template belongs to the database schema version, not to an individual
logical project or its temporary physical database name.
"""
return settings.WNDB_TEMPLATE_DB_NAME
def get_project_pgconn_string(db_name: str | None = None) -> str:
routing = get_active_project_routing()
if routing is None:
return get_pgconn_string(db_name=db_name)
if db_name is None or db_name == routing.project_code:
return routing.business_dsn
return _dsn_for_database(routing.business_dsn, db_name)
def get_project_timescale_pgconn_string(db_name: str | None = None) -> str:
routing = get_active_project_routing()
if routing is None:
return get_timescaledb_pgconn_string(db_name=db_name)
if routing.timescale_dsn is None:
raise RuntimeError(
f"TimescaleDB routing is not configured for project {routing.project_code}"
)
# Legacy simulation code used to derive the Timescale database name from
# the project code. Project-scoped requests must instead use the complete
# iot_data DSN selected by metadata routing.
return routing.timescale_dsn