feat(stdlib): Support span streaming

Author: alexander-alderman-webb

sentry_sdk/integrations/stdlib.py

@@ -8,10 +8,13 @@
 from sentry_sdk.consts import OP, SPANDATA
 from sentry_sdk.integrations import Integration
 from sentry_sdk.scope import add_global_event_processor
+from sentry_sdk.tracing import Span
+from sentry_sdk.traces import StreamedSpan
 from sentry_sdk.tracing_utils import (
     EnvironHeaders,
     should_propagate_trace,
     add_http_request_source,
+    has_span_streaming_enabled,
 )
 from sentry_sdk.utils import (
     SENSITIVE_DATA_SUBSTITUTE,
@@ -31,6 +34,7 @@
     from typing import Dict
     from typing import Optional
     from typing import List
+    from typing import Union
 
     from sentry_sdk._types import Event, Hint
 
@@ -99,22 +103,38 @@ def putrequest(
         with capture_internal_exceptions():
             parsed_url = parse_url(real_url, sanitize=False)
 
-        span = sentry_sdk.start_span(
-            op=OP.HTTP_CLIENT,
-            name="%s %s"
-            % (method, parsed_url.url if parsed_url else SENSITIVE_DATA_SUBSTITUTE),
-            origin="auto.http.stdlib.httplib",
-        )
-        span.set_data(SPANDATA.HTTP_METHOD, method)
+        span_streaming = has_span_streaming_enabled(client.options)
+        span: "Union[Span, StreamedSpan]"
+        if span_streaming:
+            span = sentry_sdk.traces.start_span(
+                name="%s %s"
+                % (method, parsed_url.url if parsed_url else SENSITIVE_DATA_SUBSTITUTE),
+                attributes={
+                    "sentry.origin": "auto.http.stdlib.httplib",
+                    "sentry.op": OP.HTTP_CLIENT,
+                },
+            )
+            set_on_span = span.set_attribute
+
+        else:
+            span = sentry_sdk.start_span(
+                op=OP.HTTP_CLIENT,
+                name="%s %s"
+                % (method, parsed_url.url if parsed_url else SENSITIVE_DATA_SUBSTITUTE),
+                origin="auto.http.stdlib.httplib",
+            )
+            set_on_span = span.set_data
+
+        set_on_span(SPANDATA.HTTP_METHOD, method)
         if parsed_url is not None:
-            span.set_data("url", parsed_url.url)
-            span.set_data(SPANDATA.HTTP_QUERY, parsed_url.query)
-            span.set_data(SPANDATA.HTTP_FRAGMENT, parsed_url.fragment)
+            set_on_span("url", parsed_url.url)
+            set_on_span(SPANDATA.HTTP_QUERY, parsed_url.query)
+            set_on_span(SPANDATA.HTTP_FRAGMENT, parsed_url.fragment)
 
         # for proxies, these point to the proxy host/port
         if tunnel_host:
-            span.set_data(SPANDATA.NETWORK_PEER_ADDRESS, self.host)
-            span.set_data(SPANDATA.NETWORK_PEER_PORT, self.port)
+            set_on_span(SPANDATA.NETWORK_PEER_ADDRESS, self.host)
+            set_on_span(SPANDATA.NETWORK_PEER_PORT, self.port)
 
         rv = real_putrequest(self, method, url, *args, **kwargs)
 
@@ -139,14 +159,23 @@ def putrequest(
     def getresponse(self: "HTTPConnection", *args: "Any", **kwargs: "Any") -> "Any":
         span = getattr(self, "_sentrysdk_span", None)
 
+        print("span is here")
+
         if span is None:
             return real_getresponse(self, *args, **kwargs)
 
         try:
             rv = real_getresponse(self, *args, **kwargs)
 
-            span.set_http_status(int(rv.status))
-            span.set_data("reason", rv.reason)
+            if isinstance(span, StreamedSpan):
+                span.set_attribute("reason", rv.reason)
+
+                status_code = int(rv.status)
+                span.status = "error" if status_code >= 400 else "ok"
+                span.set_attribute("http.response.status_code", status_code)
+            else:
+                span.set_http_status(int(rv.status))
+                span.set_data("reason", rv.reason)
         finally:
             span.finish()
 
@@ -226,11 +255,22 @@ def sentry_patched_popen_init(
 
         env = None
 
-        with sentry_sdk.start_span(
-            op=OP.SUBPROCESS,
-            name=description,
-            origin="auto.subprocess.stdlib.subprocess",
-        ) as span:
+        span_streaming = has_span_streaming_enabled(sentry_sdk.get_client().options)
+        span: "Union[Span, StreamedSpan]"
+        if span_streaming:
+            span = sentry_sdk.start_span(
+                op=OP.SUBPROCESS,
+                name=description,
+                origin="auto.subprocess.stdlib.subprocess",
+            )
+        else:
+            span = sentry_sdk.start_span(
+                op=OP.SUBPROCESS,
+                name=description,
+                origin="auto.subprocess.stdlib.subprocess",
+            )
+
+        with span:
             for k, v in sentry_sdk.get_current_scope().iter_trace_propagation_headers(
                 span=span
             ):

tests/integrations/stdlib/test_httplib.py

@@ -202,14 +202,15 @@ def test_httplib_misuse(sentry_init, capture_events, request):
     )
 
 
-def test_outgoing_trace_headers(sentry_init, monkeypatch):
+def test_outgoing_trace_headers(sentry_init, capture_events, monkeypatch):
     # 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)
+    events = capture_events()
 
     headers = {
         "sentry-trace": "771a43a4192642f0b136d5159a501700-1234567890abcdef-1",
@@ -237,38 +238,42 @@ def test_outgoing_trace_headers(sentry_init, monkeypatch):
                 key, val = line.split(": ")
                 request_headers[key] = val
 
-        request_span = transaction._span_recorder.spans[-1]
-        expected_sentry_trace = "{trace_id}-{parent_span_id}-{sampled}".format(
-            trace_id=transaction.trace_id,
-            parent_span_id=request_span.span_id,
-            sampled=1,
-        )
-        assert request_headers["sentry-trace"] == expected_sentry_trace
-
-        expected_outgoing_baggage = (
-            "sentry-trace_id=771a43a4192642f0b136d5159a501700,"
-            "sentry-public_key=49d0f7386ad645858ae85020e393bef3,"
-            "sentry-sample_rate=1.0,"
-            "sentry-user_id=Am%C3%A9lie,"
-            "sentry-sample_rand=0.132521102938283"
-        )
+    (event,) = events
+    request_span = event["spans"][-1]
+    expected_sentry_trace = "{trace_id}-{parent_span_id}-{sampled}".format(
+        trace_id=transaction.trace_id,
+        parent_span_id=request_span.span_id,
+        sampled=1,
+    )
+    assert request_headers["sentry-trace"] == expected_sentry_trace
+
+    expected_outgoing_baggage = (
+        "sentry-trace_id=771a43a4192642f0b136d5159a501700,"
+        "sentry-public_key=49d0f7386ad645858ae85020e393bef3,"
+        "sentry-sample_rate=1.0,"
+        "sentry-user_id=Am%C3%A9lie,"
+        "sentry-sample_rand=0.132521102938283"
+    )
 
-        assert request_headers["baggage"] == expected_outgoing_baggage
+    assert request_headers["baggage"] == expected_outgoing_baggage
 
 
-def test_outgoing_trace_headers_head_sdk(sentry_init, monkeypatch):
+def test_outgoing_trace_headers_head_sdk(sentry_init, capture_events, monkeypatch):
     # 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=0.5, release="foo")
+    events = capture_events()
+
     with mock.patch("sentry_sdk.tracing_utils.Random.randrange", return_value=250000):
         transaction = continue_trace({})
 
     with start_transaction(transaction=transaction, name="Head SDK tx") as transaction:
-        HTTPSConnection("www.squirrelchasers.com").request("GET", "/top-chasers")
+        connection = HTTPSConnection("www.squirrelchasers.com")
+        connection.request("GET", "/top-chasers")
 
         (request_str,) = mock_send.call_args[0]
         request_headers = {}
@@ -277,24 +282,25 @@ def test_outgoing_trace_headers_head_sdk(sentry_init, monkeypatch):
                 key, val = line.split(": ")
                 request_headers[key] = val
 
-        request_span = transaction._span_recorder.spans[-1]
-        expected_sentry_trace = "{trace_id}-{parent_span_id}-{sampled}".format(
-            trace_id=transaction.trace_id,
-            parent_span_id=request_span.span_id,
-            sampled=1,
-        )
-        assert request_headers["sentry-trace"] == expected_sentry_trace
-
-        expected_outgoing_baggage = (
-            "sentry-trace_id=%s,"
-            "sentry-sample_rand=0.250000,"
-            "sentry-environment=production,"
-            "sentry-release=foo,"
-            "sentry-sample_rate=0.5,"
-            "sentry-sampled=%s"
-        ) % (transaction.trace_id, "true" if transaction.sampled else "false")
-
-        assert request_headers["baggage"] == expected_outgoing_baggage
+    (event,) = events
+    request_span = event["spans"][0]
+    expected_sentry_trace = "{trace_id}-{parent_span_id}-{sampled}".format(
+        trace_id=transaction.trace_id,
+        parent_span_id=request_span.span_id,
+        sampled=1,
+    )
+    assert request_headers["sentry-trace"] == expected_sentry_trace
+
+    expected_outgoing_baggage = (
+        "sentry-trace_id=%s,"
+        "sentry-sample_rand=0.250000,"
+        "sentry-environment=production,"
+        "sentry-release=foo,"
+        "sentry-sample_rate=0.5,"
+        "sentry-sampled=%s"
+    ) % (transaction.trace_id, "true" if transaction.sampled else "false")
+
+    assert request_headers["baggage"] == expected_outgoing_baggage
 
 
 @pytest.mark.parametrize(