stdlib

Author: sentrivana

sentry_sdk/integrations/starlette.py

@@ -170,44 +170,37 @@ async def _create_span_call(
         middleware_name = app.__class__.__name__
 
         if span_streaming:
-            span_ctx = sentry_sdk.traces.start_span(name=middleware_name)
+            middleware_span = sentry_sdk.traces.start_span(name=middleware_name)
+            middleware_span.set_op(OP.MIDDLEWARE_STARLETTE)
+            middleware_span.set_origin(StarletteIntegration.origin)
+            middleware_span.set_attribute("starlette.middleware_name", middleware_name)
         else:
-            span_ctx = sentry_sdk.start_span(
+            middleware_span = sentry_sdk.start_span(
                 op=OP.MIDDLEWARE_STARLETTE,
                 name=middleware_name,
                 origin=StarletteIntegration.origin,
             )
+            middleware_span.set_tag("starlette.middleware_name", middleware_name)
 
-        with span_ctx as middleware_span:
-            if span_streaming:
-                middleware_span.set_op(OP.MIDDLEWARE_STARLETTE)
-                middleware_span.set_origin(StarletteIntegration.origin)
-                middleware_span.set_attribute(
-                    "starlette.middleware_name", middleware_name
-                )
-            else:
-                middleware_span.set_tag("starlette.middleware_name", middleware_name)
-
+        with middleware_span:
             # Creating spans for the "receive" callback
             async def _sentry_receive(*args: "Any", **kwargs: "Any") -> "Any":
                 if span_streaming:
-                    span_ctx = sentry_sdk.tracing.start_span(
+                    span = sentry_sdk.traces.start_span(
                         name=getattr(receive, "__qualname__", str(receive)),
                     )
+                    span.set_origin(StarletteIntegration.origin)
+                    span.set_op(OP.MIDDLEWARE_STARLETTE_RECEIVE)
+                    span.set_attribute("starlette.middleware_name", middleware_name)
                 else:
-                    span_ctx = sentry_sdk.start_span(
+                    span = sentry_sdk.start_span(
                         op=OP.MIDDLEWARE_STARLETTE_RECEIVE,
                         name=getattr(receive, "__qualname__", str(receive)),
                         origin=StarletteIntegration.origin,
                     )
+                    span.set_tag("starlette.middleware_name", middleware_name)
 
-                with span_ctx as span:
-                    if span_streaming:
-                        span.set_origin(StarletteIntegration.origin)
-                        span.set_op(OP.MIDDLEWARE_STARLETTE_RECEIVE)
-                        span.set_attribute("starlette.middleware_name", middleware_name)
-                    else:
-                        span.set_tag("starlette.middleware_name", middleware_name)
+                with span:
                     return await receive(*args, **kwargs)
 
             receive_name = getattr(receive, "__name__", str(receive))
@@ -217,24 +210,21 @@ async def _sentry_receive(*args: "Any", **kwargs: "Any") -> "Any":
             # Creating spans for the "send" callback
             async def _sentry_send(*args: "Any", **kwargs: "Any") -> "Any":
                 if span_streaming:
-                    span_ctx = sentry_sdk.tracing.start_span(
+                    span = sentry_sdk.traces.start_span(
                         name=getattr(send, "__qualname__", str(send)),
                     )
+                    span.set_op(OP.MIDDLEWARE_STARLETTE_SEND)
+                    span.set_origin(StarletteIntegration.origin)
+                    span.set_attribute("starlette.middleware_name", middleware_name)
                 else:
-                    span_ctx = sentry_sdk.start_span(
+                    span = sentry_sdk.start_span(
                         op=OP.MIDDLEWARE_STARLETTE_SEND,
                         name=getattr(send, "__qualname__", str(send)),
                         origin=StarletteIntegration.origin,
                     )
+                    span.set_tag("starlette.middleware_name", middleware_name)
 
-                with span_ctx as span:
-                    if span_streaming:
-                        span.set_op(OP.MIDDLEWARE_STARLETTE_SEND)
-                        span.set_origin(StarletteIntegration.origin)
-                        span.set_attribute("starlette.middleware_name", middleware_name)
-                    else:
-                        span.set_tag("starlette.middleware_name", middleware_name)
-
+                with span:
                     return await send(*args, **kwargs)
 
             send_name = getattr(send, "__name__", str(send))

sentry_sdk/integrations/stdlib.py

@@ -8,10 +8,12 @@
 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.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,
@@ -86,6 +88,8 @@ def putrequest(
         ):
             return real_putrequest(self, method, url, *args, **kwargs)
 
+        span_streaming = has_span_streaming_enabled(client.options)
+
         real_url = url
         if real_url is None or not real_url.startswith(("http://", "https://")):
             real_url = "%s://%s%s%s" % (
@@ -99,22 +103,44 @@ 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)
-        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)
-
-        # 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)
+        if span_streaming:
+            span = sentry_sdk.traces.start_span(
+                name=f"{method} {parsed_url.url if parsed_url else SENSITIVE_DATA_SUBSTITUTE}"
+            )
+            span.set_op(OP.HTTP_CLIENT)
+            span.set_origin("auto.http.stdlib.httplib")
+
+            span.set_attribute(SPANDATA.HTTP_METHOD, method)
+            if parsed_url is not None:
+                span.set_attribute("url", parsed_url.url)
+                span.set_attribute(SPANDATA.HTTP_QUERY, parsed_url.query)
+                span.set_attribute(SPANDATA.HTTP_FRAGMENT, parsed_url.fragment)
+
+            # for proxies, these point to the proxy host/port
+            if tunnel_host:
+                span.set_attribute(SPANDATA.NETWORK_PEER_ADDRESS, self.host)
+                span.set_attribute(SPANDATA.NETWORK_PEER_PORT, self.port)
+
+            span.start()
+
+        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",
+            )
+
+            span.set_data(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)
+
+            # 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)
 
         rv = real_putrequest(self, method, url, *args, **kwargs)
 
@@ -146,7 +172,10 @@ def getresponse(self: "HTTPConnection", *args: "Any", **kwargs: "Any") -> "Any":
             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)
+            else:
+                span.set_data("reason", rv.reason)
         finally:
             span.finish()
 
@@ -226,11 +255,23 @@ def sentry_patched_popen_init(
 
         env = None
 
-        with sentry_sdk.start_span(
-            op=OP.SUBPROCESS,
-            name=description,
-            origin="auto.subprocess.stdlib.subprocess",
-        ) as span:
+        client = sentry_sdk.get_client()
+        span_streaming = has_span_streaming_enabled(client.options)
+
+        if span_streaming:
+            span = sentry_sdk.traces.start_span(
+                name=description,
+            )
+            span.set_op(OP.SUBPROCESS)
+            span.set_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
             ):
@@ -245,11 +286,18 @@ def sentry_patched_popen_init(
                 env["SUBPROCESS_" + k.upper().replace("-", "_")] = v
 
             if cwd:
-                span.set_data("subprocess.cwd", cwd)
+                if span_streaming:
+                    span.set_attribute("subprocess.cwd", cwd)
+                else:
+                    span.set_data("subprocess.cwd", cwd)
 
             rv = old_popen_init(self, *a, **kw)
 
-            span.set_tag("subprocess.pid", self.pid)
+            if span_streaming:
+                span.set_attribute("subprocess.pid", self.pid)
+            else:
+                span.set_tag("subprocess.pid", self.pid)
+
             return rv
 
     subprocess.Popen.__init__ = sentry_patched_popen_init  # type: ignore
@@ -260,11 +308,22 @@ def sentry_patched_popen_init(
     def sentry_patched_popen_wait(
         self: "subprocess.Popen[Any]", *a: "Any", **kw: "Any"
     ) -> "Any":
-        with sentry_sdk.start_span(
-            op=OP.SUBPROCESS_WAIT,
-            origin="auto.subprocess.stdlib.subprocess",
-        ) as span:
+        client = sentry_sdk.get_client()
+        span_streaming = has_span_streaming_enabled(client.options)
+
+        if span_streaming:
+            span = sentry_sdk.traces.start_span(name="subprocess popen")
+            span.set_op(OP.SUBPROCESS_WAIT)
+            span.set_origin("auto.subprocess.stdlib.subprocess")
+            span.set_tag("subprocess.pid", self.pid)
+        else:
+            span = sentry_sdk.start_span(
+                op=OP.SUBPROCESS_WAIT,
+                origin="auto.subprocess.stdlib.subprocess",
+            )
             span.set_tag("subprocess.pid", self.pid)
+
+        with span:
             return old_popen_wait(self, *a, **kw)
 
     subprocess.Popen.wait = sentry_patched_popen_wait  # type: ignore
@@ -275,11 +334,24 @@ def sentry_patched_popen_wait(
     def sentry_patched_popen_communicate(
         self: "subprocess.Popen[Any]", *a: "Any", **kw: "Any"
     ) -> "Any":
-        with sentry_sdk.start_span(
-            op=OP.SUBPROCESS_COMMUNICATE,
-            origin="auto.subprocess.stdlib.subprocess",
-        ) as span:
+        client = sentry_sdk.get_client()
+        span_streaming = has_span_streaming_enabled(client.options)
+
+        if span_streaming:
+            span = sentry_sdk.traces.start_span(
+                name="subprocess communicate",
+            )
+            span.set_op(OP.SUBPROCESS_COMMUNICATE)
+            span.set_origin("auto.subprocess.stdlib.subprocess")
             span.set_tag("subprocess.pid", self.pid)
+        else:
+            span = sentry_sdk.start_span(
+                op=OP.SUBPROCESS_COMMUNICATE,
+                origin="auto.subprocess.stdlib.subprocess",
+            )
+            span.set_tag("subprocess.pid", self.pid)
+
+        with span:
             return old_popen_communicate(self, *a, **kw)
 
     subprocess.Popen.communicate = sentry_patched_popen_communicate  # type: ignore

sentry_sdk/traces.py

@@ -30,7 +30,7 @@
 )
 
 if TYPE_CHECKING:
-    from typing import Any, Callable, Optional, ParamSpec, TypeVar, Union
+    from typing import Any, Callable, Iterator, Optional, ParamSpec, TypeVar, Union
     from sentry_sdk._types import Attributes, AttributeValue, SamplingContext
     from sentry_sdk.scope import Scope
 
@@ -40,6 +40,9 @@
 
 FLAGS_CAPACITY = 10
 
+BAGGAGE_HEADER_NAME = "baggage"
+SENTRY_TRACE_HEADER_NAME = "sentry-trace"
+
 
 class SpanStatus(str, Enum):
     OK = "ok"
@@ -362,7 +365,7 @@ def start(self):
         """
         self.__enter__()
 
-    def end(self, end_timestamp: "Optional[Union[float, datetime]]" = None) -> None:
+    def finish(self, end_timestamp: "Optional[Union[float, datetime]]" = None) -> None:
         """
         Finish this span and queue it for sending.
 
@@ -455,18 +458,12 @@ def set_flag(self, flag: str, result: bool) -> None:
             self._flags[flag] = result
 
     def set_op(self, op: "Union[OP, str]") -> None:
-        if isinstance(op, Enum):
-            op = op.value
-
         self.set_attribute("sentry.op", op)
 
     def set_origin(self, origin: str) -> None:
         self.set_attribute("sentry.origin", origin)
 
-    def set_source(self, source: "Union[SegmentSource, str]") -> None:
-        if isinstance(source, Enum):
-            source = source.value
-
+    def set_source(self, source: "SegmentSource") -> None:
         self.set_attribute("sentry.span.source", source)
 
     def is_segment(self) -> bool:
@@ -497,7 +494,7 @@ def sampled(self) -> "Optional[bool]":
         return self._sampled
 
     def dynamic_sampling_context(self) -> "dict[str, str]":
-        return self.segment._get_baggage().dynamic_sampling_context()
+        return self.segment.get_baggage().dynamic_sampling_context()
 
     def to_traceparent(self) -> str:
         if self.sampled is True:
@@ -515,9 +512,19 @@ def to_traceparent(self) -> str:
 
     def to_baggage(self) -> "Optional[Baggage]":
         if self.segment:
-            return self.segment._get_baggage()
+            return self.segment.get_baggage()
         return None
 
+    def iter_headers(self) -> "Iterator[tuple[str, str]]":
+        if not self.segment:
+            return
+
+        yield SENTRY_TRACE_HEADER_NAME, self.to_traceparent()
+
+        baggage = self.segment.get_baggage().serialize()
+        if baggage:
+            yield BAGGAGE_HEADER_NAME, baggage
+
     def _update_active_thread(self) -> None:
         thread_id, thread_name = get_current_thread_meta()
         self._set_thread(thread_id, thread_name)
@@ -543,7 +550,7 @@ def set_http_status(self, http_status: int) -> None:
         else:
             self.set_status(SpanStatus.OK)
 
-    def _get_baggage(self) -> "Baggage":
+    def get_baggage(self) -> "Baggage":
         """
         Return the :py:class:`~sentry_sdk.tracing_utils.Baggage` associated with
         the segment.

sentry_sdk/utils.py

@@ -14,6 +14,7 @@
 from collections import namedtuple
 from datetime import datetime, timezone
 from decimal import Decimal
+from enum import Enum
 from functools import partial, partialmethod, wraps
 from numbers import Real
 from urllib.parse import parse_qs, unquote, urlencode, urlsplit, urlunsplit
@@ -2059,6 +2060,9 @@ def format_attribute(val: "Any") -> "AttributeValue":
     they're serialized further into the actual format the protocol expects:
     https://develop.sentry.dev/sdk/telemetry/attributes/
     """
+    if isinstance(val, Enum):
+        val = val.value
+
     if isinstance(val, (bool, int, float, str)):
         return val