50 lines
1.4 KiB
Python
50 lines
1.4 KiB
Python
from sqlalchemy import create_engine
|
|
from sqlalchemy.orm import declarative_base, sessionmaker
|
|
from sqlalchemy.exc import OperationalError
|
|
from fastapi import HTTPException, status
|
|
|
|
from .config import settings
|
|
|
|
DATABASE_URL = settings.DATABASE_URL
|
|
|
|
connect_args = {}
|
|
if DATABASE_URL.startswith("mysql+pymysql://"):
|
|
connect_args["connect_timeout"] = 5
|
|
|
|
engine = create_engine(
|
|
DATABASE_URL,
|
|
pool_pre_ping=True,
|
|
pool_recycle=3600,
|
|
pool_size=settings.DB_POOL_SIZE if settings.DB_POOL_SIZE else 5,
|
|
max_overflow=settings.DB_POOL_MAX_OVERFLOW if settings.DB_POOL_MAX_OVERFLOW else 10,
|
|
pool_timeout=30,
|
|
future=True,
|
|
connect_args=connect_args,
|
|
)
|
|
SessionLocal = sessionmaker(bind=engine, autocommit=False, autoflush=False, future=True)
|
|
Base = declarative_base()
|
|
|
|
|
|
def get_db():
|
|
db = SessionLocal()
|
|
try:
|
|
yield db
|
|
db.commit()
|
|
except OperationalError as exc:
|
|
import logging
|
|
logger = logging.getLogger("db")
|
|
logger.error(f"Database OperationalError: {exc}", exc_info=True)
|
|
db.rollback()
|
|
raise HTTPException(
|
|
status_code=status.HTTP_503_SERVICE_UNAVAILABLE,
|
|
detail=f"Database unavailable: {str(exc)}",
|
|
) from exc
|
|
except Exception as exc:
|
|
import logging
|
|
logger = logging.getLogger("db")
|
|
logger.error(f"Database Exception: {exc}", exc_info=True)
|
|
db.rollback()
|
|
raise
|
|
finally:
|
|
db.close()
|