WealthWise/backend/app/core/db.py

"""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,
    )