utils/database.py

# utils/database.py

from sqlalchemy import create_engine
from sqlalchemy.orm import sessionmaker
from contextlib import contextmanager
import sys  # Add sys import
import os  # Add os import

# Add project root to Python path to ensure local modules are prioritized
SCRIPT_DIR = os.path.dirname(os.path.abspath(__file__))
PROJECT_ROOT = os.path.dirname(SCRIPT_DIR)
if PROJECT_ROOT not in sys.path:
    sys.path.insert(0, PROJECT_ROOT)

from config import conf
from utils.logger import Logger

log = Logger()


def get_db_engine():
    """Create DB engine with error handling for HF Spaces"""
    try:
        engine = create_engine(
            conf.db_url,
            pool_pre_ping=True,
            pool_size=5,
            max_overflow=10,
            connect_args={"connect_timeout": 10},
        )
        log.info("Database engine created successfully")
        return engine
    except Exception as e:
        log.error(f"Failed to create database engine: {e}")


engine = get_db_engine()
SessionLocal = sessionmaker(bind=engine)


@contextmanager
def get_db():
    """Session manager for HF Spaces"""
    db = SessionLocal()
    try:
        yield db
        db.commit()
    except Exception as e:
        db.rollback()
        log.error(f"Database error: {e}")
        raise
    finally:
        db.close()


def initialize_database():
    """Initialize tables with HF Spaces compatibility"""
    try:
        from data.models import Base

        Base.metadata.create_all(bind=engine)
        log.info("Tables created successfully")
    except Exception as e:
        log.error(f"Table creation failed: {e}")