implement instrumentaton for async views

untitaker · untitaker · commit a57a6c2f33c2 · 2020-10-05T15:16:37.000+02:00
diff --git a/sentry_sdk/integrations/django/asgi.py b/sentry_sdk/integrations/django/asgi.py
@@ -6,10 +6,9 @@
 `django.core.handlers.asgi`.
 """
 
-from sentry_sdk import Hub
+from sentry_sdk import Hub, _functools
 from sentry_sdk._types import MYPY
 
-from sentry_sdk.integrations.django import DjangoIntegration
 from sentry_sdk.integrations.asgi import SentryAsgiMiddleware
 
 if MYPY:
@@ -21,6 +20,9 @@
 
 def patch_django_asgi_handler_impl(cls):
     # type: (Any) -> None
+
+    from sentry_sdk.integrations.django import DjangoIntegration
+
     old_app = cls.__call__
 
     async def sentry_patched_asgi_handler(self, scope, receive, send):
@@ -50,6 +52,9 @@ async def sentry_patched_get_response_async(self, request):
 
 def patch_channels_asgi_handler_impl(cls):
     # type: (Any) -> None
+
+    from sentry_sdk.integrations.django import DjangoIntegration
+
     old_app = cls.__call__
 
     async def sentry_patched_asgi_handler(self, receive, send):
@@ -64,3 +69,16 @@ async def sentry_patched_asgi_handler(self, receive, send):
         return await middleware(self.scope)(receive, send)
 
     cls.__call__ = sentry_patched_asgi_handler
+
+
+def wrap_async_view(hub, callback):
+    @_functools.wraps(callback)
+    async def sentry_wrapped_callback(request, *args, **kwargs):
+        # type: (Any, *Any, **Any) -> Any
+
+        with hub.start_span(
+            op="django.view", description=request.resolver_match.view_name
+        ):
+            return await callback(request, *args, **kwargs)
+
+    return sentry_wrapped_callback
diff --git a/sentry_sdk/integrations/django/views.py b/sentry_sdk/integrations/django/views.py
@@ -6,6 +6,18 @@
     from typing import Any
 
 
+try:
+    from asyncio import iscoroutinefunction
+except ImportError:
+    iscoroutinefunction = None
+
+
+try:
+    from sentry_sdk.integrations.django.asgi import wrap_async_view
+except (ImportError, SyntaxError):
+    wrap_async_view = None
+
+
 def patch_views():
     # type: () -> None
 
@@ -27,17 +39,27 @@ def sentry_patched_make_view_atomic(self, *args, **kwargs):
 
         if integration is not None and integration.middleware_spans:
 
-            @_functools.wraps(callback)
-            def sentry_wrapped_callback(request, *args, **kwargs):
-                # type: (Any, *Any, **Any) -> Any
-                with hub.start_span(
-                    op="django.view", description=request.resolver_match.view_name
-                ):
-                    return callback(request, *args, **kwargs)
+            if iscoroutinefunction is not None and wrap_async_view is not None and iscoroutinefunction(callback):
+                sentry_wrapped_callback = wrap_async_view(hub, callback)
+            else:
+                print((iscoroutinefunction, wrap_async_view, iscoroutinefunction(callback)))
+                sentry_wrapped_callback = _wrap_sync_view(hub, callback)
 
         else:
             sentry_wrapped_callback = callback
 
         return sentry_wrapped_callback
 
     BaseHandler.make_view_atomic = sentry_patched_make_view_atomic
+
+
+def _wrap_sync_view(hub, callback):
+    @_functools.wraps(callback)
+    def sentry_wrapped_callback(request, *args, **kwargs):
+        # type: (Any, *Any, **Any) -> Any
+        with hub.start_span(
+            op="django.view", description=request.resolver_match.view_name
+        ):
+            return callback(request, *args, **kwargs)
+
+    return sentry_wrapped_callback
diff --git a/tests/integrations/django/asgi/test_asgi.py b/tests/integrations/django/asgi/test_asgi.py
@@ -51,6 +51,17 @@ async def test_async_views(sentry_init, capture_events, application):
 
     events = capture_events()
 
-    comm = HttpCommunicator(application, "GET", "/async_ok")
+    comm = HttpCommunicator(application, "GET", "/async_message")
     response = await comm.get_response()
     assert response["status"] == 200
+
+    event, = events
+
+    assert event["transaction"] == "/async_message"
+    assert event["request"] == {
+        "cookies": {},
+        "headers": {},
+        "method": "GET",
+        "query_string": None,
+        "url": "/async_message",
+    }
diff --git a/tests/integrations/django/myapp/urls.py b/tests/integrations/django/myapp/urls.py
@@ -55,7 +55,7 @@ def path(path, *args, **kwargs):
         views.csrf_hello_not_exempt,
         name="csrf_hello_not_exempt",
     ),
-    path("async_ok", views.async_ok, name="async_ok"),
+    path("async_message", views.async_message, name="async_message"),
 ]
 
 
diff --git a/tests/integrations/django/myapp/views.py b/tests/integrations/django/myapp/views.py
@@ -122,5 +122,10 @@ def csrf_hello_not_exempt(*args, **kwargs):
     return HttpResponse("ok")
 
 
-async def async_ok(request):
-    return HttpResponse("ok")
+try:
+    exec("""async def async_message(request):
+    sentry_sdk.capture_message("hi")
+    return HttpResponse("ok")""")
+
+except SyntaxError:
+    async_message = None

Original file line number	Diff line number	Diff line change
`@@ -55,7 +55,7 @@ def path(path, args, *kwargs):`
`55`	`55`	`views.csrf_hello_not_exempt,`
`56`	`56`	`name="csrf_hello_not_exempt",`
`57`	`57`	`),`
`58`		`- path("async_ok", views.async_ok, name="async_ok"),`
	`58`	`+ path("async_message", views.async_message, name="async_message"),`
`59`	`59`	`]`
`60`	`60`
`61`	`61`