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