a78339a040
Co-authored-by: autofix-ci[bot] <114827586+autofix-ci[bot]@users.noreply.github.com> Co-authored-by: -LAN- <laipz8200@outlook.com>
56 lines
1.5 KiB
Python
56 lines
1.5 KiB
Python
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):
|
|
"""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():
|
|
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): # 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()
|