67 lines
1.8 KiB
Python
67 lines
1.8 KiB
Python
"""数据库引擎 / 会话 / 依赖注入 / 自动建库。"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import logging
|
|
from collections.abc import Generator
|
|
|
|
from sqlalchemy import create_engine, text
|
|
from sqlalchemy.engine import Engine
|
|
from sqlalchemy.orm import Session, sessionmaker
|
|
|
|
from app.core.config import settings
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
def _build_engine() -> Engine:
|
|
if settings.DB_DRIVER.startswith("sqlite"):
|
|
return create_engine(
|
|
settings.database_url,
|
|
echo=settings.DB_ECHO,
|
|
connect_args={"check_same_thread": False},
|
|
)
|
|
return create_engine(
|
|
settings.database_url,
|
|
echo=settings.DB_ECHO,
|
|
pool_size=settings.DB_POOL_SIZE,
|
|
pool_recycle=settings.DB_POOL_RECYCLE,
|
|
pool_pre_ping=True,
|
|
future=True,
|
|
)
|
|
|
|
|
|
engine: Engine = _build_engine()
|
|
SessionLocal = sessionmaker(bind=engine, autocommit=False, autoflush=False, class_=Session)
|
|
|
|
|
|
def ensure_database_exists() -> None:
|
|
"""MySQL 下若目标库不存在则自动创建，省去手工建库。"""
|
|
if settings.DB_DRIVER.startswith("sqlite"):
|
|
return
|
|
root_engine = create_engine(settings.server_url_without_db, future=True)
|
|
try:
|
|
with root_engine.connect() as conn:
|
|
conn.execute(
|
|
text(
|
|
f"CREATE DATABASE IF NOT EXISTS `{settings.DB_NAME}` "
|
|
f"DEFAULT CHARACTER SET utf8mb4 COLLATE utf8mb4_general_ci"
|
|
)
|
|
)
|
|
conn.commit()
|
|
finally:
|
|
root_engine.dispose()
|
|
|
|
|
|
def get_db() -> Generator[Session, None, None]:
|
|
"""FastAPI 依赖：每个请求一个会话，结束即关闭。"""
|
|
db = SessionLocal()
|
|
try:
|
|
yield db
|
|
finally:
|
|
db.close()
|
|
|
|
|
|
def dialect_name() -> str:
|
|
return engine.dialect.name
|