xxx

src/shared/database/database.py

@@ -13,6 +13,24 @@ from shared.utils.logger import get_logger
 
 logger = get_logger(__name__)
 
+
+def _get_pool_config(role: str = "web") -> dict:
+    """按角色返回连接池参数。
+
+    web 入口：适中并发；celery worker：少量长连接。
+    通过 DB_POOL_SIZE / DB_MAX_OVERFLOW 环境变量可覆盖默认值。
+    """
+    defaults = {
+        "web":    {"pool_size": 10, "max_overflow": 20},
+        "celery": {"pool_size": 5,  "max_overflow": 10},
+    }
+    role_cfg = defaults.get(role, defaults["web"])
+    # 允许环境变量覆盖
+    pool_size = int(os.getenv("DB_POOL_SIZE", str(role_cfg["pool_size"])))
+    max_overflow = int(os.getenv("DB_MAX_OVERFLOW", str(role_cfg["max_overflow"])))
+    return {"pool_size": pool_size, "max_overflow": max_overflow}
+
+
 class DatabaseManager:
     """数据库管理器"""
     
@@ -21,21 +39,27 @@ class DatabaseManager:
         self.async_session = None
         self.is_connected = False
     
-    async def connect(self):
-        """连接数据库"""
+    async def connect(self, role: str = "web"):
+        """连接数据库
+
+        Args:
+            role: 连接角色，"web" 或 "celery"，决定连接池大小
+        """
         if not settings.DATABASE_URL:
             logger.warning("未配置数据库连接，跳过数据库初始化")
             self.is_connected = False
             return
             
         try:
+            pool_cfg = _get_pool_config(role)
             # 创建异步引擎
             self.engine = create_async_engine(
                 settings.DATABASE_URL,
                 echo=settings.DEBUG,
-                pool_size=20,
-                max_overflow=30,
-                pool_recycle=3600
+                pool_size=pool_cfg["pool_size"],
+                max_overflow=pool_cfg["max_overflow"],
+                pool_recycle=3600,
+                pool_pre_ping=True,  # 自动检测失效连接，避免 PG 断连报错
             )
             
             # 创建异步会话工厂
@@ -50,7 +74,10 @@ class DatabaseManager:
                 await conn.execute(text("SELECT 1"))
             
             self.is_connected = True
-            logger.info("数据库连接成功")
+            logger.info(
+                "数据库连接成功 (pool_size=%d, max_overflow=%d)",
+                pool_cfg["pool_size"], pool_cfg["max_overflow"],
+            )
             
         except Exception as e:
             logger.error(f"数据库连接失败: {e}")
@@ -66,18 +93,22 @@ class DatabaseManager:
     
     @asynccontextmanager
     async def session(self):
-        """获取数据库会话的异步上下文管理器"""
+        """获取数据库会话的异步上下文管理器（用于后台任务/Celery）"""
         if not self.is_connected:
             await self.connect()
         
         session = self.async_session()
         try:
             yield session
+            await session.commit()
+        except Exception:
+            await session.rollback()
+            raise
         finally:
             await session.close()
     
     async def get_session(self) -> AsyncSession:
-        """获取数据库会话"""
+        """获取数据库会话（非上下文管理器，配合 get_db_session 依赖使用）"""
         if not self.is_connected:
             await self.connect()
         
@@ -100,9 +131,20 @@ db_manager = DatabaseManager()
 
 # 数据库依赖注入
 async def get_db_session():
-    """获取数据库会话的依赖函数"""
+    """获取数据库会话的依赖函数
+
+    统一事务边界：
+    - 路由正常返回 → 自动 commit
+    - 路由抛出异常 → 自动 rollback
+
+    路由中应使用 flush() 代替 commit()，以便在提交前仍能 refresh()。
+    """
     session = await db_manager.get_session()
     try:
         yield session
+        await session.commit()
+    except Exception:
+        await session.rollback()
+        raise
     finally:
-        await session.close()
+        await session.close()