From d1d98a90819a5f387a51b6d1fce920efedfac768 Mon Sep 17 00:00:00 2001
From: InfinityPacer <160988576+InfinityPacer@users.noreply.github.com>
Date: Sun, 20 Oct 2024 03:10:08 +0800
Subject: [PATCH] feat(db): add pool class configuration and adjust connection
 settings

---
 app/core/config.py | 22 ++++++++++++----------
 app/db/__init__.py | 38 +++++++++++++++++++++++---------------
 config/app.env     |  6 +++---
 3 files changed, 38 insertions(+), 28 deletions(-)

diff --git a/app/core/config.py b/app/core/config.py
index e916af31b..86ceb8979 100644
--- a/app/core/config.py
+++ b/app/core/config.py
@@ -54,16 +54,18 @@ class Config:
     DEV: bool = False
     # 是否在控制台输出 SQL 语句，默认关闭
     DB_ECHO: bool = False
-    # 是否在获取连接时进行预先 ping 操作，默认开启
-    DB_POOL_PRE_PING: bool = True
-    # 数据库连接池的大小，默认 100
-    DB_POOL_SIZE: int = 100
-    # 数据库连接的回收时间（秒），默认 3600 秒（1 小时）
-    DB_POOL_RECYCLE: int = 3600
-    # 数据库连接池获取连接的超时时间（秒），默认 180 秒
-    DB_POOL_TIMEOUT: int = 180
-    # 数据库连接池最大溢出连接数，默认 5
-    DB_MAX_OVERFLOW: int = 5
+    # 数据库连接池类型，QueuePool, NullPool
+    DB_POOL_TYPE: str = "QueuePool"
+    # 是否在获取连接时进行预先 ping 操作，默认关闭
+    DB_POOL_PRE_PING: bool = False
+    # 数据库连接池的大小，默认 20
+    DB_POOL_SIZE: int = 20
+    # 数据库连接的回收时间（秒），默认 1800 秒
+    DB_POOL_RECYCLE: int = 1800
+    # 数据库连接池获取连接的超时时间（秒），默认 30 秒
+    DB_POOL_TIMEOUT: int = 30
+    # 数据库连接池最大溢出连接数，默认 10
+    DB_MAX_OVERFLOW: int = 10
     # SQLite 的 busy_timeout 参数，默认为 60 秒
     DB_TIMEOUT: int = 60
     # 配置文件目录
diff --git a/app/db/__init__.py b/app/db/__init__.py
index 0fdb01954..cf7ab9acb 100644
--- a/app/db/__init__.py
+++ b/app/db/__init__.py
@@ -1,24 +1,32 @@
-from typing import Any, Self, List, Tuple, Optional, Generator
+from typing import Any, Generator, List, Optional, Self, Tuple
 
-from sqlalchemy import create_engine, QueuePool, and_, inspect
-from sqlalchemy.orm import declared_attr, sessionmaker, Session, scoped_session, as_declarative
+from sqlalchemy import NullPool, QueuePool, and_, create_engine, inspect
+from sqlalchemy.orm import Session, as_declarative, declared_attr, scoped_session, sessionmaker
 
 from app.core.config import settings
 
-# 数据库引擎
-Engine = create_engine(
-    url=f"sqlite:///{settings.CONFIG_PATH}/user.db",
-    pool_pre_ping=settings.DB_POOL_PRE_PING,
-    echo=settings.DB_ECHO,
-    poolclass=QueuePool,
-    pool_size=settings.DB_POOL_SIZE,
-    pool_recycle=settings.DB_POOL_RECYCLE,
-    pool_timeout=settings.DB_POOL_TIMEOUT,
-    max_overflow=settings.DB_MAX_OVERFLOW,
-    connect_args={
+# 根据池类型设置 poolclass 和相关参数
+pool_class = NullPool if settings.DB_POOL_TYPE == "NullPool" else QueuePool
+kwargs = {
+    "url": f"sqlite:///{settings.CONFIG_PATH}/user.db",
+    "pool_pre_ping": settings.DB_POOL_PRE_PING,
+    "echo": settings.DB_ECHO,
+    "poolclass": pool_class,
+    "pool_recycle": settings.DB_POOL_RECYCLE,
+    "connect_args": {
+        # "check_same_thread": False,
         "timeout": settings.DB_TIMEOUT
     }
-)
+}
+# 当使用 QueuePool 时，添加 QueuePool 特有的参数
+if pool_class == QueuePool:
+    kwargs.update({
+        "pool_size": settings.DB_POOL_SIZE,
+        "pool_timeout": settings.DB_POOL_TIMEOUT,
+        "max_overflow": settings.DB_MAX_OVERFLOW
+    })
+# 创建数据库引擎
+Engine = create_engine(**kwargs)
 
 # 会话工厂
 SessionFactory = sessionmaker(bind=Engine)
diff --git a/config/app.env b/config/app.env
index 1352e3064..8c34a1d6e 100644
--- a/config/app.env
+++ b/config/app.env
@@ -9,10 +9,10 @@ DEBUG=false
 DEV=false
 # 日志级别（DEBUG、INFO、WARNING、ERROR等），当DEBUG=true时，此配置项将被忽略，日志级别始终为DEBUG
 LOG_LEVEL=INFO
-# 数据库连接池的大小，可适当降低如10-50以减少I/O压力
-DB_POOL_SIZE=100
+# 数据库连接池的大小，可适当降低如5以减少I/O压力
+DB_POOL_SIZE=20
 # 数据库连接池最大溢出连接数，可适当降低如0以减少I/O压力
-DB_MAX_OVERFLOW=5
+DB_MAX_OVERFLOW=10
 # SQLite 的 busy_timeout 参数，可适当增加如180以减少锁定错误
 DB_TIMEOUT=60
 # 【*】超级管理员，设置后一但重启将固化到数据库中，修改将无效（初始化超级管理员密码仅会生成一次，请在日志中查看并自行登录系统修改）