feat: support SQLALCHEMY_POOL_RESET_ON_RETURN config (#31156)

api/.env.example

@@ -98,6 +98,8 @@ DB_DATABASE=dify
 
 SQLALCHEMY_POOL_PRE_PING=true
 SQLALCHEMY_POOL_TIMEOUT=30
+# Connection pool reset behavior on return
+SQLALCHEMY_POOL_RESET_ON_RETURN=rollback
 
 # Storage configuration
 # use for store upload files, private keys...

api/configs/middleware/__init__.py

@@ -114,7 +114,7 @@ class SQLAlchemyEngineOptionsDict(TypedDict):
     pool_pre_ping: bool
     connect_args: dict[str, str]
     pool_use_lifo: bool
-    pool_reset_on_return: None
+    pool_reset_on_return: Literal["commit", "rollback", None]
     pool_timeout: int
 
 
@@ -223,6 +223,11 @@ class DatabaseConfig(BaseSettings):
         default=30,
     )
 
+    SQLALCHEMY_POOL_RESET_ON_RETURN: Literal["commit", "rollback", None] = Field(
+        description="Connection pool reset behavior on return. Options: 'commit', 'rollback', or None",
+        default="rollback",
+    )
+
     RETRIEVAL_SERVICE_EXECUTORS: NonNegativeInt = Field(
         description="Number of processes for the retrieval service, default to CPU cores.",
         default=os.cpu_count() or 1,
@@ -252,7 +257,7 @@ class DatabaseConfig(BaseSettings):
             "pool_pre_ping": self.SQLALCHEMY_POOL_PRE_PING,
             "connect_args": connect_args,
             "pool_use_lifo": self.SQLALCHEMY_POOL_USE_LIFO,
-            "pool_reset_on_return": None,
+            "pool_reset_on_return": self.SQLALCHEMY_POOL_RESET_ON_RETURN,
             "pool_timeout": self.SQLALCHEMY_POOL_TIMEOUT,
         }
         return result

api/tests/unit_tests/configs/test_dify_config.py

@@ -114,8 +114,8 @@ def test_flask_configs(monkeypatch: pytest.MonkeyPatch):
         "pool_recycle": 3600,
         "pool_size": 30,
         "pool_use_lifo": False,
-        "pool_reset_on_return": None,
         "pool_timeout": 30,
+        "pool_reset_on_return": "rollback",
     }
 
     assert config["CONSOLE_WEB_URL"] == "https://example.com"

docker/.env.example

@@ -268,6 +268,8 @@ SQLALCHEMY_POOL_USE_LIFO=false
 # Number of seconds to wait for a connection from the pool before raising a timeout error.
 # Default is 30
 SQLALCHEMY_POOL_TIMEOUT=30
+# Connection pool reset behavior on return
+SQLALCHEMY_POOL_RESET_ON_RETURN=rollback
 
 # Maximum number of connections to the database
 # Default is 100

docker/docker-compose.yaml

@@ -70,6 +70,7 @@ x-shared-env: &shared-api-worker-env
   SQLALCHEMY_POOL_PRE_PING: ${SQLALCHEMY_POOL_PRE_PING:-false}
   SQLALCHEMY_POOL_USE_LIFO: ${SQLALCHEMY_POOL_USE_LIFO:-false}
   SQLALCHEMY_POOL_TIMEOUT: ${SQLALCHEMY_POOL_TIMEOUT:-30}
+  SQLALCHEMY_POOL_RESET_ON_RETURN: ${SQLALCHEMY_POOL_RESET_ON_RETURN:-rollback}
   POSTGRES_MAX_CONNECTIONS: ${POSTGRES_MAX_CONNECTIONS:-200}
   POSTGRES_SHARED_BUFFERS: ${POSTGRES_SHARED_BUFFERS:-128MB}
   POSTGRES_WORK_MEM: ${POSTGRES_WORK_MEM:-4MB}