Bladeren bron

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>
tags/1.8.0
Amy 2 maanden geleden
bovenliggende
commit
b5c2756261
No account linked to committer's email address

+ 1
- 0
api/configs/middleware/__init__.py Bestand weergeven

@@ -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,
}



+ 49
- 0
api/extensions/ext_database.py Bestand weergeven

@@ -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()

+ 1
- 0
api/tests/unit_tests/configs/test_dify_config.py Bestand weergeven

@@ -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"

Laden…
Annuleren
Opslaan