62 lines
1.9 KiB
Python
62 lines
1.9 KiB
Python
from fastapi import Depends
|
||
from typing import Annotated
|
||
from sqlmodel import Session, create_engine
|
||
from core.config.settings import settings
|
||
from sqlalchemy.orm import sessionmaker
|
||
|
||
# 1.引擎配置
|
||
engine = create_engine(
|
||
settings.mysql_url, # 数据库连接url
|
||
echo=settings.DATABASE_ECHO, # echo=True 打印生成的SQL语句，开发调试打开，生产关闭
|
||
pool_recycle = 3600, # 1小时回收连接，解决mysql 8小时断开问题
|
||
pool_pre_ping=True, # 连接前检查
|
||
pool_size=settings.DATABASE_POOL_SIZE, # 连接池大小
|
||
max_overflow=settings.DATABASE_MAX_OVERFLOW, # 最大溢出连接
|
||
)
|
||
|
||
# 2. 创建 sessionmaker（工厂）
|
||
SessionLocal = sessionmaker(autocommit=False, autoflush=False, bind=engine)
|
||
|
||
# 3.依赖注入函数
|
||
def get_db():
|
||
"""FastAPI 依赖注入用"""
|
||
db = SessionLocal() # 使用工厂创建 Session
|
||
try:
|
||
yield db
|
||
db.commit() # 无异常则提交
|
||
except Exception:
|
||
db.rollback() # 异常则回滚
|
||
raise
|
||
finally:
|
||
db.close() # 确保关闭
|
||
|
||
# 4. 定义依赖类型
|
||
SessionDep = Annotated[Session, Depends(get_db)]
|
||
|
||
# 5. 供定时任务使用的上下文管理器
|
||
from contextlib import contextmanager
|
||
|
||
@contextmanager
|
||
def get_db_session():
|
||
"""定时任务、脚本用"""
|
||
db = SessionLocal()
|
||
try:
|
||
yield db
|
||
db.commit()
|
||
except Exception:
|
||
db.rollback()
|
||
raise
|
||
finally:
|
||
db.close()
|
||
|
||
# 使用方式
|
||
# def crawl_task():
|
||
# data = fetch_lottery_data()
|
||
#
|
||
# with get_db_session() as db:
|
||
# # 检查是否存在
|
||
# existing = db.query(LotteryResult).filter_by(draw_no=data['draw_no']).first()
|
||
# if not existing:
|
||
# record = LotteryResult(**data)
|
||
# db.add(record)
|
||
# # with 块结束时自动 commit |