feat(tracing): Add more sampling context for asgi, celery, rq, and wsgi (#906)

Author: lobsterkatie

sentry_sdk/integrations/asgi.py

@@ -139,7 +139,9 @@ async def _run_app(self, scope, callback):
                 transaction.name = _DEFAULT_TRANSACTION_NAME
                 transaction.set_tag("asgi.type", ty)
 
-                with hub.start_transaction(transaction):
+                with hub.start_transaction(
+                    transaction, custom_sampling_context={"asgi_scope": scope}
+                ):
                     # XXX: Would be cool to have correct span status, but we
                     # would have to wrap send(). That is a bit hard to do with
                     # the current abstraction over ASGI 2/3.

sentry_sdk/integrations/celery.py

@@ -159,7 +159,18 @@ def _inner(*args, **kwargs):
             if transaction is None:
                 return f(*args, **kwargs)
 
-            with hub.start_transaction(transaction):
+            with hub.start_transaction(
+                transaction,
+                custom_sampling_context={
+                    "celery_job": {
+                        "task": task.name,
+                        # for some reason, args[1] is a list if non-empty but a
+                        # tuple if empty
+                        "args": list(args[1]),
+                        "kwargs": args[2],
+                    }
+                },
+            ):
                 return f(*args, **kwargs)
 
     return _inner  # type: ignore

sentry_sdk/integrations/rq.py

@@ -70,7 +70,9 @@ def sentry_patched_perform_job(self, job, *args, **kwargs):
                 with capture_internal_exceptions():
                     transaction.name = job.func_name
 
-                with hub.start_transaction(transaction):
+                with hub.start_transaction(
+                    transaction, custom_sampling_context={"rq_job": job}
+                ):
                     rv = old_perform_job(self, job, *args, **kwargs)
 
             if self.is_horse:

sentry_sdk/integrations/wsgi.py

@@ -117,7 +117,9 @@ def __call__(self, environ, start_response):
                         environ, op="http.server", name="generic WSGI request"
                     )
 
-                    with hub.start_transaction(transaction):
+                    with hub.start_transaction(
+                        transaction, custom_sampling_context={"wsgi_environ": environ}
+                    ):
                         try:
                             rv = self.app(
                                 environ,

tests/integrations/asgi/test_asgi.py

@@ -8,6 +8,11 @@
 from starlette.testclient import TestClient
 from starlette.websockets import WebSocket
 
+try:
+    from unittest import mock  # python 3.3 and above
+except ImportError:
+    import mock  # python < 3.3
+
 
 @pytest.fixture
 def app():
@@ -202,3 +207,47 @@ def handler(*args, **kwargs):
     (exception,) = event["exception"]["values"]
     assert exception["type"] == "ValueError"
     assert exception["value"] == "oh no"
+
+
+def test_transaction(app, sentry_init, capture_events):
+    sentry_init(traces_sample_rate=1.0)
+    events = capture_events()
+
+    @app.route("/tricks/kangaroo")
+    def kangaroo_handler(request):
+        return PlainTextResponse("dogs are great")
+
+    client = TestClient(app)
+    client.get("/tricks/kangaroo")
+
+    event = events[0]
+    assert event["type"] == "transaction"
+    assert (
+        event["transaction"]
+        == "tests.integrations.asgi.test_asgi.test_transaction.<locals>.kangaroo_handler"
+    )
+
+
+def test_traces_sampler_gets_scope_in_sampling_context(
+    app, sentry_init, DictionaryContaining  # noqa: N803
+):
+    traces_sampler = mock.Mock()
+    sentry_init(traces_sampler=traces_sampler)
+
+    @app.route("/tricks/kangaroo")
+    def kangaroo_handler(request):
+        return PlainTextResponse("dogs are great")
+
+    client = TestClient(app)
+    client.get("/tricks/kangaroo")
+
+    traces_sampler.assert_any_call(
+        DictionaryContaining(
+            {
+                # starlette just uses a dictionary to hold the scope
+                "asgi_scope": DictionaryContaining(
+                    {"method": "GET", "path": "/tricks/kangaroo"}
+                )
+            }
+        )
+    )

tests/integrations/celery/test_celery.py

@@ -11,6 +11,11 @@
 from celery import Celery, VERSION
 from celery.bin import worker
 
+try:
+    from unittest import mock  # python 3.3 and above
+except ImportError:
+    import mock  # python < 3.3
+
 
 @pytest.fixture
 def connect_signal(request):
@@ -379,3 +384,26 @@ def dummy_task(self, x, y):
 
     assert dummy_task.apply(kwargs={"x": 1, "y": 1}).wait() == 1
     assert celery_invocation(dummy_task, 1, 1)[0].wait() == 1
+
+
+def test_traces_sampler_gets_task_info_in_sampling_context(
+    init_celery, celery_invocation, DictionaryContaining  # noqa:N803
+):
+    traces_sampler = mock.Mock()
+    celery = init_celery(traces_sampler=traces_sampler)
+
+    @celery.task(name="dog_walk")
+    def walk_dogs(x, y):
+        dogs, route = x
+        num_loops = y
+        return dogs, route, num_loops
+
+    _, args_kwargs = celery_invocation(
+        walk_dogs, [["Maisey", "Charlie", "Bodhi", "Cory"], "Dog park round trip"], 1
+    )
+
+    traces_sampler.assert_any_call(
+        # depending on the iteration of celery_invocation, the data might be
+        # passed as args or as kwargs, so make this generic
+        DictionaryContaining({"celery_job": dict(task="dog_walk", **args_kwargs)})
+    )

tests/integrations/rq/test_rq.py

@@ -5,6 +5,11 @@
 from fakeredis import FakeStrictRedis
 import rq
 
+try:
+    from unittest import mock  # python 3.3 and above
+except ImportError:
+    import mock  # python < 3.3
+
 
 @pytest.fixture(autouse=True)
 def _patch_rq_get_server_version(monkeypatch):
@@ -28,6 +33,14 @@ def crashing_job(foo):
     1 / 0
 
 
+def chew_up_shoes(dog, human, shoes):
+    raise Exception("{}!! Why did you eat {}'s {}??".format(dog, human, shoes))
+
+
+def do_trick(dog, trick):
+    return "{}, can you {}? Good dog!".format(dog, trick)
+
+
 def test_basic(sentry_init, capture_events):
     sentry_init(integrations=[RqIntegration()])
     events = capture_events()
@@ -71,3 +84,96 @@ def test_transport_shutdown(sentry_init, capture_events_forksafe):
 
     (exception,) = event["exception"]["values"]
     assert exception["type"] == "ZeroDivisionError"
+
+
+def test_transaction_with_error(
+    sentry_init, capture_events, DictionaryContaining  # noqa:N803
+):
+
+    sentry_init(integrations=[RqIntegration()], traces_sample_rate=1.0)
+    events = capture_events()
+
+    queue = rq.Queue(connection=FakeStrictRedis())
+    worker = rq.SimpleWorker([queue], connection=queue.connection)
+
+    queue.enqueue(chew_up_shoes, "Charlie", "Katie", shoes="flip-flops")
+    worker.work(burst=True)
+
+    error_event, envelope = events
+
+    assert error_event["transaction"] == "tests.integrations.rq.test_rq.chew_up_shoes"
+    assert error_event["contexts"]["trace"]["op"] == "rq.task"
+    assert error_event["exception"]["values"][0]["type"] == "Exception"
+    assert (
+        error_event["exception"]["values"][0]["value"]
+        == "Charlie!! Why did you eat Katie's flip-flops??"
+    )
+
+    assert envelope["type"] == "transaction"
+    assert envelope["contexts"]["trace"] == error_event["contexts"]["trace"]
+    assert envelope["transaction"] == error_event["transaction"]
+    assert envelope["extra"]["rq-job"] == DictionaryContaining(
+        {
+            "args": ["Charlie", "Katie"],
+            "kwargs": {"shoes": "flip-flops"},
+            "func": "tests.integrations.rq.test_rq.chew_up_shoes",
+            "description": "tests.integrations.rq.test_rq.chew_up_shoes('Charlie', 'Katie', shoes='flip-flops')",
+        }
+    )
+
+
+def test_transaction_no_error(
+    sentry_init, capture_events, DictionaryContaining  # noqa:N803
+):
+    sentry_init(integrations=[RqIntegration()], traces_sample_rate=1.0)
+    events = capture_events()
+
+    queue = rq.Queue(connection=FakeStrictRedis())
+    worker = rq.SimpleWorker([queue], connection=queue.connection)
+
+    queue.enqueue(do_trick, "Maisey", trick="kangaroo")
+    worker.work(burst=True)
+
+    envelope = events[0]
+
+    assert envelope["type"] == "transaction"
+    assert envelope["contexts"]["trace"]["op"] == "rq.task"
+    assert envelope["transaction"] == "tests.integrations.rq.test_rq.do_trick"
+    assert envelope["extra"]["rq-job"] == DictionaryContaining(
+        {
+            "args": ["Maisey"],
+            "kwargs": {"trick": "kangaroo"},
+            "func": "tests.integrations.rq.test_rq.do_trick",
+            "description": "tests.integrations.rq.test_rq.do_trick('Maisey', trick='kangaroo')",
+        }
+    )
+
+
+def test_traces_sampler_gets_correct_values_in_sampling_context(
+    sentry_init, DictionaryContaining, ObjectDescribedBy  # noqa:N803
+):
+    traces_sampler = mock.Mock(return_value=True)
+    sentry_init(integrations=[RqIntegration()], traces_sampler=traces_sampler)
+
+    queue = rq.Queue(connection=FakeStrictRedis())
+    worker = rq.SimpleWorker([queue], connection=queue.connection)
+
+    queue.enqueue(do_trick, "Bodhi", trick="roll over")
+    worker.work(burst=True)
+
+    traces_sampler.assert_any_call(
+        DictionaryContaining(
+            {
+                "rq_job": ObjectDescribedBy(
+                    type=rq.job.Job,
+                    attrs={
+                        "description": "tests.integrations.rq.test_rq.do_trick('Bodhi', trick='roll over')",
+                        "result": "Bodhi, can you roll over? Good dog!",
+                        "func_name": "tests.integrations.rq.test_rq.do_trick",
+                        "args": ("Bodhi",),
+                        "kwargs": {"trick": "roll over"},
+                    },
+                ),
+            }
+        )
+    )

tests/integrations/wsgi/test_wsgi.py

@@ -3,6 +3,11 @@
 
 from sentry_sdk.integrations.wsgi import SentryWsgiMiddleware
 
+try:
+    from unittest import mock  # python 3.3 and above
+except ImportError:
+    import mock  # python < 3.3
+
 
 @pytest.fixture
 def crashing_app():
@@ -109,3 +114,90 @@ def test_keyboard_interrupt_is_captured(sentry_init, capture_events):
     assert exc["type"] == "KeyboardInterrupt"
     assert exc["value"] == ""
     assert event["level"] == "error"
+
+
+def test_transaction_with_error(
+    sentry_init, crashing_app, capture_events, DictionaryContaining  # noqa:N803
+):
+    def dogpark(environ, start_response):
+        raise Exception("Fetch aborted. The ball was not returned.")
+
+    sentry_init(send_default_pii=True, traces_sample_rate=1.0)
+    app = SentryWsgiMiddleware(dogpark)
+    client = Client(app)
+    events = capture_events()
+
+    with pytest.raises(Exception):
+        client.get("http://dogs.are.great/sit/stay/rollover/")
+
+    error_event, envelope = events
+
+    assert error_event["transaction"] == "generic WSGI request"
+    assert error_event["contexts"]["trace"]["op"] == "http.server"
+    assert error_event["exception"]["values"][0]["type"] == "Exception"
+    assert (
+        error_event["exception"]["values"][0]["value"]
+        == "Fetch aborted. The ball was not returned."
+    )
+
+    assert envelope["type"] == "transaction"
+
+    # event trace context is a subset of envelope trace context
+    assert envelope["contexts"]["trace"] == DictionaryContaining(
+        error_event["contexts"]["trace"]
+    )
+    assert envelope["contexts"]["trace"]["status"] == "internal_error"
+    assert envelope["transaction"] == error_event["transaction"]
+    assert envelope["request"] == error_event["request"]
+
+
+def test_transaction_no_error(
+    sentry_init, capture_events, DictionaryContaining  # noqa:N803
+):
+    def dogpark(environ, start_response):
+        start_response("200 OK", [])
+        return ["Go get the ball! Good dog!"]
+
+    sentry_init(send_default_pii=True, traces_sample_rate=1.0)
+    app = SentryWsgiMiddleware(dogpark)
+    client = Client(app)
+    events = capture_events()
+
+    client.get("/dogs/are/great/")
+
+    envelope = events[0]
+
+    assert envelope["type"] == "transaction"
+    assert envelope["transaction"] == "generic WSGI request"
+    assert envelope["contexts"]["trace"]["op"] == "http.server"
+    assert envelope["request"] == DictionaryContaining(
+        {"method": "GET", "url": "http://localhost/dogs/are/great/"}
+    )
+
+
+def test_traces_sampler_gets_correct_values_in_sampling_context(
+    sentry_init, DictionaryContaining, ObjectDescribedBy  # noqa:N803
+):
+    def app(environ, start_response):
+        start_response("200 OK", [])
+        return ["Go get the ball! Good dog!"]
+
+    traces_sampler = mock.Mock(return_value=True)
+    sentry_init(send_default_pii=True, traces_sampler=traces_sampler)
+    app = SentryWsgiMiddleware(app)
+    client = Client(app)
+
+    client.get("/dogs/are/great/")
+
+    traces_sampler.assert_any_call(
+        DictionaryContaining(
+            {
+                "wsgi_environ": DictionaryContaining(
+                    {
+                        "PATH_INFO": "/dogs/are/great/",
+                        "REQUEST_METHOD": "GET",
+                    },
+                ),
+            }
+        )
+    )