feat(tracing): Add aiohttp request object to sampling context (getsentry#888)

lobsterkatie · web-flow · commit 617c516d7261 · 2020-10-30T07:05:48.000-07:00
diff --git a/sentry_sdk/integrations/aiohttp.py b/sentry_sdk/integrations/aiohttp.py
@@ -106,8 +106,9 @@ async def sentry_app_handle(self, request, *args, **kwargs):
                     # URL resolver did not find a route or died trying.
                     name="generic AIOHTTP request",
                 )
-
-                with hub.start_transaction(transaction):
+                with hub.start_transaction(
+                    transaction, custom_sampling_context={"aiohttp_request": request}
+                ):
                     try:
                         response = await old_handle(self, request)
                     except HTTPException as e:
diff --git a/tests/integrations/aiohttp/test_aiohttp.py b/tests/integrations/aiohttp/test_aiohttp.py
@@ -5,9 +5,15 @@
 import pytest
 from aiohttp import web
 from aiohttp.client import ServerDisconnectedError
+from aiohttp.web_request import Request
 
 from sentry_sdk.integrations.aiohttp import AioHttpIntegration
 
+try:
+    from unittest import mock  # python 3.3 and above
+except ImportError:
+    import mock  # python < 3.3
+
 
 async def test_basic(sentry_init, aiohttp_client, loop, capture_events):
     sentry_init(integrations=[AioHttpIntegration()])
@@ -223,3 +229,35 @@ async def hello(request):
 
     assert event["type"] == "transaction"
     assert event["transaction"] == expected_transaction
+
+
+async def test_traces_sampler_gets_request_object_in_sampling_context(
+    sentry_init,
+    aiohttp_client,
+    DictionaryContaining,  # noqa:N803
+    ObjectDescribedBy,  # noqa:N803
+):
+    traces_sampler = mock.Mock()
+    sentry_init(
+        integrations=[AioHttpIntegration()],
+        traces_sampler=traces_sampler,
+    )
+
+    async def kangaroo_handler(request):
+        return web.Response(text="dogs are great")
+
+    app = web.Application()
+    app.router.add_get("/tricks/kangaroo", kangaroo_handler)
+
+    client = await aiohttp_client(app)
+    await client.get("/tricks/kangaroo")
+
+    traces_sampler.assert_any_call(
+        DictionaryContaining(
+            {
+                "aiohttp_request": ObjectDescribedBy(
+                    type=Request, attrs={"method": "GET", "path": "/tricks/kangaroo"}
+                )
+            }
+        )
+    )
