From 9eb0f40c08a9716891162a7a153b0c28fe52685d Mon Sep 17 00:00:00 2001 From: John Lancaster <32917998+jsl12@users.noreply.github.com> Date: Thu, 30 Jul 2026 22:27:21 -0500 Subject: [PATCH] session and engine --- src/transcription/db/engine.py | 60 +++++++++++++++++++++++++ src/transcription/db/session.py | 79 +++++++++++++++++++++++++++++++++ 2 files changed, 139 insertions(+) create mode 100644 src/transcription/db/engine.py create mode 100644 src/transcription/db/session.py diff --git a/src/transcription/db/engine.py b/src/transcription/db/engine.py new file mode 100644 index 0000000..89c541c --- /dev/null +++ b/src/transcription/db/engine.py @@ -0,0 +1,60 @@ +from functools import cache +from typing import Any + +from sqlalchemy import URL +from sqlalchemy import StaticPool +from sqlalchemy.ext.asyncio import AsyncEngine +from sqlalchemy.ext.asyncio import create_async_engine + +from ..config import PostgresSettings +from ..config import Settings +from ..config import SqliteSettings +from ..config import get_settings + + +def get_database_url(settings: Settings) -> str: + match settings.database: + case SqliteSettings(path=path): + url = URL.create( + drivername="sqlite+aiosqlite", + database=path, + ) + case PostgresSettings() as database: + url = URL.create( + drivername="postgresql+asyncpg", + host=database.host, + port=database.port, + database=database.database, + username=database.user, + password=database.password.get_secret_value(), + ) + return url.render_as_string(hide_password=False) + + +def resolve_engine(settings: Settings | None = None) -> AsyncEngine: + active_settings = settings or get_settings() + return get_engine(get_database_url(active_settings)) + + +@cache +def get_engine(database_url: str) -> AsyncEngine: + kwargs: dict[str, Any] = {"echo": False, "pool_pre_ping": True} + if database_url.startswith("sqlite"): + kwargs["connect_args"] = {"check_same_thread": False} + if ":memory:" in database_url: + kwargs["poolclass"] = StaticPool + + return create_async_engine(database_url, **kwargs) + + +async def dispose_engine(database_url: str) -> None: + engine = get_engine(database_url) + try: + await engine.dispose() + finally: + get_engine.cache_clear() + + +async def refresh_engine(database_url: str) -> AsyncEngine: + await dispose_engine(database_url) + return get_engine(database_url) diff --git a/src/transcription/db/session.py b/src/transcription/db/session.py new file mode 100644 index 0000000..9135f16 --- /dev/null +++ b/src/transcription/db/session.py @@ -0,0 +1,79 @@ +from collections.abc import AsyncGenerator +from contextlib import asynccontextmanager +from functools import cache +from typing import Annotated + +from fastapi import Depends +from sqlalchemy.ext.asyncio import AsyncSessionTransaction +from sqlalchemy.ext.asyncio import async_sessionmaker +from sqlmodel.ext.asyncio.session import AsyncSession + +from ..config import get_settings +from .engine import dispose_engine +from .engine import get_database_url +from .engine import get_engine + +type SessionFactory = async_sessionmaker[AsyncSession] + + +@cache +def get_session_factory(database_url: str) -> SessionFactory: + return async_sessionmaker( + bind=get_engine(database_url), + class_=AsyncSession, + expire_on_commit=False, + ) + + +def resolve_session_factory(database_url: str | None = None) -> SessionFactory: + return get_session_factory(database_url or get_database_url(get_settings())) + + +type SessionFactoryDep = Annotated[SessionFactory, Depends(resolve_session_factory)] + + +async def dispose_session_factory(database_url: str) -> None: + get_session_factory.cache_clear() + await dispose_engine(database_url) + + +@asynccontextmanager +async def session_scope( + *, + database_url: str | None = None, + session: AsyncSession | None = None, +) -> AsyncGenerator[AsyncSession]: + if session is not None: + yield session + return + + session_factory = resolve_session_factory(database_url) + async with session_factory() as owned_session: + yield owned_session + + +type SessionScopeDep = Annotated[AsyncSession, Depends(session_scope)] + + +@asynccontextmanager +async def transaction_scope( + *, + database_url: str | None = None, + session: AsyncSessionTransaction | None = None, +) -> AsyncGenerator[AsyncSessionTransaction]: + match session: + case AsyncSession() as async_session: + if not async_session.in_transaction(): + raise RuntimeError("A supplied session must have an active transaction") + yield async_session + return + case AsyncSessionTransaction() as async_transaction: + yield async_transaction + return + + session_factory = resolve_session_factory(database_url) + async with session_factory().begin() as owned_session: + yield owned_session + + +type TransactionScopeDep = Annotated[AsyncSessionTransaction, Depends(transaction_scope)]