66 lines
2.0 KiB
Python
66 lines
2.0 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 make_conninfo
|
|
|
|
from app.core.config import get_pgconn_string, get_timescaledb_pgconn_string
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class ActiveProjectRouting:
|
|
project_code: str
|
|
business_dsn: str
|
|
timescale_dsn: str | None = None
|
|
|
|
|
|
_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_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
|