Implement new Propagator.inject for OTLPIntegration

sentry_sdk/integrations/otlp.py

@@ -1,16 +1,45 @@
+from sentry_sdk import get_client
 from sentry_sdk.integrations import Integration, DidNotEnable
 from sentry_sdk.scope import register_external_propagation_context
 from sentry_sdk.utils import logger, Dsn
 from sentry_sdk.consts import VERSION, EndpointType
+from sentry_sdk.tracing_utils import Baggage
+from sentry_sdk.tracing import (
+    BAGGAGE_HEADER_NAME,
+    SENTRY_TRACE_HEADER_NAME,
+)
 
 try:
-    from opentelemetry import trace
     from opentelemetry.propagate import set_global_textmap
     from opentelemetry.sdk.trace import TracerProvider
     from opentelemetry.sdk.trace.export import BatchSpanProcessor
     from opentelemetry.exporter.otlp.proto.http.trace_exporter import OTLPSpanExporter
 
+    from opentelemetry.trace import (
+        get_current_span,
+        get_tracer_provider,
+        set_tracer_provider,
+        format_trace_id,
+        format_span_id,
+        SpanContext,
+        INVALID_SPAN_ID,
+        INVALID_TRACE_ID,
+    )
+
+    from opentelemetry.context import (
+        Context,
+        get_current,
+        get_value,
+    )
+
+    from opentelemetry.propagators.textmap import (
+        CarrierT,
+        Setter,
+        default_setter,
+    )
+
     from sentry_sdk.integrations.opentelemetry.propagator import SentryPropagator
+    from sentry_sdk.integrations.opentelemetry.consts import SENTRY_BAGGAGE_KEY
 except ImportError:
     raise DidNotEnable("opentelemetry-distro[otlp] is not installed")
 
@@ -25,22 +54,22 @@ def otel_propagation_context():
     """
     Get the (trace_id, span_id) from opentelemetry if exists.
     """
-    ctx = trace.get_current_span().get_span_context()
+    ctx = get_current_span().get_span_context()
 
-    if ctx.trace_id == trace.INVALID_TRACE_ID or ctx.span_id == trace.INVALID_SPAN_ID:
+    if ctx.trace_id == INVALID_TRACE_ID or ctx.span_id == INVALID_SPAN_ID:
         return None
 
-    return (trace.format_trace_id(ctx.trace_id), trace.format_span_id(ctx.span_id))
+    return (format_trace_id(ctx.trace_id), format_span_id(ctx.span_id))
 
 
 def setup_otlp_traces_exporter(dsn=None):
     # type: (Optional[str]) -> None
-    tracer_provider = trace.get_tracer_provider()
+    tracer_provider = get_tracer_provider()
 
     if not isinstance(tracer_provider, TracerProvider):
         logger.debug("[OTLP] No TracerProvider configured by user, creating a new one")
         tracer_provider = TracerProvider()
-        trace.set_tracer_provider(tracer_provider)
+        set_tracer_provider(tracer_provider)
 
     endpoint = None
     headers = None
@@ -55,6 +84,56 @@ def setup_otlp_traces_exporter(dsn=None):
     tracer_provider.add_span_processor(span_processor)
 
 
+class SentryOTLPPropagator(SentryPropagator):
+    """
+    We need to override the inject of the older propagator since that
+    is SpanProcessor based.
+
+    !!! Note regarding baggage:
+    We cannot meaningfully populate a new baggage as a head SDK
+    when we are using OTLP since we don't have any sort of transaction semantic to
+    track state across a group of spans.
+
+    For incoming baggage, we just pass it on as is so that case is correctly handled.
+    """
+
+    def inject(self, carrier, context=None, setter=default_setter):
+        # type: (CarrierT, Optional[Context], Setter[CarrierT]) -> None
+        otlp_integration = get_client().get_integration(OTLPIntegration)
+        if otlp_integration is None:
+            return
+
+        if context is None:
+            context = get_current()
+
+        current_span = get_current_span(context)
+        current_span_context = current_span.get_span_context()
+
+        if not current_span_context.is_valid:
+            return
+
+        sentry_trace = _to_traceparent(current_span_context)
+        setter.set(carrier, SENTRY_TRACE_HEADER_NAME, sentry_trace)
+
+        baggage = get_value(SENTRY_BAGGAGE_KEY, context)
+        if baggage is not None and isinstance(baggage, Baggage):
+            baggage_data = baggage.serialize()
+            if baggage_data:
+                setter.set(carrier, BAGGAGE_HEADER_NAME, baggage.serialize())
+
+
+def _to_traceparent(span_context):
+    # type: (SpanContext) -> str
+    """
+    Helper method to generate the sentry-trace header.
+    """
+    span_id = format_span_id(span_context.span_id)
+    trace_id = format_trace_id(span_context.trace_id)
+    sampled = span_context.trace_flags.sampled
+
+    return f"{trace_id}-{span_id}-{'1' if sampled else '0'}"
+
+
 class OTLPIntegration(Integration):
     identifier = "otlp"
 
@@ -79,4 +158,4 @@ def setup_once_with_options(self, options=None):
         if self.setup_propagator:
             logger.debug("[OTLP] Setting up propagator for distributed tracing")
             # TODO-neel better propagator support, chain with existing ones if possible instead of replacing
-            set_global_textmap(SentryPropagator())
+            set_global_textmap(SentryOTLPPropagator())

sentry_sdk/scope.py

@@ -165,6 +165,11 @@ def get_external_propagation_context():
     )
 
 
+def has_external_propagation_context():
+    # type: () -> bool
+    return _external_propagation_context_fn is not None
+
+
 def _attr_setter(fn):
     # type: (Any) -> Any
     return property(fset=fn, doc=fn.__doc__)
@@ -637,6 +642,10 @@ def iter_trace_propagation_headers(self, *args, **kwargs):
         if has_tracing_enabled(client.options) and span is not None:
             for header in span.iter_headers():
                 yield header
+        elif has_external_propagation_context():
+            # when we have an external_propagation_context (otlp)
+            # we leave outgoing propagation to the propagator
+            return
         else:
             for header in self.get_active_propagation_context().iter_headers():
                 yield header

tests/integrations/otlp/test_otlp.py

@@ -8,15 +8,16 @@
     ProxyTracerProvider,
     format_span_id,
     format_trace_id,
+    get_current_span,
 )
+from opentelemetry.context import attach, detach
 from opentelemetry.propagate import get_global_textmap, set_global_textmap
 from opentelemetry.util._once import Once
 from opentelemetry.sdk.trace import TracerProvider
 from opentelemetry.sdk.trace.export import BatchSpanProcessor
 from opentelemetry.exporter.otlp.proto.http.trace_exporter import OTLPSpanExporter
 
-from sentry_sdk.integrations.otlp import OTLPIntegration
-from sentry_sdk.integrations.opentelemetry import SentryPropagator
+from sentry_sdk.integrations.otlp import OTLPIntegration, SentryOTLPPropagator
 from sentry_sdk.scope import get_external_propagation_context
 
 
@@ -111,7 +112,7 @@ def test_sets_propagator(sentry_init):
     )
 
     propagator = get_global_textmap()
-    assert isinstance(get_global_textmap(), SentryPropagator)
+    assert isinstance(get_global_textmap(), SentryOTLPPropagator)
     assert propagator is not original_propagator
 
 
@@ -122,7 +123,7 @@ def test_does_not_set_propagator_if_disabled(sentry_init):
     )
 
     propagator = get_global_textmap()
-    assert not isinstance(propagator, SentryPropagator)
+    assert not isinstance(propagator, SentryOTLPPropagator)
     assert propagator is original_propagator
 
 
@@ -152,3 +153,72 @@ def test_otel_propagation_context(sentry_init):
     assert trace_id == format_trace_id(root_span.get_span_context().trace_id)
     assert trace_id == format_trace_id(span.get_span_context().trace_id)
     assert span_id == format_span_id(span.get_span_context().span_id)
+
+
+def test_propagator_inject_head_of_trace(sentry_init):
+    sentry_init(
+        dsn="https://mysecret@bla.ingest.sentry.io/12312012",
+        integrations=[OTLPIntegration()],
+    )
+
+    tracer = trace.get_tracer(__name__)
+    propagator = get_global_textmap()
+    carrier = {}
+
+    with tracer.start_as_current_span("foo") as span:
+        propagator.inject(carrier)
+
+        span_context = span.get_span_context()
+        trace_id = format_trace_id(span_context.trace_id)
+        span_id = format_span_id(span_context.span_id)
+
+        assert "sentry-trace" in carrier
+        assert carrier["sentry-trace"] == f"{trace_id}-{span_id}-1"
+
+        #! we cannot populate baggage in otlp as head SDK yet
+        assert "baggage" not in carrier
+
+
+def test_propagator_inject_continue_trace(sentry_init):
+    sentry_init(
+        dsn="https://mysecret@bla.ingest.sentry.io/12312012",
+        integrations=[OTLPIntegration()],
+    )
+
+    tracer = trace.get_tracer(__name__)
+    propagator = get_global_textmap()
+    carrier = {}
+
+    incoming_headers = {
+        "sentry-trace": "771a43a4192642f0b136d5159a501700-1234567890abcdef-1",
+        "baggage": (
+            "sentry-trace_id=771a43a4192642f0b136d5159a501700,sentry-sampled=true"
+        ),
+    }
+
+    ctx = propagator.extract(incoming_headers)
+    token = attach(ctx)
+
+    parent_span_context = get_current_span().get_span_context()
+    assert (
+        format_trace_id(parent_span_context.trace_id)
+        == "771a43a4192642f0b136d5159a501700"
+    )
+    assert format_span_id(parent_span_context.span_id) == "1234567890abcdef"
+
+    with tracer.start_as_current_span("foo") as span:
+        propagator.inject(carrier)
+
+        span_context = span.get_span_context()
+        trace_id = format_trace_id(span_context.trace_id)
+        span_id = format_span_id(span_context.span_id)
+
+        assert trace_id == "771a43a4192642f0b136d5159a501700"
+
+        assert "sentry-trace" in carrier
+        assert carrier["sentry-trace"] == f"{trace_id}-{span_id}-1"
+
+        assert "baggage" in carrier
+        assert carrier["baggage"] == incoming_headers["baggage"]
+
+    detach(token)

tests/test_scope.py

@@ -997,11 +997,15 @@ def external_propagation_context():
 
     register_external_propagation_context(external_propagation_context)
 
-    trace_context = sentry_sdk.get_current_scope().get_trace_context()
+    scope = sentry_sdk.get_current_scope()
 
+    trace_context = scope.get_trace_context()
     assert trace_context["trace_id"] == "trace_id_foo"
     assert trace_context["span_id"] == "span_id_bar"
 
+    headers = list(scope.iter_trace_propagation_headers())
+    assert not headers
+
     remove_external_propagation_context()