From 6996d50a4e067157afa2ea2b273e949fce5777fc Mon Sep 17 00:00:00 2001
From: pushiwuhua <45622023+pushiwuhua7@users.noreply.github.com>
Date: Thu, 4 Jun 2026 18:45:11 +0800
Subject: [PATCH] feat(db): configurable connection pool tuning for MySQL /
 PostgreSQL

Without pool_pre_ping=True, asyncpg leaves cancelled-task connections
in a broken state inside the pool, causing the next checkout to hang
indefinitely under concurrent load. Expose the standard SQLAlchemy
pool knobs (size / max_overflow / recycle / pre_ping / timeout) via
env vars with sensible defaults, applied to non-SQLite engines only.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
---
 .env.example                           |  9 +++++++++
 backend/src/analytics_agent/config.py  |  9 +++++++++
 backend/src/analytics_agent/db/base.py | 20 +++++++++++++++-----
 3 files changed, 33 insertions(+), 5 deletions(-)

diff --git a/.env.example b/.env.example
index 52c80ef..a871a00 100644
--- a/.env.example
+++ b/.env.example
@@ -34,6 +34,15 @@ AWS_REGION=us-west-2
 # SQLite (local dev):
 DATABASE_URL=sqlite+aiosqlite:///./data/dev.db
 
+# Connection pool tuning — applied to MySQL / PostgreSQL only (SQLite ignored).
+# DB_POOL_PRE_PING is required for asyncpg: connections broken by task
+# cancellation otherwise remain pooled and stall the next checkout.
+# DB_POOL_SIZE=10
+# DB_MAX_OVERFLOW=20
+# DB_POOL_RECYCLE=1800
+# DB_POOL_PRE_PING=true
+# DB_POOL_TIMEOUT=10
+
 # Engine config
 ENGINES_CONFIG=./config.yaml
 SQL_ROW_LIMIT=500
diff --git a/backend/src/analytics_agent/config.py b/backend/src/analytics_agent/config.py
index 3fe2d31..a93b783 100644
--- a/backend/src/analytics_agent/config.py
+++ b/backend/src/analytics_agent/config.py
@@ -245,6 +245,15 @@ def get_api_key(self) -> str:
     # Database — defaults to the user config dir; override via DATABASE_URL env var
     database_url: str = f"sqlite+aiosqlite:///{_CONFIG_DIR}/data/agent.db"
 
+    # Connection pool — applied only to non-SQLite engines (MySQL / PostgreSQL).
+    # pool_pre_ping is critical for asyncpg: without it, a connection broken by
+    # task cancellation stays in the pool and the next checkout hangs forever.
+    db_pool_size: int = 10
+    db_max_overflow: int = 20
+    db_pool_recycle: int = 1800
+    db_pool_pre_ping: bool = True
+    db_pool_timeout: int = 10
+
     # Engine config — defaults to the user config dir; override via ENGINES_CONFIG env var
     engines_config: str = str(_CONFIG_DIR / "config.yaml")
     sql_row_limit: int = 500
diff --git a/backend/src/analytics_agent/db/base.py b/backend/src/analytics_agent/db/base.py
index baf038a..866fdf8 100644
--- a/backend/src/analytics_agent/db/base.py
+++ b/backend/src/analytics_agent/db/base.py
@@ -13,11 +13,21 @@
 def _get_engine():
     global _engine, _AsyncSessionFactory
     if _engine is None:
-        _engine = create_async_engine(
-            settings.database_url,
-            echo=settings.log_level == "DEBUG",
-            connect_args={"check_same_thread": False} if "sqlite" in settings.database_url else {},
-        )
+        url = settings.database_url
+        is_sqlite = "sqlite" in url
+        kwargs: dict = {
+            "echo": settings.log_level == "DEBUG",
+            "connect_args": {"check_same_thread": False} if is_sqlite else {},
+        }
+        if not is_sqlite:
+            kwargs.update(
+                pool_size=settings.db_pool_size,
+                max_overflow=settings.db_max_overflow,
+                pool_recycle=settings.db_pool_recycle,
+                pool_pre_ping=settings.db_pool_pre_ping,
+                pool_timeout=settings.db_pool_timeout,
+            )
+        _engine = create_async_engine(url, **kwargs)
         _AsyncSessionFactory = async_sessionmaker(_engine, expire_on_commit=False)
     return _engine