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>
This commit is contained in:
Amy
2025-08-26 18:14:06 +08:00
committed by GitHub
parent fa753239ad
commit b5c2756261
3 changed files with 51 additions and 0 deletions

View File

@@ -215,6 +215,7 @@ class DatabaseConfig(BaseSettings):
"pool_pre_ping": self.SQLALCHEMY_POOL_PRE_PING, "pool_pre_ping": self.SQLALCHEMY_POOL_PRE_PING,
"connect_args": connect_args, "connect_args": connect_args,
"pool_use_lifo": self.SQLALCHEMY_POOL_USE_LIFO, "pool_use_lifo": self.SQLALCHEMY_POOL_USE_LIFO,
"pool_reset_on_return": None,
} }

View File

@@ -1,6 +1,55 @@
import logging
import gevent
from sqlalchemy import event
from sqlalchemy.pool import Pool
from dify_app import DifyApp from dify_app import DifyApp
from models import db 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): def init_app(app: DifyApp):
db.init_app(app) db.init_app(app)
_setup_gevent_compatibility()

View File

@@ -90,6 +90,7 @@ def test_flask_configs(monkeypatch):
"pool_recycle": 3600, "pool_recycle": 3600,
"pool_size": 30, "pool_size": 30,
"pool_use_lifo": False, "pool_use_lifo": False,
"pool_reset_on_return": None,
} }
assert config["CONSOLE_WEB_URL"] == "https://example.com" assert config["CONSOLE_WEB_URL"] == "https://example.com"