From 8132c444dc8331e09e0025de6b40985f2a2cd9be Mon Sep 17 00:00:00 2001 From: wangxiaolei Date: Fri, 8 May 2026 20:25:46 +0800 Subject: [PATCH] feat: support SQLALCHEMY_POOL_RESET_ON_RETURN config (#31156) --- api/.env.example | 2 ++ api/configs/middleware/__init__.py | 9 +++++++-- api/tests/unit_tests/configs/test_dify_config.py | 2 +- docker/.env.example | 2 ++ docker/docker-compose.yaml | 1 + 5 files changed, 13 insertions(+), 3 deletions(-) diff --git a/api/.env.example b/api/.env.example index f6f65011ea..bec3c67ac8 100644 --- a/api/.env.example +++ b/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... diff --git a/api/configs/middleware/__init__.py b/api/configs/middleware/__init__.py index c392b8840f..ee8b93aa9f 100644 --- a/api/configs/middleware/__init__.py +++ b/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 diff --git a/api/tests/unit_tests/configs/test_dify_config.py b/api/tests/unit_tests/configs/test_dify_config.py index bad246a4bb..57dbf453de 100644 --- a/api/tests/unit_tests/configs/test_dify_config.py +++ b/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" diff --git a/docker/.env.example b/docker/.env.example index 122228cdd1..5871cacc39 100644 --- a/docker/.env.example +++ b/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 diff --git a/docker/docker-compose.yaml b/docker/docker-compose.yaml index 6dcab4a9fc..56f657f7a7 100644 --- a/docker/docker-compose.yaml +++ b/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}