Merge branch 'master' into shellmayr/cohere/1-genai-migration

Author: shellmayr

sentry_sdk/scope.py

@@ -33,7 +33,7 @@
     normalize_incoming_data,
     PropagationContext,
 )
-from sentry_sdk.traces import StreamedSpan
+from sentry_sdk.traces import _DEFAULT_PARENT_SPAN, StreamedSpan, NoOpStreamedSpan
 from sentry_sdk.tracing import (
     BAGGAGE_HEADER_NAME,
     SENTRY_TRACE_HEADER_NAME,
@@ -1174,6 +1174,59 @@ def start_span(
 
             return span
 
+    def start_streamed_span(
+        self,
+        name: str,
+        attributes: "Optional[Attributes]",
+        parent_span: "Optional[StreamedSpan]",
+        active: bool,
+    ) -> "StreamedSpan":
+        # TODO: rename to start_span once we drop the old API
+        if isinstance(parent_span, NoOpStreamedSpan):
+            # parent_span is only set if the user explicitly set it
+            logger.debug(
+                "Ignored parent span provided. Span will be parented to the "
+                "currently active span instead."
+            )
+
+        if parent_span is _DEFAULT_PARENT_SPAN or isinstance(
+            parent_span, NoOpStreamedSpan
+        ):
+            parent_span = self.span  # type: ignore
+
+        # If no eligible parent_span was provided and there is no currently
+        # active span, this is a segment
+        if parent_span is None:
+            propagation_context = self.get_active_propagation_context()
+
+            return StreamedSpan(
+                name=name,
+                attributes=attributes,
+                active=active,
+                scope=self,
+                segment=None,
+                trace_id=propagation_context.trace_id,
+                parent_span_id=propagation_context.parent_span_id,
+                parent_sampled=propagation_context.parent_sampled,
+                baggage=propagation_context.baggage,
+            )
+
+        # This is a child span; take propagation context from the parent span
+        with new_scope():
+            if isinstance(parent_span, NoOpStreamedSpan):
+                return NoOpStreamedSpan()
+
+            return StreamedSpan(
+                name=name,
+                attributes=attributes,
+                active=active,
+                scope=self,
+                segment=parent_span._segment,
+                trace_id=parent_span.trace_id,
+                parent_span_id=parent_span.span_id,
+                parent_sampled=parent_span.sampled,
+            )
+
     def continue_trace(
         self,
         environ_or_headers: "Dict[str, Any]",

sentry_sdk/traces.py

@@ -9,10 +9,12 @@
 from enum import Enum
 from typing import TYPE_CHECKING
 
+import sentry_sdk
+from sentry_sdk.tracing_utils import Baggage
 from sentry_sdk.utils import format_attribute, logger
 
 if TYPE_CHECKING:
-    from typing import Optional, Union
+    from typing import Any, Optional, Union
     from sentry_sdk._types import Attributes, AttributeValue
 
 
@@ -57,6 +59,114 @@ def __str__(self) -> str:
 }
 
 
+# Sentinel value for an unset parent_span to be able to distinguish it from
+# a None set by the user
+_DEFAULT_PARENT_SPAN = object()
+
+
+def start_span(
+    name: str,
+    attributes: "Optional[Attributes]" = None,
+    parent_span: "Optional[StreamedSpan]" = _DEFAULT_PARENT_SPAN,  # type: ignore[assignment]
+    active: bool = True,
+) -> "StreamedSpan":
+    """
+    Start a span.
+
+    The span's parent, unless provided explicitly via the `parent_span` argument,
+    will be the current active span, if any. If there is none, this span will
+    become the root of a new span tree. If you explicitly want this span to be
+    top-level without a parent, set `parent_span=None`.
+
+    `start_span()` can either be used as context manager or you can use the span
+    object it returns and explicitly end it via `span.end()`. The following is
+    equivalent:
+
+    ```python
+    import sentry_sdk
+
+    with sentry_sdk.traces.start_span(name="My Span"):
+        # do something
+
+    # The span automatically finishes once the `with` block is exited
+    ```
+
+    ```python
+    import sentry_sdk
+
+    span = sentry_sdk.traces.start_span(name="My Span")
+    # do something
+    span.end()
+    ```
+
+    To continue a trace from another service, call
+    `sentry_sdk.traces.continue_trace()` prior to creating a top-level span.
+
+    :param name: The name to identify this span by.
+    :type name: str
+
+    :param attributes: Key-value attributes to set on the span from the start.
+        These will also be accessible in the traces sampler.
+    :type attributes: "Optional[Attributes]"
+
+    :param parent_span: A span instance that the new span should consider its
+        parent. If not provided, the parent will be set to the currently active
+        span, if any. If set to `None`, this span will become a new root-level
+        span.
+    :type parent_span: "Optional[StreamedSpan]"
+
+    :param active: Controls whether spans started while this span is running
+        will automatically become its children. That's the default behavior. If
+        you want to create a span that shouldn't have any children (unless
+        provided explicitly via the `parent_span` argument), set this to `False`.
+    :type active: bool
+
+    :return: The span that has been started.
+    :rtype: StreamedSpan
+    """
+    return sentry_sdk.get_current_scope().start_streamed_span(
+        name, attributes, parent_span, active
+    )
+
+
+def continue_trace(incoming: "dict[str, Any]") -> None:
+    """
+    Continue a trace from headers or environment variables.
+
+    This function sets the propagation context on the scope. Any span started
+    in the updated scope will belong under the trace extracted from the
+    provided propagation headers or environment variables.
+
+    continue_trace() doesn't start any spans on its own. Use the start_span()
+    API for that.
+    """
+    # This is set both on the isolation and the current scope for compatibility
+    # reasons. Conceptually, it belongs on the isolation scope, and it also
+    # used to be set there in non-span-first mode. But in span first mode, we
+    # start spans on the current scope, regardless of type, like JS does, so we
+    # need to set the propagation context there.
+    sentry_sdk.get_isolation_scope().generate_propagation_context(
+        incoming,
+    )
+    sentry_sdk.get_current_scope().generate_propagation_context(
+        incoming,
+    )
+
+
+def new_trace() -> None:
+    """
+    Resets the propagation context, forcing a new trace.
+
+    This function sets the propagation context on the scope. Any span started
+    in the updated scope will start its own trace.
+
+    new_trace() doesn't start any spans on its own. Use the start_span() API
+    for that.
+    """
+    sentry_sdk.get_isolation_scope().set_new_propagation_context()
+    sentry_sdk.get_current_scope().set_new_propagation_context()
+
+
 class StreamedSpan:
     """
     A span holds timing information of a block of code.
@@ -73,7 +183,12 @@ class StreamedSpan:
         "_active",
         "_span_id",
         "_trace_id",
+        "_parent_span_id",
+        "_segment",
+        "_parent_sampled",
         "_status",
+        "_scope",
+        "_baggage",
     )
 
     def __init__(
@@ -82,7 +197,12 @@ def __init__(
         name: str,
         attributes: "Optional[Attributes]" = None,
         active: bool = True,
+        scope: "sentry_sdk.Scope",
+        segment: "Optional[StreamedSpan]" = None,
         trace_id: "Optional[str]" = None,
+        parent_span_id: "Optional[str]" = None,
+        parent_sampled: "Optional[bool]" = None,
+        baggage: "Optional[Baggage]" = None,
     ):
         self._name: str = name
         self._active: bool = active
@@ -91,8 +211,16 @@ def __init__(
             for attribute, value in attributes.items():
                 self.set_attribute(attribute, value)
 
-        self._span_id: "Optional[str]" = None
+        self._scope = scope
+
+        self._segment = segment or self
+
         self._trace_id: "Optional[str]" = trace_id
+        self._parent_span_id = parent_span_id
+        self._parent_sampled = parent_sampled
+        self._baggage = baggage
+
+        self._span_id: "Optional[str]" = None
 
         self._status = SpanStatus.OK.value
         self.set_attribute("sentry.span.source", SegmentSource.CUSTOM.value)
@@ -103,6 +231,7 @@ def __repr__(self) -> str:
             f"name={self._name}, "
             f"trace_id={self.trace_id}, "
             f"span_id={self.span_id}, "
+            f"parent_span_id={self._parent_span_id}, "
             f"active={self._active})>"
         )
 
@@ -165,8 +294,18 @@ def trace_id(self) -> str:
 
         return self._trace_id
 
+    @property
+    def sampled(self) -> "Optional[bool]":
+        return True
+
 
 class NoOpStreamedSpan(StreamedSpan):
+    def __init__(self) -> None:
+        pass
+
+    def __repr__(self) -> str:
+        return f"<{self.__class__.__name__}(sampled={self.sampled})>"
+
     def get_attributes(self) -> "Attributes":
         return {}
 
@@ -206,3 +345,7 @@ def span_id(self) -> str:
     @property
     def trace_id(self) -> str:
         return "00000000000000000000000000000000"
+
+    @property
+    def sampled(self) -> "Optional[bool]":
+        return False