38 lines
871 B
Python
38 lines
871 B
Python
from collections.abc import Generator
|
|
|
|
from sqlalchemy import create_engine
|
|
from sqlalchemy.engine import Engine
|
|
from sqlalchemy.orm import DeclarativeBase, Session, sessionmaker
|
|
|
|
from app.core.config import get_settings
|
|
|
|
|
|
class Base(DeclarativeBase):
|
|
pass
|
|
|
|
|
|
_engine: Engine | None = None
|
|
|
|
|
|
def get_engine() -> Engine:
|
|
global _engine
|
|
url = get_settings().database_url.strip()
|
|
if not url:
|
|
raise RuntimeError("数据库连接未配置")
|
|
if _engine is None:
|
|
_engine = create_engine(url, pool_pre_ping=True)
|
|
SessionLocal.configure(bind=_engine)
|
|
return _engine
|
|
|
|
|
|
SessionLocal = sessionmaker(autocommit=False, autoflush=False, expire_on_commit=False)
|
|
|
|
|
|
def get_db() -> Generator[Session, None, None]:
|
|
SessionLocal.configure(bind=get_engine())
|
|
db = SessionLocal()
|
|
try:
|
|
yield db
|
|
finally:
|
|
db.close()
|