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