124 lines
3.4 KiB
Python
124 lines
3.4 KiB
Python
"""WealthWise database configuration and session management.
|
|
|
|
This module provides:
|
|
- Async SQLAlchemy engine with connection pooling
|
|
- Async session factory for database operations
|
|
- FastAPI dependency for session injection
|
|
- Connection health checking utilities
|
|
"""
|
|
|
|
from typing import AsyncGenerator
|
|
|
|
from sqlalchemy.exc import SQLAlchemyError
|
|
from sqlalchemy.ext.asyncio import (
|
|
AsyncEngine,
|
|
AsyncSession,
|
|
async_sessionmaker,
|
|
create_async_engine,
|
|
)
|
|
from sqlalchemy.pool import NullPool
|
|
from sqlmodel.ext.asyncio.session import AsyncSession as SQLModelAsyncSession
|
|
|
|
from app.core.config import get_settings
|
|
|
|
settings = get_settings()
|
|
|
|
|
|
# Create async engine with connection pooling
|
|
# Supabase Transaction Pooler (port 6543) supports high concurrency
|
|
engine: AsyncEngine = create_async_engine(
|
|
settings.DATABASE_URL,
|
|
echo=settings.DB_ECHO,
|
|
pool_size=settings.DB_POOL_SIZE,
|
|
max_overflow=settings.DB_MAX_OVERFLOW,
|
|
pool_pre_ping=settings.DB_POOL_PRE_PING,
|
|
# Connection pool settings for production
|
|
pool_recycle=3600, # Recycle connections after 1 hour
|
|
pool_timeout=30, # Wait up to 30 seconds for available connection
|
|
# Async-specific settings
|
|
future=True,
|
|
)
|
|
|
|
# Create async session factory
|
|
# expire_on_commit=False allows accessing attributes after session closes
|
|
AsyncSessionLocal = async_sessionmaker(
|
|
engine,
|
|
class_=SQLModelAsyncSession,
|
|
expire_on_commit=False,
|
|
autoflush=False,
|
|
autocommit=False,
|
|
)
|
|
|
|
|
|
async def get_session() -> AsyncGenerator[SQLModelAsyncSession, None]:
|
|
"""FastAPI dependency that provides an async database session.
|
|
|
|
This generator yields a database session for use in API endpoints.
|
|
It automatically handles transaction rollback on errors and ensures
|
|
proper session cleanup.
|
|
|
|
Usage:
|
|
@app.get("/items")
|
|
async def get_items(session: AsyncSession = Depends(get_session)):
|
|
result = await session.execute(select(Item))
|
|
return result.scalars().all()
|
|
|
|
Yields:
|
|
AsyncSession: Database session instance
|
|
|
|
Raises:
|
|
SQLAlchemyError: Re-raised after rollback if database error occurs
|
|
"""
|
|
session: SQLModelAsyncSession = AsyncSessionLocal()
|
|
try:
|
|
yield session
|
|
await session.commit()
|
|
except SQLAlchemyError as e:
|
|
await session.rollback()
|
|
raise e
|
|
finally:
|
|
await session.close()
|
|
|
|
|
|
async def close_engine() -> None:
|
|
"""Close all database connections.
|
|
|
|
Call this during application shutdown to properly release
|
|
all database connections in the pool.
|
|
"""
|
|
await engine.dispose()
|
|
|
|
|
|
async def check_db_connection() -> bool:
|
|
"""Verify database connectivity by executing a simple query.
|
|
|
|
Returns:
|
|
True if database is accessible, False otherwise
|
|
"""
|
|
try:
|
|
async with AsyncSessionLocal() as session:
|
|
result = await session.execute("SELECT 1")
|
|
return result.scalar() == 1
|
|
except Exception:
|
|
return False
|
|
|
|
|
|
# For Alembic migrations (sync operations)
|
|
def create_sync_engine():
|
|
"""Create synchronous engine for Alembic migrations.
|
|
|
|
Returns:
|
|
SyncEngine: Synchronous SQLAlchemy engine
|
|
"""
|
|
from sqlalchemy import create_engine
|
|
|
|
# Convert async URL to sync URL
|
|
sync_url = settings.DATABASE_URL.replace(
|
|
"postgresql+asyncpg://", "postgresql://"
|
|
)
|
|
|
|
return create_engine(
|
|
sync_url,
|
|
echo=settings.DB_ECHO,
|
|
)
|