session and engine

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)

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