simplify get parent

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

@@ -57,27 +57,21 @@ def set_span_contexts(wrapped_span, span_instance=None):
         impl_wrapper.set_current_span(span_instance)
 
 
-def get_parent_and_original_contexts(kwargs):
+def get_parent_span(parent_span):
     # type: (Any) -> Tuple(AbstractSpan, AbstractSpan, Any)
     """
     Returns the current span so that the function's span will be its child. It will create a new span if there is
-    no current span in any of the context. Because it creates a new span in those cases, also returns the original
-    span as well. If it does not create a new span then the orig_wrapped_span should be the same as the parent span.
+    no current span in any of the context.
 
-    `AbstractSpan = azure.core.tracing.abstract_span.AbstractSpan`
-
-    :param kwargs: The dictionary of named variables that the user passed into the top level azure function
-    :param kwargs: dict
-    :returns: The parent_span, the original wrapped span, the original implementation span
-    :rtype: `AbstractSpan`, `AbstractSpan`, implementation span
+    :param parent_span: The parent_span arg that the user passes into the top level function
+    :returns: the parent_span of the function to be traced
+    :rtype: `azure.core.tracing.abstract_span.AbstractSpan`
     """
-    parent_span = kwargs.pop("parent_span", None)  # type: AbstractSpan
-    orig_wrapped_span = tracing_context.current_span.get()
-
     wrapper_class = settings.tracing_implementation()
     if wrapper_class is None:
-        return None, orig_wrapped_span, None
-    original_wrapper_span_instance = wrapper_class.get_current_span()
+        return None
+
+    orig_wrapped_span = tracing_context.current_span.get()
     # parent span is given, get from my context, get from the implementation context or make our own
     parent_span = orig_wrapped_span if parent_span is None else wrapper_class(parent_span)
     if parent_span is None:
@@ -88,7 +82,7 @@ def get_parent_and_original_contexts(kwargs):
             else wrapper_class(name="azure-sdk-for-python-first_parent_span")
         )
 
-    return parent_span, orig_wrapped_span, original_wrapper_span_instance
+    return parent_span
 
 
 def should_use_trace(parent_span):

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

@@ -28,16 +28,22 @@
 import functools
 
 import azure.core.tracing.common as common
+from azure.core.settings import settings
+from azure.core.tracing.context import tracing_context
 
 
 def distributed_tracing_decorator(func):
     # type: (Callable[[Any], Any]) -> Callable[[Any], Any]
     @functools.wraps(func)
     def wrapper_use_tracer(self, *args, **kwargs):
         # type: (Any) -> Any
-        parent_span, original_span_from_sdk_context, original_span_instance = common.get_parent_and_original_contexts(
-            kwargs
-        )
+        passed_in_parent = kwargs.pop("parent_span", None)
+        orig_wrapped_span = tracing_context.current_span.get()
+        wrapper_class = settings.tracing_implementation()
+        original_span_instance = None
+        if wrapper_class is not None:
+            original_span_instance = wrapper_class.get_current_span()
+        parent_span = common.get_parent_span(passed_in_parent)
         ans = None
         if common.should_use_trace(parent_span):
             common.set_span_contexts(parent_span)
@@ -48,9 +54,9 @@ def wrapper_use_tracer(self, *args, **kwargs):
             ans = func(self, *args, **kwargs)
             child.finish()
             common.set_span_contexts(parent_span)
-            if getattr(parent_span, "was_created_by_azure_sdk", False):
+            if orig_wrapped_span is None and passed_in_parent is None:
                 parent_span.finish()
-            common.set_span_contexts(original_span_from_sdk_context, span_instance=original_span_instance)
+            common.set_span_contexts(orig_wrapped_span, span_instance=original_span_instance)
         else:
             ans = func(self, *args, **kwargs)
         return ans

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

@@ -28,16 +28,22 @@
 import functools
 
 import azure.core.tracing.common as common
+from azure.core.settings import settings
+from azure.core.tracing.context import tracing_context
 
 
 def distributed_tracing_decorator_async(func):
     # type: (Callable[[Any], Any]) -> Callable[[Any], Any]
     @functools.wraps(func)
     async def wrapper_use_tracer(self, *args, **kwargs):
         # type: (Any) -> Any
-        parent_span, original_span_from_sdk_context, original_span_instance = common.get_parent_and_original_contexts(
-            kwargs
-        )
+        passed_in_parent = kwargs.pop("parent_span", None)
+        orig_wrapped_span = tracing_context.current_span.get()
+        wrapper_class = settings.tracing_implementation()
+        original_span_instance = None
+        if wrapper_class is not None:
+            original_span_instance = wrapper_class.get_current_span()
+        parent_span = common.get_parent_span(passed_in_parent)
         ans = None
         if common.should_use_trace(parent_span):
             common.set_span_contexts(parent_span)
@@ -48,9 +54,9 @@ async def wrapper_use_tracer(self, *args, **kwargs):
             ans = await func(self, *args, **kwargs)
             child.finish()
             common.set_span_contexts(parent_span)
-            if getattr(parent_span, "was_created_by_azure_sdk", False):
+            if orig_wrapped_span is None and passed_in_parent is None:
                 parent_span.finish()
-            common.set_span_contexts(original_span_from_sdk_context, span_instance=original_span_instance)
+            common.set_span_contexts(orig_wrapped_span, span_instance=original_span_instance)
         else:
             ans = await func(self, *args, **kwargs)
         return ans

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

@@ -84,41 +84,30 @@ def test_set_span_context(self):
             assert parent.span_instance == wrapper.get_current_span()
             assert tracing_context.current_span.get() == parent
 
-    def test_get_parent_and_original_contexts(self):
+    def test_get_parent_span(self):
         with ContextHelper():
             opencensus = sys.modules["opencensus"]
             del sys.modules["opencensus"]
 
-            parent, orig_tracing_context, orig_span_inst = common.get_parent_and_original_contexts({})
-            assert orig_span_inst is None
+            parent = common.get_parent_span(None)
             assert parent is None
-            assert orig_tracing_context is None
 
             sys.modules["opencensus"] = opencensus
-            parent, orig_tracing_context, orig_span_inst = common.get_parent_and_original_contexts({})
-            assert orig_span_inst is None
+            parent = common.get_parent_span(None)
             assert parent.span_instance.name == "azure-sdk-for-python-first_parent_span"
-            assert orig_tracing_context is None
 
             tracer = tracer_module.Tracer(sampler=AlwaysOnSampler())
-            parent, orig_tracing_context, orig_span_inst = common.get_parent_and_original_contexts({})
-            assert orig_span_inst is None
+            parent = common.get_parent_span(None)
             assert parent.span_instance.name == "azure-sdk-for-python-first_parent_span"
-            assert orig_tracing_context is None
             parent.finish()
 
             some_span = tracer.start_span(name="some_span")
-            new_parent, orig_tracing_context, orig_span_inst = common.get_parent_and_original_contexts({})
-            assert orig_span_inst == some_span
+            new_parent = common.get_parent_span(None)
             assert new_parent.span_instance.name == "some_span"
-            assert orig_tracing_context is None
+            some_span.finish()
 
-            kwarg = {"parent_span": parent.span_instance}
-            should_be_old_parent, orig_tracing_context, orig_span_inst = common.get_parent_and_original_contexts(kwarg)
-            assert kwarg.get("parent_span") is None
-            assert orig_span_inst == some_span
+            should_be_old_parent = common.get_parent_span(parent.span_instance)
             assert should_be_old_parent.span_instance == parent.span_instance
-            assert orig_tracing_context is None
 
     def test_should_use_trace(self):
         with ContextHelper(environ={"AZURE_TRACING_ONLY_PROPAGATE": "yes"}):