feat: Spans for Django middleware calls (#498)

* feat: Spans for Django middleware calls

* fix: Work under Django 1.6

* test: Add tests for django middlewares

Author: untitaker

sentry_sdk/integrations/django/__init__.py

@@ -48,6 +48,7 @@
 from sentry_sdk.integrations._wsgi_common import RequestExtractor
 from sentry_sdk.integrations.django.transactions import LEGACY_RESOLVER
 from sentry_sdk.integrations.django.templates import get_template_frame_from_exception
+from sentry_sdk.integrations.django.middleware import patch_django_middlewares
 
 
 if DJANGO_VERSION < (1, 10):
@@ -68,16 +69,18 @@ class DjangoIntegration(Integration):
     identifier = "django"
 
     transaction_style = None
+    middleware_spans = None
 
-    def __init__(self, transaction_style="url"):
-        # type: (str) -> None
+    def __init__(self, transaction_style="url", middleware_spans=True):
+        # type: (str, bool) -> None
         TRANSACTION_STYLE_VALUES = ("function_name", "url")
         if transaction_style not in TRANSACTION_STYLE_VALUES:
             raise ValueError(
                 "Invalid value for transaction_style: %s (must be in %s)"
                 % (transaction_style, TRANSACTION_STYLE_VALUES)
             )
         self.transaction_style = transaction_style
+        self.middleware_spans = middleware_spans
 
     @staticmethod
     def setup_once():
@@ -208,6 +211,8 @@ def _django_queryset_repr(value, hint):
                 id(value),
             )
 
+        patch_django_middlewares()
+
 
 _DRF_PATCHED = False
 _DRF_PATCH_LOCK = threading.Lock()

sentry_sdk/integrations/django/middleware.py

@@ -0,0 +1,106 @@
+"""
+Create spans from Django middleware invocations
+"""
+
+from functools import wraps
+
+from django import VERSION as DJANGO_VERSION  # type: ignore
+
+from sentry_sdk import Hub
+from sentry_sdk.utils import ContextVar, transaction_from_function
+
+_import_string_should_wrap_middleware = ContextVar(
+    "import_string_should_wrap_middleware"
+)
+
+if DJANGO_VERSION < (1, 7):
+    import_string_name = "import_by_path"
+else:
+    import_string_name = "import_string"
+
+
+def patch_django_middlewares():
+    from django.core.handlers import base
+
+    old_import_string = getattr(base, import_string_name)
+
+    def sentry_patched_import_string(dotted_path):
+        rv = old_import_string(dotted_path)
+
+        if _import_string_should_wrap_middleware.get(None):
+            rv = _wrap_middleware(rv, dotted_path)
+
+        return rv
+
+    setattr(base, import_string_name, sentry_patched_import_string)
+
+    old_load_middleware = base.BaseHandler.load_middleware
+
+    def sentry_patched_load_middleware(self):
+        _import_string_should_wrap_middleware.set(True)
+        try:
+            return old_load_middleware(self)
+        finally:
+            _import_string_should_wrap_middleware.set(False)
+
+    base.BaseHandler.load_middleware = sentry_patched_load_middleware
+
+
+def _wrap_middleware(middleware, middleware_name):
+    from sentry_sdk.integrations.django import DjangoIntegration
+
+    def _get_wrapped_method(old_method):
+        @wraps(old_method)
+        def sentry_wrapped_method(*args, **kwargs):
+            hub = Hub.current
+            integration = hub.get_integration(DjangoIntegration)
+            if integration is None or not integration.middleware_spans:
+                return old_method(*args, **kwargs)
+
+            function_name = transaction_from_function(old_method)
+
+            description = middleware_name
+            function_basename = getattr(old_method, "__name__", None)
+            if function_basename:
+                description = "{}.{}".format(description, function_basename)
+
+            with hub.start_span(
+                op="django.middleware", description=description
+            ) as span:
+                span.set_tag("django.function_name", function_name)
+                span.set_tag("django.middleware_name", middleware_name)
+                return old_method(*args, **kwargs)
+
+        return sentry_wrapped_method
+
+    class SentryWrappingMiddleware(object):
+        def __init__(self, *args, **kwargs):
+            self._inner = middleware(*args, **kwargs)
+            self._call_method = None
+
+        # We need correct behavior for `hasattr()`, which we can only determine
+        # when we have an instance of the middleware we're wrapping.
+        def __getattr__(self, method_name):
+            if method_name not in (
+                "process_request",
+                "process_view",
+                "process_template_response",
+                "process_response",
+                "process_exception",
+            ):
+                raise AttributeError()
+
+            old_method = getattr(self._inner, method_name)
+            rv = _get_wrapped_method(old_method)
+            self.__dict__[method_name] = rv
+            return rv
+
+        def __call__(self, *args, **kwargs):
+            if self._call_method is None:
+                self._call_method = _get_wrapped_method(self._inner.__call__)
+            return self._call_method(*args, **kwargs)
+
+    if hasattr(middleware, "__name__"):
+        SentryWrappingMiddleware.__name__ = middleware.__name__
+
+    return SentryWrappingMiddleware

tests/integrations/django/test_basic.py

@@ -2,6 +2,7 @@
 import json
 
 from werkzeug.test import Client
+from django import VERSION as DJANGO_VERSION
 from django.contrib.auth.models import User
 from django.core.management import execute_from_command_line
 from django.db.utils import OperationalError, ProgrammingError, DataError
@@ -495,3 +496,49 @@ def test_does_not_capture_403(sentry_init, client, capture_events, endpoint):
     assert status.lower() == "403 forbidden"
 
     assert not events
+
+
+def test_middleware_spans(sentry_init, client, capture_events):
+    sentry_init(integrations=[DjangoIntegration()], traces_sample_rate=1.0)
+    events = capture_events()
+
+    _content, status, _headers = client.get(reverse("message"))
+
+    message, transaction = events
+
+    assert message["message"] == "hi"
+
+    for middleware in transaction["spans"]:
+        assert middleware["op"] == "django.middleware"
+
+    if DJANGO_VERSION >= (1, 10):
+        reference_value = [
+            "tests.integrations.django.myapp.settings.TestMiddleware.__call__",
+            "django.contrib.auth.middleware.AuthenticationMiddleware.__call__",
+            "django.contrib.sessions.middleware.SessionMiddleware.__call__",
+        ]
+    else:
+        reference_value = [
+            "django.contrib.sessions.middleware.SessionMiddleware.process_request",
+            "django.contrib.auth.middleware.AuthenticationMiddleware.process_request",
+            "tests.integrations.django.myapp.settings.TestMiddleware.process_request",
+            "tests.integrations.django.myapp.settings.TestMiddleware.process_response",
+            "django.contrib.sessions.middleware.SessionMiddleware.process_response",
+        ]
+
+    assert [t["description"] for t in transaction["spans"]] == reference_value
+
+
+def test_middleware_spans_disabled(sentry_init, client, capture_events):
+    sentry_init(
+        integrations=[DjangoIntegration(middleware_spans=False)], traces_sample_rate=1.0
+    )
+    events = capture_events()
+
+    _content, status, _headers = client.get(reverse("message"))
+
+    message, transaction = events
+
+    assert message["message"] == "hi"
+
+    assert not transaction["spans"]