51 lines
1.0 KiB
Python
51 lines
1.0 KiB
Python
"""SQLModel 数据库配置"""
|
|
from typing import Generator
|
|
from sqlmodel import Session, create_engine
|
|
from sqlalchemy.orm import sessionmaker
|
|
from .config import settings
|
|
|
|
|
|
# 创建数据库引擎
|
|
engine = create_engine(
|
|
settings.DATABASE_URL,
|
|
echo=settings.DEBUG,
|
|
pool_pre_ping=True,
|
|
pool_size=10,
|
|
max_overflow=20,
|
|
)
|
|
|
|
# 创建 SessionLocal
|
|
SessionLocal = sessionmaker(
|
|
autocommit=False,
|
|
autoflush=False,
|
|
bind=engine,
|
|
class_=Session,
|
|
)
|
|
|
|
|
|
def get_session() -> Generator[Session, None, None]:
|
|
"""
|
|
获取数据库会话（依赖注入）
|
|
"""
|
|
with SessionLocal() as session:
|
|
yield session
|
|
|
|
|
|
def init_db() -> None:
|
|
"""初始化数据库（创建所有表）"""
|
|
from sqlmodel import SQLModel
|
|
from ..models.job import Job, Step, JobLog # noqa: F401
|
|
|
|
SQLModel.metadata.create_all(engine)
|
|
print("✓ Database initialized")
|
|
|
|
|
|
def drop_db() -> None:
|
|
"""删除所有表（开发用）"""
|
|
from sqlmodel import SQLModel
|
|
|
|
SQLModel.metadata.drop_all(engine)
|
|
print("✓ Database dropped")
|
|
|
|
|