CI Test
50 lines
1.2 KiB
Python
50 lines
1.2 KiB
Python
"""Database session management and engine configuration.
|
|
|
|
统一使用 app.config 中的数据库配置，移除重复的 DatabaseSettings。
|
|
"""
|
|
|
|
from contextlib import contextmanager
|
|
from typing import Generator
|
|
|
|
from app.config import settings
|
|
from sqlalchemy import create_engine
|
|
from sqlalchemy.orm import Session, sessionmaker
|
|
|
|
engine = create_engine(
|
|
settings.database_url,
|
|
pool_size=settings.DATABASE_POOL_SIZE,
|
|
max_overflow=settings.DATABASE_MAX_OVERFLOW,
|
|
pool_timeout=settings.DATABASE_POOL_TIMEOUT,
|
|
pool_recycle=settings.DATABASE_POOL_RECYCLE,
|
|
)
|
|
|
|
SessionLocal = sessionmaker(autocommit=False, autoflush=False, bind=engine)
|
|
|
|
|
|
def get_db() -> Generator[Session, None, None]:
|
|
"""Dependency for getting database sessions."""
|
|
db = SessionLocal()
|
|
try:
|
|
yield db
|
|
finally:
|
|
db.close()
|
|
|
|
|
|
@contextmanager
|
|
def get_db_context() -> Generator[Session, None, None]:
|
|
"""Context manager for database sessions.
|
|
|
|
Usage:
|
|
with get_db_context() as db:
|
|
db.query(Model).all()
|
|
"""
|
|
db = SessionLocal()
|
|
try:
|
|
yield db
|
|
db.commit()
|
|
except Exception:
|
|
db.rollback()
|
|
raise
|
|
finally:
|
|
db.close()
|