File size: 2,138 Bytes
bc1cd44
 
1ef4e10
 
 
f7ef7d3
 
86cf81a
f7ef7d3
 
 
 
 
 
 
1ef4e10
 
 
0fb23b8
1ef4e10
 
 
 
aad2e16
 
 
 
 
 
 
 
 
 
 
 
 
 
1ef4e10
 
 
 
 
 
 
 
 
 
86cf81a
 
 
 
 
 
17ebb8a
 
86cf81a
 
1ef4e10
 
e24b889
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1ef4e10
 
86cf81a
 
1ef4e10
86cf81a
 
 
 
17ebb8a
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
# 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