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
import sentry_sdk

# 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}")
    #     sentry_sdk.capture_exception(e)
    #     raise


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}")
    #     sentry_sdk.capture_exception(e)  # Add Sentry logging
    #     raise  # Re-raise for debugging
    # finally:
    db.close()


@contextmanager
def get_db_readonly():
    """Read-only session manager that doesn't auto-commit"""
    db = SessionLocal()
    try:
        yield db
    except Exception as e:
        db.rollback()
        log.error(f"Database error: {e}")
        sentry_sdk.capture_exception(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}")
    #     sentry_sdk.capture_exception(e)
    #     raise