backend/db/session.py

"""
Async Database Session Management

Provides async SQLAlchemy engine and session management.
"""

import os
from contextlib import asynccontextmanager
from typing import AsyncGenerator, Optional

from sqlalchemy.ext.asyncio import AsyncEngine, AsyncSession, create_async_engine
from sqlalchemy.orm import sessionmaker
from sqlalchemy.pool import NullPool

from backend.core.config import settings
from backend.core.exceptions import DatabaseError
from backend.logging.logger import get_logger
from sqlalchemy import text


# Initialize logger
logger = get_logger("database_session", component="db")


# Global engine instance
_engine: Optional[AsyncEngine] = None
_session_factory: Optional[sessionmaker] = None


def _is_production() -> bool:
    """Check if running in production mode."""
    return os.getenv("AEGISLM_ENV", "development").lower() == "production"


def create_database_engine() -> AsyncEngine:
    """
    Create async database engine.
    
    Uses asyncpg for PostgreSQL with connection pooling.
    SQLite is NOT allowed in production - PostgreSQL is required.
    """
    global _engine
    
    if _engine is not None:
        return _engine
    
    is_prod = _is_production()
    
    try:
        # Determine if using SQLite
        is_sqlite = settings.database_url.startswith("sqlite")
        
        # PRODUCTION: SQLite is NOT allowed
        if is_sqlite and is_prod:
            raise DatabaseError(
                "CRITICAL: SQLite is not allowed in production. "
                "Please use PostgreSQL with postgresql+asyncpg:// driver.",
                details={"database_url": "SQLite URLs are blocked in production"}
            )
        
        # DEVELOPMENT: Warn about SQLite
        if is_sqlite:
            logger.warning(
                "SQLite is being used. This is only suitable for development. "
                "For production, use PostgreSQL with postgresql+asyncpg://"
            )
        
        # Fix SQLite URL for async compatibility (development only)
        # sqlite:// -> sqlite+aiosqlite://
        db_url = settings.database_url
        if is_sqlite and "+aiosqlite" not in db_url:
            db_url = db_url.replace("sqlite://", "sqlite+aiosqlite://")
        
        # Create async engine with appropriate pool settings
        # Note: When using NullPool, we can't use pool_size/max_overflow/pool_pre_ping
        engine = create_async_engine(
            db_url,
            echo=False,  # Set to True for SQL debugging
            poolclass=NullPool,  # Use NullPool for asyncpg compatibility
        )
        
        # Enable SQLite foreign keys - defer to first use instead of at init
        # This avoids event loop issues with TestClient
        if is_sqlite:
            # Store engine first, FK will be enabled lazily on first connection
            _engine = engine
            logger.warning("SQLite database engine created (development mode only)")
            return _engine
        
        # PostgreSQL engine created
        _engine = engine
        logger.info(f"PostgreSQL database engine created: {settings.database_url.split('@')[-1]}")
        return _engine
        
    except DatabaseError:
        # Re-raise DatabaseError as-is
        raise
    except Exception as e:
        raise DatabaseError(
            f"Failed to create database engine: {str(e)}",
            details={"database_url": settings.database_url.split("@")[-1] if "@" in settings.database_url else "unknown"}
        )


def get_session_factory() -> sessionmaker:
    """
    Get or create session factory.
    
    Returns:
        SQLAlchemy async session maker
    """
    global _session_factory, _engine
    
    if _engine is None:
        _engine = create_database_engine()
    
    if _session_factory is None:
        _session_factory = sessionmaker(
            _engine,
            class_=AsyncSession,
            expire_on_commit=False,
        )
    
    return _session_factory


async def get_db_session() -> AsyncGenerator[AsyncSession, None]:
    """
    Dependency for getting async database session.
    
    Yields:
        Async SQLAlchemy session
        
    Usage:
        @app.get("/items")
        async def get_items(db: AsyncSession = Depends(get_db_session)):
            ...
    """
    factory = get_session_factory()
    
    async with factory() as session:
        try:
            yield session
            await session.commit()
        except Exception:
            await session.rollback()
            raise
        finally:
            await session.close()


@asynccontextmanager
async def get_db_context() -> AsyncGenerator[AsyncSession, None]:
    """
    Context manager for database sessions.
    
    Usage:
        async with get_db_context() as session:
            result = await session.execute(...)
    """
    factory = get_session_factory()
    
    async with factory() as session:
        try:
            yield session
            await session.commit()
        except Exception:
            await session.rollback()
            raise
        finally:
            await session.close()


async def init_database() -> None:
    """
    Initialize database tables.
    
    Creates all tables if they don't exist.
    For production, use Alembic migrations instead.
    """
    from backend.db.models import Base
    
    engine = create_database_engine()
    
    async with engine.begin() as conn:
        # Create all tables
        await conn.run_sync(Base.metadata.create_all)
    
    # Dispose engine (will be recreated on next use)
    await engine.dispose()


async def close_database() -> None:
    """
    Close database connections and dispose engine.
    
    Call this on application shutdown.
    """
    global _engine, _session_factory
    
    if _engine is not None:
        await _engine.dispose()
        _engine = None
    
    _session_factory = None


async def check_database_connection() -> bool:
    """
    Check database connectivity.
    
    Returns:
        True if connection successful, False otherwise
    """
    try:
        engine = create_database_engine()
        async with engine.connect() as conn:
            await conn.execute(text("SELECT 1"))
        return True
    except Exception:
        return False