From b5c2756261cad335928abff44a6cb8d210310522 Mon Sep 17 00:00:00 2001
From: Amy <1530140574@qq.com>
Date: Tue, 26 Aug 2025 18:14:06 +0800
Subject: [PATCH] fix(api):safe reset in db pool, avoid rollback in gevent
 callback (#24556)

Co-authored-by: autofix-ci[bot] <114827586+autofix-ci[bot]@users.noreply.github.com>
---
 api/configs/middleware/__init__.py            |  1 +
 api/extensions/ext_database.py                | 49 +++++++++++++++++++
 .../unit_tests/configs/test_dify_config.py    |  1 +
 3 files changed, 51 insertions(+)

diff --git a/api/configs/middleware/__init__.py b/api/configs/middleware/__init__.py
index ba8bbc713..4751b9601 100644
--- a/api/configs/middleware/__init__.py
+++ b/api/configs/middleware/__init__.py
@@ -215,6 +215,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,
         }
 
 
diff --git a/api/extensions/ext_database.py b/api/extensions/ext_database.py
index 93842a303..b32616b17 100644
--- a/api/extensions/ext_database.py
+++ b/api/extensions/ext_database.py
@@ -1,6 +1,55 @@
+import logging
+
+import gevent
+from sqlalchemy import event
+from sqlalchemy.pool import Pool
+
 from dify_app import DifyApp
 from models import db
 
+logger = logging.getLogger(__name__)
+
+# Global flag to avoid duplicate registration of event listener
+_GEVENT_COMPATIBILITY_SETUP: bool = False
+
+
+def _safe_rollback(connection) -> None:
+    """Safely rollback database connection.
+
+    Args:
+        connection: Database connection object
+    """
+    try:
+        connection.rollback()
+    except Exception:  # pylint: disable=broad-exception-caught
+        logger.exception("Failed to rollback connection")
+
+
+def _setup_gevent_compatibility() -> None:
+    global _GEVENT_COMPATIBILITY_SETUP  # pylint: disable=global-statement
+
+    # Avoid duplicate registration
+    if _GEVENT_COMPATIBILITY_SETUP:
+        return
+
+    @event.listens_for(Pool, "reset")
+    def _safe_reset(dbapi_connection, connection_record, reset_state) -> None:  # pylint: disable=unused-argument
+        if reset_state.terminate_only:
+            return
+
+        # Safe rollback for connection
+        try:
+            hub = gevent.get_hub()
+            if hasattr(hub, "loop") and getattr(hub.loop, "in_callback", False):
+                gevent.spawn_later(0, lambda: _safe_rollback(dbapi_connection))
+            else:
+                _safe_rollback(dbapi_connection)
+        except (AttributeError, ImportError):
+            _safe_rollback(dbapi_connection)
+
+    _GEVENT_COMPATIBILITY_SETUP = True
+
 
 def init_app(app: DifyApp):
     db.init_app(app)
+    _setup_gevent_compatibility()
diff --git a/api/tests/unit_tests/configs/test_dify_config.py b/api/tests/unit_tests/configs/test_dify_config.py
index 0ae6a09f5..b95ed431b 100644
--- a/api/tests/unit_tests/configs/test_dify_config.py
+++ b/api/tests/unit_tests/configs/test_dify_config.py
@@ -90,6 +90,7 @@ def test_flask_configs(monkeypatch):
         "pool_recycle": 3600,
         "pool_size": 30,
         "pool_use_lifo": False,
+        "pool_reset_on_return": None,
     }
 
     assert config["CONSOLE_WEB_URL"] == "https://example.com"