[Core] Add context kwarg to start_span methods (Azure#41511)

Enables setting the parent span to something other than the current context.

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

Author: pvaneck

sdk/core/azure-core/CHANGELOG.md

@@ -5,6 +5,7 @@
 ### Features Added
 
 - Added a `start_time` keyword argument to the `start_span` and `start_as_current_span` methods in the `OpenTelemetryTracer` class. This allows users to specify a custom start time for created spans. #41106
+- Added a `context` keyword argument to the `start_span` and `start_as_current_span` methods in the `OpenTelemetryTracer` class. This allows users to specify a custom parent context for created spans. #41511
 - Added `is_generated_model` method to `azure.core.serialization`. Returns whether a given input is a model from one of our generated sdks. #41445
 
 ### Breaking Changes

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

@@ -5,7 +5,7 @@
 from __future__ import annotations
 from contextlib import contextmanager
 from contextvars import Token
-from typing import Optional, Dict, Sequence, cast, Callable, Iterator, TYPE_CHECKING
+from typing import Any, Optional, Dict, Sequence, cast, Callable, Iterator, TYPE_CHECKING
 
 from opentelemetry import context as otel_context_module, trace
 from opentelemetry.trace import (
@@ -82,6 +82,7 @@ def start_span(
         attributes: Optional[Attributes] = None,
         links: Optional[Sequence[Link]] = None,
         start_time: Optional[int] = None,
+        context: Optional[Dict[str, Any]] = None,
     ) -> Span:
         """Starts a span without setting it as the current span in the context.
 
@@ -95,14 +96,22 @@ def start_span(
         :paramtype links: list[~azure.core.tracing.Link]
         :keyword start_time: The start time of the span in nanoseconds since the epoch.
         :paramtype start_time: Optional[int]
+        :keyword context: A dictionary of context values corresponding to the parent span. If not provided,
+          the current global context will be used.
+        :paramtype context: Optional[Dict[str, any]]
         :return: The span that was started
         :rtype: ~opentelemetry.trace.Span
         """
         otel_kind = _KIND_MAPPINGS.get(kind, OpenTelemetrySpanKind.INTERNAL)
         otel_links = self._parse_links(links)
 
+        otel_context = None
+        if context:
+            otel_context = extract(context)
+
         otel_span = self._tracer.start_span(
             name,
+            context=otel_context,
             kind=otel_kind,
             attributes=attributes,
             links=otel_links,
@@ -121,6 +130,7 @@ def start_as_current_span(
         attributes: Optional[Attributes] = None,
         links: Optional[Sequence[Link]] = None,
         start_time: Optional[int] = None,
+        context: Optional[Dict[str, Any]] = None,
         end_on_exit: bool = True,
     ) -> Iterator[Span]:
         """Context manager that starts a span and sets it as the current span in the context.
@@ -141,12 +151,17 @@ def start_as_current_span(
         :paramtype links: Optional[Sequence[Link]]
         :keyword start_time: The start time of the span in nanoseconds since the epoch.
         :paramtype start_time: Optional[int]
+        :keyword context: A dictionary of context values corresponding to the parent span. If not provided,
+          the current global context will be used.
+        :paramtype context: Optional[Dict[str, any]]
         :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: Iterator[~opentelemetry.trace.Span]
         """
-        span = self.start_span(name, kind=kind, attributes=attributes, links=links, start_time=start_time)
+        span = self.start_span(
+            name, kind=kind, attributes=attributes, links=links, start_time=start_time, context=context
+        )
         with trace.use_span(  # pylint: disable=not-context-manager
             span, record_exception=False, end_on_exit=end_on_exit
         ) as span:

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

@@ -395,3 +395,26 @@ def test_start_span_with_start_time(tracing_helper):
     finished_spans = tracing_helper.exporter.get_finished_spans()
     assert len(finished_spans) == 2
     assert finished_spans[1].start_time == start_time
+
+
+def test_tracer_with_custom_context(tracing_helper):
+    """Test that the tracer can start a span with a custom context."""
+    tracer = get_tracer()
+    assert tracer
+
+    with tracer.start_as_current_span(name="foo-span") as foo_span:
+        foo_trace_context = tracer.get_trace_context()
+
+    assert "traceparent" in foo_trace_context
+
+    with tracer.start_as_current_span(name="bar-span", context=dict(foo_trace_context)) as bar_span:
+        pass
+
+    finished_spans = tracing_helper.exporter.get_finished_spans()
+    assert len(finished_spans) == 2
+    assert finished_spans[0].name == "foo-span"
+    assert finished_spans[1].name == "bar-span"
+    assert finished_spans[1].context.trace_id == finished_spans[0].context.trace_id
+
+    # foo_span should be the parent of bar_span
+    assert finished_spans[1].parent.span_id == finished_spans[0].context.span_id