CI / ruff · pytest (pull_request) Successful in 45s
Two DeclarativeBases on two engines (ADR-003): identity holds direct identifiers, clinical holds token-keyed pseudonymised data plus audit_log. Portable types (JSONB variant, tz-aware UTC timestamps rejecting naive values) so tests run on SQLite. init_db() runs create_all at startup. Guardrail #5 checked by AST and import-graph tests. Refs #6 Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
90 lines
2.9 KiB
Python
90 lines
2.9 KiB
Python
"""Engines and sessions for the two databases (ADR-003).
|
|
|
|
This module knows nothing about models: clinical code can import it without pulling in the identity
|
|
schema (guardrail #5). Schema creation lives in `app.db.bootstrap`.
|
|
"""
|
|
|
|
from collections.abc import AsyncIterator
|
|
from contextlib import AbstractAsyncContextManager, asynccontextmanager
|
|
|
|
from sqlalchemy.ext.asyncio import AsyncEngine, AsyncSession, async_sessionmaker, create_async_engine
|
|
from sqlalchemy.pool import StaticPool
|
|
|
|
from app.config import get_settings
|
|
|
|
_engines: dict[str, AsyncEngine] = {}
|
|
_makers: dict[str, async_sessionmaker[AsyncSession]] = {}
|
|
_urls: dict[str, str] = {}
|
|
|
|
|
|
def _make_engine(url: str) -> AsyncEngine:
|
|
kwargs: dict[str, object] = {}
|
|
if url.startswith("postgresql"):
|
|
# Fail fast when the DB is unreachable instead of hanging on asyncpg's 60 s default.
|
|
kwargs = {"pool_pre_ping": True, "connect_args": {"timeout": 5}}
|
|
elif url.startswith("sqlite") and (":memory:" in url or url.rstrip("/").endswith(":")):
|
|
# In-memory SQLite (tests): one shared connection, otherwise every connection sees an empty DB.
|
|
kwargs = {"poolclass": StaticPool}
|
|
return create_async_engine(url, **kwargs)
|
|
|
|
|
|
def _url(name: str) -> str:
|
|
if name in _urls:
|
|
return _urls[name]
|
|
settings = get_settings()
|
|
return settings.identity_db_url if name == "identity" else settings.clinical_db_url
|
|
|
|
|
|
def _engine(name: str) -> AsyncEngine:
|
|
if name not in _engines:
|
|
_engines[name] = _make_engine(_url(name))
|
|
_makers[name] = async_sessionmaker(_engines[name], expire_on_commit=False)
|
|
return _engines[name]
|
|
|
|
|
|
def identity_engine() -> AsyncEngine:
|
|
return _engine("identity")
|
|
|
|
|
|
def clinical_engine() -> AsyncEngine:
|
|
return _engine("clinical")
|
|
|
|
|
|
async def configure(*, identity_url: str | None = None, clinical_url: str | None = None) -> None:
|
|
"""Override DB URLs (tests, scripts). Disposes engines created so far; new ones are created lazily."""
|
|
await dispose()
|
|
_urls.clear()
|
|
if identity_url:
|
|
_urls["identity"] = identity_url
|
|
if clinical_url:
|
|
_urls["clinical"] = clinical_url
|
|
|
|
|
|
@asynccontextmanager
|
|
async def _session(name: str) -> AsyncIterator[AsyncSession]:
|
|
_engine(name)
|
|
async with _makers[name]() as session:
|
|
try:
|
|
yield session
|
|
await session.commit()
|
|
except BaseException:
|
|
await session.rollback()
|
|
raise
|
|
|
|
|
|
def get_identity_session() -> AbstractAsyncContextManager[AsyncSession]:
|
|
"""Unit of work on the identity DB: commits on success, rolls back on error."""
|
|
return _session("identity")
|
|
|
|
|
|
def get_clinical_session() -> AbstractAsyncContextManager[AsyncSession]:
|
|
"""Unit of work on the clinical DB (incl. audit_log): commits on success, rolls back on error."""
|
|
return _session("clinical")
|
|
|
|
|
|
async def dispose() -> None:
|
|
for engine in _engines.values():
|
|
await engine.dispose()
|
|
_engines.clear()
|
|
_makers.clear()
|