Litellm-proxy / app /database.py
NitinBot001's picture
Update app/database.py
87047a5 verified
from sqlalchemy import create_engine
from sqlalchemy.orm import declarative_base, sessionmaker
import os
# HuggingFace pe /data exist nahi karta — fallback current dir pe
_default_db = "sqlite:///./proxy.db"
DATABASE_URL = os.getenv("DATABASE_URL", _default_db)
# Agar SQLite hai aur path /data se start hota hai to directory banao
if "sqlite" in DATABASE_URL:
db_path = DATABASE_URL.replace("sqlite:///", "").replace("sqlite:////", "/")
db_dir = os.path.dirname(db_path)
if db_dir and db_dir != "." and not os.path.exists(db_dir):
try:
os.makedirs(db_dir, exist_ok=True)
except Exception:
# fallback: current directory mein banao
DATABASE_URL = _default_db
engine = create_engine(
DATABASE_URL,
connect_args={"check_same_thread": False} if "sqlite" in DATABASE_URL else {},
)
SessionLocal = sessionmaker(autocommit=False, autoflush=False, bind=engine)
Base = declarative_base()
def get_db():
db = SessionLocal()
try:
yield db
finally:
db.close()