enclava/backend/app/db/database.py

"""
Database connection and session management
"""

import logging
from typing import AsyncGenerator
from sqlalchemy import create_engine, MetaData
from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker, create_async_engine
from sqlalchemy.orm import sessionmaker, declarative_base
from sqlalchemy.pool import StaticPool

from app.core.config import settings

logger = logging.getLogger(__name__)

# Create async engine with optimized connection pooling
engine = create_async_engine(
    settings.DATABASE_URL.replace("postgresql://", "postgresql+asyncpg://"),
    echo=settings.APP_DEBUG,
    future=True,
    pool_pre_ping=True,
    pool_size=50,          # Increased from 20 for better concurrency
    max_overflow=100,      # Increased from 30 for burst capacity  
    pool_recycle=3600,     # Recycle connections every hour
    pool_timeout=30,       # Max time to get connection from pool
    connect_args={
        "command_timeout": 5,
        "server_settings": {
            "application_name": "enclava_backend",
        },
    },
)

# Create async session factory
async_session_factory = async_sessionmaker(
    bind=engine,
    class_=AsyncSession,
    expire_on_commit=False,
)

# Create synchronous engine and session for budget enforcement (optimized)
sync_engine = create_engine(
    settings.DATABASE_URL,
    echo=settings.APP_DEBUG,
    future=True,
    pool_pre_ping=True,
    pool_size=25,          # Increased from 10 for better performance
    max_overflow=50,       # Increased from 20 for burst capacity
    pool_recycle=3600,     # Recycle connections every hour
    pool_timeout=30,       # Max time to get connection from pool
    connect_args={
        "application_name": "enclava_backend_sync",
    },
)

# Create sync session factory
SessionLocal = sessionmaker(
    bind=sync_engine,
    expire_on_commit=False,
)

# Create base class for models
Base = declarative_base()

# Metadata for migrations
metadata = MetaData()


async def get_db() -> AsyncGenerator[AsyncSession, None]:
    """Get database session"""
    async with async_session_factory() as session:
        try:
            yield session
        except Exception as e:
            logger.error(f"Database session error: {e}")
            await session.rollback()
            raise
        finally:
            await session.close()


async def init_db():
    """Initialize database"""
    try:
        async with engine.begin() as conn:
            # Import all models to ensure they're registered
            from app.models.user import User
            from app.models.api_key import APIKey
            from app.models.usage_tracking import UsageTracking
            # Import additional models - these are available
            try:
                from app.models.budget import Budget
            except ImportError:
                logger.warning("Budget model not available yet")
            
            try:
                from app.models.audit_log import AuditLog
            except ImportError:
                logger.warning("AuditLog model not available yet")
            
            try:
                from app.models.module import Module
            except ImportError:
                logger.warning("Module model not available yet")
            
            # Tables are now created via migration container - no need to create here
            # await conn.run_sync(Base.metadata.create_all)  # DISABLED - migrations handle this
            
        # Create default admin user if no admin exists
        await create_default_admin()
            
        logger.info("Database initialized successfully")
    except Exception as e:
        logger.error(f"Failed to initialize database: {e}")
        raise


async def create_default_admin():
    """Create default admin user if none exists"""
    from app.models.user import User
    from app.core.security import get_password_hash
    from app.core.config import settings
    from sqlalchemy import select
    
    try:
        async with async_session_factory() as session:
            # Check if any admin user exists
            stmt = select(User).where(User.role == "super_admin")
            result = await session.execute(stmt)
            existing_admin = result.scalar_one_or_none()
            
            if existing_admin:
                logger.info("Admin user already exists")
                return
            
            # Create default admin user from environment variables
            admin_username = settings.ADMIN_USER
            admin_password = settings.ADMIN_PASSWORD
            admin_email = settings.ADMIN_EMAIL or f"{admin_username}@example.com"
            
            admin_user = User.create_default_admin(
                email=admin_email,
                username=admin_username,
                password_hash=get_password_hash(admin_password)
            )
            
            session.add(admin_user)
            await session.commit()
            
            logger.warning("=" * 60)
            logger.warning("DEFAULT ADMIN USER CREATED")
            logger.warning(f"Email: {admin_email}")
            logger.warning(f"Username: {admin_username}")
            logger.warning("Password: [Set via ADMIN_PASSWORD environment variable]")
            logger.warning("PLEASE CHANGE THE PASSWORD AFTER FIRST LOGIN")
            logger.warning("=" * 60)
            
    except Exception as e:
        logger.error(f"Failed to create default admin user: {e}")
        # Don't raise here as this shouldn't block the application startup