feat: Instrument ASGI under Django 3.0 (getsentry#573)

untitaker · web-flow · commit 19200af6dbb6 · 2019-12-11T13:57:44.000+01:00
* feat: Instrument ASGI under Django 3.0

* fix: Clarify doc comment
diff --git a/sentry_sdk/integrations/django/__init__.py b/sentry_sdk/integrations/django/__init__.py
@@ -109,6 +109,8 @@ def sentry_patched_wsgi_handler(self, environ, start_response):
 
         WSGIHandler.__call__ = sentry_patched_wsgi_handler
 
+        _patch_django_asgi_handler()
+
         # patch get_response, because at that point we have the Django request
         # object
         from django.core.handlers.base import BaseHandler
@@ -314,6 +316,30 @@ def sentry_patched_asgi_handler(self, receive, send):
     AsgiHandler.__call__ = sentry_patched_asgi_handler
 
 
+def _patch_django_asgi_handler():
+    # type: () -> None
+    try:
+        from django.core.handlers.asgi import ASGIHandler
+    except ImportError:
+        return
+
+    if not HAS_REAL_CONTEXTVARS:
+        # We better have contextvars or we're going to leak state between
+        # requests.
+        #
+        # We cannot hard-raise here because Django may not be used at all in
+        # the current process.
+        logger.warning(
+            "We detected that you are using Django 3. To get proper "
+            "instrumentation for ASGI requests, the Sentry SDK requires "
+            "Python 3.7+ or the aiocontextvars package from PyPI."
+        )
+
+    from sentry_sdk.integrations.django.asgi import patch_django_asgi_handler_impl
+
+    patch_django_asgi_handler_impl(ASGIHandler)
+
+
 def _make_event_processor(weak_request, integration):
     # type: (Callable[[], WSGIRequest], DjangoIntegration) -> EventProcessor
     def event_processor(event, hint):
diff --git a/sentry_sdk/integrations/django/asgi.py b/sentry_sdk/integrations/django/asgi.py
@@ -0,0 +1,31 @@
+"""
+Instrumentation for Django 3.0
+
+Since this file contains `async def` it is conditionally imported in
+`sentry_sdk.integrations.django` (depending on the existence of
+`django.core.handlers.asgi`.
+"""
+
+from sentry_sdk import Hub
+from sentry_sdk._types import MYPY
+
+from sentry_sdk.integrations.django import DjangoIntegration
+from sentry_sdk.integrations.asgi import SentryAsgiMiddleware
+
+if MYPY:
+    from typing import Any
+
+
+def patch_django_asgi_handler_impl(cls):
+    # type: (Any) -> None
+    old_app = cls.__call__
+
+    async def sentry_patched_asgi_handler(self, scope, receive, send):
+        # type: (Any, Any, Any, Any) -> Any
+        if Hub.current.get_integration(DjangoIntegration) is None:
+            return await old_app(self, scope, receive, send)
+
+        middleware = SentryAsgiMiddleware(old_app.__get__(self, cls))._run_asgi3
+        return await middleware(scope, receive, send)
+
+    cls.__call__ = sentry_patched_asgi_handler
diff --git a/tests/integrations/django/asgi/__init__.py b/tests/integrations/django/asgi/__init__.py
diff --git a/tests/integrations/django/asgi/test_asgi.py b/tests/integrations/django/asgi/test_asgi.py
@@ -1,16 +1,24 @@
 import pytest
 
+import django
 
 from channels.testing import HttpCommunicator
 
 from sentry_sdk import capture_message
 from sentry_sdk.integrations.django import DjangoIntegration
 
-from tests.integrations.django.myapp.asgi import application
+from tests.integrations.django.myapp.asgi import channels_application
 
+APPS = [channels_application]
+if django.VERSION >= (3, 0):
+    from tests.integrations.django.myapp.asgi import asgi_application
 
+    APPS += [asgi_application]
+
+
+@pytest.mark.parametrize("application", APPS)
 @pytest.mark.asyncio
-async def test_basic(sentry_init, capture_events):
+async def test_basic(sentry_init, capture_events, application):
     sentry_init(integrations=[DjangoIntegration()], send_default_pii=True)
     events = capture_events()
 
diff --git a/tests/integrations/django/myapp/asgi.py b/tests/integrations/django/myapp/asgi.py
@@ -12,4 +12,9 @@
 )
 
 django.setup()
-application = get_default_application()
+channels_application = get_default_application()
+
+if django.VERSION >= (3, 0):
+    from django.core.asgi import get_asgi_application
+
+    asgi_application = get_asgi_application()
