rebase

app/db.py

@@ -0,0 +1,70 @@
+"""Async SQLAlchemy database session setup."""
+
+from __future__ import annotations
+
+from collections.abc import AsyncGenerator
+
+from sqlalchemy.ext.asyncio import (
+    AsyncSession,
+    async_sessionmaker,
+    create_async_engine,
+)
+from sqlalchemy.orm import DeclarativeBase
+
+from app.config import get_settings
+
+
+class Base(DeclarativeBase):
+    """Declarative base for all ORM models."""
+
+
+_engine = None
+_sessionmaker = None
+
+
+def get_engine():
+    """Lazy-create the global async engine."""
+    global _engine
+    if _engine is None:
+        cfg = get_settings()
+        _engine = create_async_engine(
+            cfg.database_url,
+            echo=cfg.db_echo,
+            pool_size=cfg.db_pool_size,
+            max_overflow=cfg.db_max_overflow,
+            future=True,
+        )
+    return _engine
+
+
+def get_sessionmaker() -> async_sessionmaker[AsyncSession]:
+    """Lazy-create the global session factory."""
+    global _sessionmaker
+    if _sessionmaker is None:
+        _sessionmaker = async_sessionmaker(
+            get_engine(),
+            class_=AsyncSession,
+            expire_on_commit=False,
+            autoflush=False,
+        )
+    return _sessionmaker
+
+
+async def get_db() -> AsyncGenerator[AsyncSession, None]:
+    """FastAPI dependency: yields an async session and rolls back on error."""
+    sm = get_sessionmaker()
+    async with sm() as session:
+        try:
+            yield session
+        except Exception:
+            await session.rollback()
+            raise
+
+
+async def dispose_engine() -> None:
+    """Dispose engine on application shutdown."""
+    global _engine, _sessionmaker
+    if _engine is not None:
+        await _engine.dispose()
+        _engine = None
+        _sessionmaker = None