fix(tracing): Get HTTP headers from span rather than transaction if possible (getsentry#1035)

lobsterkatie · web-flow · commit 51987c571571 · 2021-02-26T18:17:36.000-08:00
diff --git a/sentry_sdk/hub.py b/sentry_sdk/hub.py
@@ -682,15 +682,19 @@ def flush(
         if client is not None:
             return client.flush(timeout=timeout, callback=callback)
 
-    def iter_trace_propagation_headers(self):
-        # type: () -> Generator[Tuple[str, str], None, None]
-        # TODO: Document
-        client, scope = self._stack[-1]
-        span = scope.span
-
-        if span is None:
+    def iter_trace_propagation_headers(self, span=None):
+        # type: (Optional[Span]) -> Generator[Tuple[str, str], None, None]
+        """
+        Return HTTP headers which allow propagation of trace data. Data taken
+        from the span representing the request, if available, or the current
+        span on the scope if not.
+        """
+        span = span or self.scope.span
+        if not span:
             return
 
+        client = self._stack[-1][0]
+
         propagate_traces = client and client.options["propagate_traces"]
         if not propagate_traces:
             return
diff --git a/sentry_sdk/integrations/celery.py b/sentry_sdk/integrations/celery.py
@@ -96,9 +96,9 @@ def apply_async(*args, **kwargs):
         hub = Hub.current
         integration = hub.get_integration(CeleryIntegration)
         if integration is not None and integration.propagate_traces:
-            with hub.start_span(op="celery.submit", description=args[0].name):
+            with hub.start_span(op="celery.submit", description=args[0].name) as span:
                 with capture_internal_exceptions():
-                    headers = dict(hub.iter_trace_propagation_headers())
+                    headers = dict(hub.iter_trace_propagation_headers(span))
 
                     if headers:
                         # Note: kwargs can contain headers=None, so no setdefault!
diff --git a/sentry_sdk/integrations/stdlib.py b/sentry_sdk/integrations/stdlib.py
@@ -85,7 +85,7 @@ def putrequest(self, method, url, *args, **kwargs):
 
         rv = real_putrequest(self, method, url, *args, **kwargs)
 
-        for key, value in hub.iter_trace_propagation_headers():
+        for key, value in hub.iter_trace_propagation_headers(span):
             self.putheader(key, value)
 
         self._sentrysdk_span = span
@@ -178,12 +178,15 @@ def sentry_patched_popen_init(self, *a, **kw):
 
         env = None
 
-        for k, v in hub.iter_trace_propagation_headers():
-            if env is None:
-                env = _init_argument(a, kw, "env", 10, lambda x: dict(x or os.environ))
-            env["SUBPROCESS_" + k.upper().replace("-", "_")] = v
-
         with hub.start_span(op="subprocess", description=description) as span:
+
+            for k, v in hub.iter_trace_propagation_headers(span):
+                if env is None:
+                    env = _init_argument(
+                        a, kw, "env", 10, lambda x: dict(x or os.environ)
+                    )
+                env["SUBPROCESS_" + k.upper().replace("-", "_")] = v
+
             if cwd:
                 span.set_data("subprocess.cwd", cwd)
 
diff --git a/tests/conftest.py b/tests/conftest.py
@@ -368,15 +368,21 @@ def __init__(self, substring):
             self.substring = substring
 
             try:
-                # unicode only exists in python 2
+                # the `unicode` type only exists in python 2, so if this blows up,
+                # we must be in py3 and have the `bytes` type
                 self.valid_types = (str, unicode)  # noqa
             except NameError:
-                self.valid_types = (str,)
+                self.valid_types = (str, bytes)
 
         def __eq__(self, test_string):
             if not isinstance(test_string, self.valid_types):
                 return False
 
+            # this is safe even in py2 because as of 2.6, `bytes` exists in py2
+            # as an alias for `str`
+            if isinstance(test_string, bytes):
+                test_string = test_string.decode()
+
             if len(self.substring) > len(test_string):
                 return False
 
diff --git a/tests/integrations/stdlib/test_httplib.py b/tests/integrations/stdlib/test_httplib.py
@@ -17,7 +17,12 @@
     # py3
     from http.client import HTTPSConnection
 
-from sentry_sdk import capture_message
+try:
+    from unittest import mock  # python 3.3 and above
+except ImportError:
+    import mock  # python < 3.3
+
+from sentry_sdk import capture_message, start_transaction
 from sentry_sdk.integrations.stdlib import StdlibIntegration
 
 
@@ -110,3 +115,35 @@ def test_httplib_misuse(sentry_init, capture_events):
         "status_code": 200,
         "reason": "OK",
     }
+
+
+def test_outgoing_trace_headers(
+    sentry_init, monkeypatch, StringContaining  # noqa: N803
+):
+    # HTTPSConnection.send is passed a string containing (among other things)
+    # the headers on the request. Mock it so we can check the headers, and also
+    # so it doesn't try to actually talk to the internet.
+    mock_send = mock.Mock()
+    monkeypatch.setattr(HTTPSConnection, "send", mock_send)
+
+    sentry_init(traces_sample_rate=1.0)
+
+    with start_transaction(
+        name="/interactions/other-dogs/new-dog",
+        op="greeting.sniff",
+        trace_id="12312012123120121231201212312012",
+    ) as transaction:
+
+        HTTPSConnection("www.squirrelchasers.com").request("GET", "/top-chasers")
+
+        request_span = transaction._span_recorder.spans[-1]
+
+        expected_sentry_trace = (
+            "sentry-trace: {trace_id}-{parent_span_id}-{sampled}".format(
+                trace_id=transaction.trace_id,
+                parent_span_id=request_span.span_id,
+                sampled=1,
+            )
+        )
+
+        mock_send.assert_called_with(StringContaining(expected_sentry_trace))
diff --git a/tests/integrations/stdlib/test_subprocess.py b/tests/integrations/stdlib/test_subprocess.py
@@ -183,9 +183,6 @@ def test_subprocess_invalid_args(sentry_init):
     sentry_init(integrations=[StdlibIntegration()])
 
     with pytest.raises(TypeError) as excinfo:
-        subprocess.Popen()
+        subprocess.Popen(1)
 
-    if PY2:
-        assert "__init__() takes at least 2 arguments (1 given)" in str(excinfo.value)
-    else:
-        assert "missing 1 required positional argument: 'args" in str(excinfo.value)
+    assert "'int' object is not iterable" in str(excinfo.value)
diff --git a/tests/tracing/test_integration_tests.py b/tests/tracing/test_integration_tests.py
@@ -58,7 +58,7 @@ def test_continue_from_headers(sentry_init, capture_events, sampled, sample_rate
     with start_transaction(name="hi", sampled=True if sample_rate == 0 else None):
         with start_span() as old_span:
             old_span.sampled = sampled
-            headers = dict(Hub.current.iter_trace_propagation_headers())
+            headers = dict(Hub.current.iter_trace_propagation_headers(old_span))
 
     # test that the sampling decision is getting encoded in the header correctly
     header = headers["sentry-trace"]
