初始化

core/db.py

@@ -0,0 +1,62 @@
+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