app/database.py

"""Database connection and session management."""
from sqlalchemy import create_engine
from sqlalchemy.orm import sessionmaker, Session
from app.config import get_settings
from app.models import Base
import logging

logger = logging.getLogger(__name__)
settings = get_settings()

# Create engine with connection pooling
try:
    engine = create_engine(
        settings.DATABASE_URL,
        pool_pre_ping=True,
        pool_size=10,
        max_overflow=20
    )
    SessionLocal = sessionmaker(autocommit=False, autoflush=False, bind=engine)
    logger.info("Database engine created successfully")
except Exception as e:
    logger.warning(f"Failed to create database engine: {e}")
    engine = None
    SessionLocal = None


def init_db():
    """Initialize database tables."""
    if engine is None:
        logger.warning("Database not configured - skipping initialization")
        return
    
    try:
        Base.metadata.create_all(bind=engine)
        logger.info("Database tables created successfully")
        _run_migrations()
    except Exception as e:
        logger.error(f"Failed to initialize database: {e}")
        raise


def _run_migrations():
    """Run any pending column migrations."""
    with engine.connect() as conn:
        # Add file_content column if it doesn't exist
        try:
            conn.execute(
                __import__('sqlalchemy').text(
                    "ALTER TABLE documents ADD COLUMN IF NOT EXISTS file_content BYTEA"
                )
            )
            conn.commit()
            logger.info("Migration: file_content column ensured")
        except Exception as e:
            logger.warning(f"Migration note: {e}")
            conn.rollback()

        # Add status and error_message columns
        for col_sql in [
            "ALTER TABLE documents ADD COLUMN IF NOT EXISTS status VARCHAR DEFAULT 'done'",
            "ALTER TABLE documents ADD COLUMN IF NOT EXISTS error_message TEXT",
            "ALTER TABLE folders ADD COLUMN IF NOT EXISTS parent_id VARCHAR REFERENCES folders(id)",
        ]:
            try:
                conn.execute(__import__('sqlalchemy').text(col_sql))
                conn.commit()
            except Exception as e:
                logger.warning(f"Migration note: {e}")
                conn.rollback()
        logger.info("Migration: status/error_message columns ensured")


def get_db() -> Session:
    """Dependency for getting database session."""
    if SessionLocal is None:
        raise RuntimeError("Database not configured")
    
    db = SessionLocal()
    try:
        yield db
    finally:
        db.close()