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:
@@ -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,
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|
||||||
|
@@ -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()
|
||||||
|
@@ -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"
|
||||||
|
Reference in New Issue
Block a user