70 lines
1.9 KiB
Python
70 lines
1.9 KiB
Python
|
"""Async SQLAlchemy engine, session factory, and FastAPI dependency.
|
||
|
|
|
||
|
|
Driver-agnostic: works with both aiosqlite (dev) and asyncpg (prod).
|
||
|
|
"""
|
||
|
|
|
||
|
|
from __future__ import annotations
|
||
|
|
|
||
|
|
from collections.abc import AsyncGenerator
|
||
|
|
from typing import Any
|
||
|
|
|
||
|
|
from sqlalchemy.ext.asyncio import (
|
||
|
|
AsyncEngine,
|
||
|
|
AsyncSession,
|
||
|
|
async_sessionmaker,
|
||
|
|
create_async_engine,
|
||
|
|
)
|
||
|
|
from sqlalchemy.orm import DeclarativeBase
|
||
|
|
|
||
|
|
from app.core.config import get_settings
|
||
|
|
|
||
|
|
|
||
|
|
class Base(DeclarativeBase):
|
||
|
|
"""Declarative base for all ORM models.
|
||
|
|
|
||
|
|
Re-exported here so Alembic env.py can import it from a single location.
|
||
|
|
"""
|
||
|
|
|
||
|
|
|
||
|
|
# === Engine & Session Factory ===
|
||
|
|
|
||
|
|
_settings = get_settings()
|
||
|
|
|
||
|
|
# SQLite needs check_same_thread=False even for async — handled by aiosqlite.
|
||
|
|
# echo=False in production; controlled by env if needed later.
|
||
|
|
_engine_kwargs: dict[str, Any] = {"echo": False, "future": True}
|
||
|
|
|
||
|
|
# For PostgreSQL asyncpg, pool sizing matters; for SQLite it's irrelevant.
|
||
|
|
if not _settings.DATABASE_URL.startswith("sqlite"):
|
||
|
|
_engine_kwargs["pool_size"] = 5
|
||
|
|
_engine_kwargs["max_overflow"] = 10
|
||
|
|
_engine_kwargs["pool_pre_ping"] = True
|
||
|
|
|
||
|
|
engine: AsyncEngine = create_async_engine(_settings.DATABASE_URL, **_engine_kwargs)
|
||
|
|
|
||
|
|
AsyncSessionLocal: async_sessionmaker[AsyncSession] = async_sessionmaker(
|
||
|
|
bind=engine,
|
||
|
|
expire_on_commit=False,
|
||
|
|
class_=AsyncSession,
|
||
|
|
autoflush=False,
|
||
|
|
)
|
||
|
|
|
||
|
|
|
||
|
|
# === Dependency ===
|
||
|
|
|
||
|
|
|
||
|
|
async def get_db() -> AsyncGenerator[AsyncSession, None]:
|
||
|
|
"""FastAPI dependency that yields an AsyncSession and ensures cleanup."""
|
||
|
|
async with AsyncSessionLocal() as session:
|
||
|
|
try:
|
||
|
|
yield session
|
||
|
|
except Exception:
|
||
|
|
await session.rollback()
|
||
|
|
raise
|
||
|
|
# No explicit close needed — async context manager handles it.
|
||
|
|
|
||
|
|
|
||
|
|
async def dispose_engine() -> None:
|
||
|
|
"""Dispose of the engine on application shutdown."""
|
||
|
|
await engine.dispose()
|