Files
PTSRaiAgent/app/db/session.py
ElebobaandClaude Opus 5.5 121b607414
CI / ruff · pytest (pull_request) Successful in 45s
feat(backend): data model — identity / clinical / audit, repositories, create_all
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>
2026-10-03 14:58:26 +03:00

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()