Merge branch 'main' into feat/rag-2

api/extensions/ext_database.py

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

api/extensions/ext_login.py

@@ -21,7 +21,7 @@ login_manager = flask_login.LoginManager()
 def load_user_from_request(request_from_flask_login):
     """Load user based on the request."""
     # Skip authentication for documentation endpoints
-    if request.path.endswith("/docs") or request.path.endswith("/swagger.json"):
+    if dify_config.SWAGGER_UI_ENABLED and request.path.endswith((dify_config.SWAGGER_UI_PATH, "/swagger.json")):
         return None
 
     auth_header = request.headers.get("Authorization", "")

api/extensions/ext_mail.py

@@ -6,6 +6,8 @@ from flask import Flask
 from configs import dify_config
 from dify_app import DifyApp
 
+logger = logging.getLogger(__name__)
+
 
 class Mail:
     def __init__(self):
@@ -18,7 +20,7 @@ class Mail:
     def init_app(self, app: Flask):
         mail_type = dify_config.MAIL_TYPE
         if not mail_type:
-            logging.warning("MAIL_TYPE is not set")
+            logger.warning("MAIL_TYPE is not set")
             return
 
         if dify_config.MAIL_DEFAULT_SEND_FROM:

api/extensions/ext_otel.py

@@ -16,6 +16,8 @@ from dify_app import DifyApp
 from libs.helper import extract_tenant_id
 from models import Account, EndUser
 
+logger = logging.getLogger(__name__)
+
 
 @user_logged_in.connect
 @user_loaded_from_request.connect
@@ -33,7 +35,7 @@ def on_user_loaded(_sender, user: Union["Account", "EndUser"]):
                     current_span.set_attribute("service.tenant.id", tenant_id)
                     current_span.set_attribute("service.user.id", user.id)
             except Exception:
-                logging.exception("Error setting tenant and user attributes")
+                logger.exception("Error setting tenant and user attributes")
                 pass
 
 
@@ -74,12 +76,12 @@ def init_app(app: DifyApp):
                         attributes[SpanAttributes.HTTP_METHOD] = str(request.method)
                     _http_response_counter.add(1, attributes)
                 except Exception:
-                    logging.exception("Error setting status and attributes")
+                    logger.exception("Error setting status and attributes")
                     pass
 
         instrumentor = FlaskInstrumentor()
         if dify_config.DEBUG:
-            logging.info("Initializing Flask instrumentor")
+            logger.info("Initializing Flask instrumentor")
         instrumentor.instrument_app(app, response_hook=response_hook)
 
     def init_sqlalchemy_instrumentor(app: DifyApp):
@@ -253,5 +255,5 @@ def init_celery_worker(*args, **kwargs):
         tracer_provider = get_tracer_provider()
         metric_provider = get_meter_provider()
         if dify_config.DEBUG:
-            logging.info("Initializing OpenTelemetry for Celery worker")
+            logger.info("Initializing OpenTelemetry for Celery worker")
         CeleryInstrumentor(tracer_provider=tracer_provider, meter_provider=metric_provider).instrument()

api/extensions/ext_request_logging.py

@@ -8,7 +8,7 @@ from flask.signals import request_finished, request_started
 
 from configs import dify_config
 
-_logger = logging.getLogger(__name__)
+logger = logging.getLogger(__name__)
 
 
 def _is_content_type_json(content_type: str) -> bool:
@@ -20,20 +20,20 @@ def _is_content_type_json(content_type: str) -> bool:
 
 def _log_request_started(_sender, **_extra):
     """Log the start of a request."""
-    if not _logger.isEnabledFor(logging.DEBUG):
+    if not logger.isEnabledFor(logging.DEBUG):
         return
 
     request = flask.request
     if not (_is_content_type_json(request.content_type) and request.data):
-        _logger.debug("Received Request %s -> %s", request.method, request.path)
+        logger.debug("Received Request %s -> %s", request.method, request.path)
         return
     try:
         json_data = json.loads(request.data)
     except (TypeError, ValueError):
-        _logger.exception("Failed to parse JSON request")
+        logger.exception("Failed to parse JSON request")
         return
     formatted_json = json.dumps(json_data, ensure_ascii=False, indent=2)
-    _logger.debug(
+    logger.debug(
         "Received Request %s -> %s, Request Body:\n%s",
         request.method,
         request.path,
@@ -43,21 +43,21 @@ def _log_request_started(_sender, **_extra):
 
 def _log_request_finished(_sender, response, **_extra):
     """Log the end of a request."""
-    if not _logger.isEnabledFor(logging.DEBUG) or response is None:
+    if not logger.isEnabledFor(logging.DEBUG) or response is None:
         return
 
     if not _is_content_type_json(response.content_type):
-        _logger.debug("Response %s %s", response.status, response.content_type)
+        logger.debug("Response %s %s", response.status, response.content_type)
         return
 
     response_data = response.get_data(as_text=True)
     try:
         json_data = json.loads(response_data)
     except (TypeError, ValueError):
-        _logger.exception("Failed to parse JSON response")
+        logger.exception("Failed to parse JSON response")
         return
     formatted_json = json.dumps(json_data, ensure_ascii=False, indent=2)
-    _logger.debug(
+    logger.debug(
         "Response %s %s, Response Body:\n%s",
         response.status,
         response.content_type,