[Core] OpenTelemetryTracer updates (Azure#40024)

* [Core] OpenTelemetryTracer updates

- Add `end_on_exit` keyword argument to `start_as_current_span`
- Add `use_span` class method to activate a span in the current tracing
  context.

Signed-off-by: Paul Van Eck <[email protected]>

* Update docstrings

Signed-off-by: Paul Van Eck <[email protected]>

---------

Signed-off-by: Paul Van Eck <[email protected]>

Author: pvaneck

sdk/core/azure-core/CHANGELOG.md

@@ -1,6 +1,6 @@
 # Release History
 
-## 1.33.0 (2025-03-10)
+## 1.33.0 (Unreleased)
 
 ### Features Added
 
@@ -27,6 +27,8 @@
     - "x-ms-client-request-id" is now "az.client_request_id"
     - "x-ms-request-id" is now "az.service_request_id"
 
+### Bugs Fixed
+
 ### Other Changes
 
 - Added `opentelemetry-api` as an optional dependency for tracing. This can be installed with `pip install azure-core[tracing]`. #39563

sdk/core/azure-core/azure/core/tracing/opentelemetry.py

@@ -114,11 +114,10 @@ def start_as_current_span(
         kind: SpanKind = _SpanKind.INTERNAL,
         attributes: Optional[Attributes] = None,
         links: Optional[Sequence[Link]] = None,
+        end_on_exit: bool = True,
     ) -> Iterator[Span]:
         """Context manager that starts a span and sets it as the current span in the context.
 
-        Exiting the context manager will call the span's end method.
-
         .. code:: python
 
             with tracer.start_as_current_span("span_name") as span:
@@ -133,13 +132,34 @@ def start_as_current_span(
         :paramtype attributes: Optional[Attributes]
         :keyword links: Links to add to the span.
         :paramtype links: Optional[Sequence[Link]]
+        :keyword end_on_exit: Whether to end the span when exiting the context manager. Defaults to True.
+        :paramtype end_on_exit: bool
         :return: The span that was started
-        :rtype: ~opentelemetry.trace.Span
+        :rtype: Iterator[~opentelemetry.trace.Span]
         """
         span = self.start_span(name, kind=kind, attributes=attributes, links=links)
-        with trace.use_span(span, record_exception=False, end_on_exit=True) as span:  # type: ignore[attr-defined]  # pylint: disable=not-context-manager
+        with trace.use_span(  # pylint: disable=not-context-manager
+            span, record_exception=False, end_on_exit=end_on_exit
+        ) as span:
             yield span
 
+    @classmethod
+    @contextmanager
+    def use_span(cls, span: Span, *, end_on_exit: bool = True) -> Iterator[Span]:
+        """Context manager that takes a non-active span and activates it in the current context.
+
+        :param span: The span to set as the current span
+        :type span: ~opentelemetry.trace.Span
+        :keyword end_on_exit: Whether to end the span when exiting the context manager. Defaults to True.
+        :paramtype end_on_exit: bool
+        :return: The span that was activated.
+        :rtype: Iterator[~opentelemetry.trace.Span]
+        """
+        with trace.use_span(  # pylint: disable=not-context-manager
+            span, record_exception=False, end_on_exit=end_on_exit
+        ) as active_span:
+            yield active_span
+
     def _parse_links(self, links: Optional[Sequence[Link]]) -> Optional[Sequence[OpenTelemetryLink]]:
         if not links:
             return None
@@ -189,7 +209,7 @@ def call_with_current_context(*args, **kwargs):
 
     @classmethod
     def get_trace_context(cls) -> Dict[str, str]:
-        """Returns the Trace Context header values associated with the span.
+        """Returns the Trace Context header values associated with the current span.
 
         These are generally the W3C Trace Context headers (i.e. "traceparent" and "tracestate").
 

sdk/core/azure-core/tests/test_tracer_otel.py

@@ -94,6 +94,21 @@ def test_start_span_with_attributes(tracing_helper):
     assert finished_spans[0].attributes["biz"] == 123
 
 
+def test_start_as_current_span_no_exit(tracing_helper):
+    tracer = get_tracer()
+    assert tracer
+
+    with tracer.start_as_current_span(name="foo-span", kind=SpanKind.INTERNAL, end_on_exit=False) as span:
+        assert span.is_recording()
+        assert tracer.get_current_span() == span
+
+    assert span.is_recording()
+    span.end()
+    assert not span.is_recording()
+    finished_spans = tracing_helper.exporter.get_finished_spans()
+    assert len(finished_spans) == 1
+
+
 def test_tracer_get_current_span():
     """Test that the tracer can get the current OpenTelemetry span instance."""
     tracer = get_tracer()
@@ -129,6 +144,37 @@ def test_end_span(tracing_helper):
     assert finished_spans[0].name == "foo-span"
 
 
+def test_use_span(tracing_helper):
+    tracer = get_tracer()
+    assert tracer
+
+    assert isinstance(tracer.get_current_span(), NonRecordingSpan)
+    with tracer.start_as_current_span(name="root", kind=SpanKind.INTERNAL) as root:
+        span = tracer.start_span(name="foo-span", kind=SpanKind.INTERNAL)
+        assert span.is_recording()
+        assert tracer.get_current_span() == root
+
+        with tracer.use_span(span):
+            assert tracer.get_current_span() == span
+
+        assert not span.is_recording()
+
+        span2 = tracer.start_span(name="bar-span", kind=SpanKind.INTERNAL)
+        assert span2.is_recording()
+        assert tracer.get_current_span() == root
+        with tracer.use_span(span2, end_on_exit=False):
+            assert tracer.get_current_span() == span2
+            assert span2.is_recording()
+
+        assert span2.is_recording()
+        span2.end()
+        assert not span2.is_recording()
+        assert tracer.get_current_span() == root
+
+    finished_spans = tracing_helper.exporter.get_finished_spans()
+    assert len(finished_spans) == 3
+
+
 def test_get_trace_context():
     """Test that the tracer can get the trace context and it contains the traceparent header."""
     tracer = get_tracer()