Spaces:
Running
Running
File size: 2,101 Bytes
09801ca | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 | """
Database Session Module — Async engine and session factory.
"""
from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker, create_async_engine
from app.core.config import get_settings
def _create_engine():
"""Create the async SQLAlchemy engine from settings."""
settings = get_settings()
url = settings.DATABASE_URL
if url.startswith("postgres://"):
url = url.replace("postgres://", "postgresql+asyncpg://", 1)
elif url.startswith("postgresql://") and not url.startswith("postgresql+asyncpg://"):
url = url.replace("postgresql://", "postgresql+asyncpg://", 1)
return create_async_engine(
url,
pool_size=settings.DB_POOL_SIZE,
max_overflow=settings.DB_MAX_OVERFLOW,
pool_pre_ping=settings.DB_POOL_PRE_PING,
pool_recycle=settings.DB_POOL_RECYCLE,
echo=settings.DB_ECHO,
)
_engine = None
_session_factory = None
def get_engine():
"""Get or create the global engine singleton."""
global _engine
if _engine is None:
_engine = _create_engine()
return _engine
def get_session_factory() -> async_sessionmaker[AsyncSession]:
"""Get or create the global session factory."""
global _session_factory
if _session_factory is None:
_session_factory = async_sessionmaker(
bind=get_engine(),
class_=AsyncSession,
expire_on_commit=False,
)
return _session_factory
async def get_db() -> AsyncSession:
"""
FastAPI dependency — yields an async database session.
Usage:
@router.get("/users")
async def list_users(db: AsyncSession = Depends(get_db)):
...
"""
factory = get_session_factory()
async with factory() as session:
try:
yield session
except Exception:
await session.rollback()
raise
finally:
await session.close()
async def dispose_engine():
"""Dispose the engine on shutdown."""
global _engine
if _engine:
await _engine.dispose()
_engine = None
|