Files

57 lines
1.3 KiB
Python
Raw Permalink Normal View History

"""Async SQLAlchemy database engine, session factory, and base model."""
from sqlalchemy.ext.asyncio import (
AsyncSession,
async_sessionmaker,
create_async_engine,
)
from sqlalchemy.orm import DeclarativeBase
from app.config import settings
class Base(DeclarativeBase):
"""Declarative base for all SQLAlchemy models."""
pass
engine = create_async_engine(
settings.DATABASE_URL,
echo=(settings.APP_ENV == "development"),
pool_pre_ping=True,
pool_size=10,
max_overflow=20,
)
async_session_factory = async_sessionmaker(
engine,
class_=AsyncSession,
expire_on_commit=False,
)
async def get_db() -> AsyncSession:
"""FastAPI dependency that yields an async DB session."""
async with async_session_factory() as session:
try:
yield session
await session.commit()
except Exception:
await session.rollback()
raise
finally:
await session.close()
async def init_db():
"""Create all tables. Used for testing and first-run setup."""
async with engine.begin() as conn:
await conn.run_sync(Base.metadata.create_all)
async def drop_db():
"""Drop all tables. Used in test teardown."""
async with engine.begin() as conn:
await conn.run_sync(Base.metadata.drop_all)