Add decorator

sdk/core/azure-core/azure/core/settings.py

@@ -33,6 +33,8 @@
 import os
 from typing import Any, Union
 
+from azure.core.tracing import AbstractSpan
+
 
 __all__ = ("settings",)
 
@@ -110,6 +112,47 @@ def convert_logging(value):
     return level
 
 
+def _get_opencensus_wrapper():
+    # type: () -> OpencensusWrapper
+    """Returns the OpencensusWrapper if opencensus is installed else returns None"""
+    try:
+        from azure.core.tracing.ext.opencensus_wrapper import OpencensusWrapper
+
+        return OpencensusWrapper
+    except ImportError:
+        return None
+
+
+def convert_tracing_impl(value):
+    # type: (Union[str, AbstractSpan]) -> AbstractSpan
+    """Convert a string to AbstractSpan
+
+    If a AbstractSpan is passed in, it is returned as-is. Otherwise the function
+    understands the following strings, ignoring case:
+
+    * "opencensus"
+
+    :param value: the value to convert
+    :type value: string
+    :returns: AbstractSpan
+    :raises ValueError: If conversion to AbstractSpan fails
+
+    """
+    if issubclass(value.__class__, AbstractSpan):
+        return value
+
+    _impl_dict = {"opencensus": _get_opencensus_wrapper()}
+    wrapper_class = _impl_dict.get(value, None)
+    if wrapper_class is None:
+        raise ValueError(
+            "Cannot convert {} to AbstractSpan, valid values are: {}".format(
+                value, ", ".join(_impl_dict)
+            )
+        )
+
+    return wrapper_class
+
+
 class PrioritizedSetting(object):
     """Return a value for a global setting according to configuration precedence.
 
@@ -361,5 +404,19 @@ def _config(self, props):  # pylint: disable=no-self-use
         default=False,
     )
 
+    tracing_implementation = PrioritizedSetting(
+        "tracing_implementation",
+        env_var="AZURE_SDK_TRACING_IMPLEMENTATION",
+        convert=convert_tracing_impl,
+        default=None,
+    )
+
+    tracing_should_only_propagate = PrioritizedSetting(
+        "tracing_should_only_propagate",
+        env_var="AZURE_TRACING_ONLY_PROPAGATE",
+        convert=convert_bool,
+        default=False,
+    )
+
 
 settings = Settings()

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

@@ -4,6 +4,16 @@
 # ------------------------------------
 from azure.core.tracing.abstract_span import AbstractSpan
 from azure.core.tracing.context import tracing_context
+from azure.core.tracing.decorator import distributed_tracing_decorator
 
-__all__ = ["tracing_context", "AbstractSpan"]
+try:
+    from azure.core.tracing.decorator_async import distributed_tracing_decorator_async
+except ImportError:
+    distributed_tracing_decorator_async = lambda x: x
 
+__all__ = [
+    "tracing_context",
+    "AbstractSpan",
+    "distributed_tracing_decorator",
+    "distributed_tracing_decorator_async",
+]

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

@@ -0,0 +1,91 @@
+# --------------------------------------------------------------------------
+#
+# Copyright (c) Microsoft Corporation. All rights reserved.
+#
+# The MIT License (MIT)
+#
+# Permission is hereby granted, free of charge, to any person obtaining a copy
+# of this software and associated documentation files (the ""Software""), to deal
+# in the Software without restriction, including without limitation the rights
+# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+# copies of the Software, and to permit persons to whom the Software is
+# furnished to do so, subject to the following conditions:
+#
+# The above copyright notice and this permission notice shall be included in
+# all copies or substantial portions of the Software.
+#
+# THE SOFTWARE IS PROVIDED *AS IS*, WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
+# THE SOFTWARE.
+#
+# --------------------------------------------------------------------------
+from os import environ
+import re
+
+from azure.core.tracing.context import tracing_context
+from azure.core.tracing.abstract_span import AbstractSpan
+from azure.core.settings import settings
+
+
+def _get_opencensus_wrapper():
+    # type: () -> OpencensusWrapper
+    """Returns the OpencensusWrapper if opencensus is installed else returns None"""
+    try:
+        from azure.core.tracing.ext.opencensus_wrapper  import OpencensusWrapper
+
+        return OpencensusWrapper
+    except ImportError:
+        return None
+
+
+def set_span_contexts(wrapped_span, span_instance=None, impl_wrapper=None):
+    # type: (AbstractSpan, AbstractSpan) -> None
+    tracing_context.current_span.set(wrapped_span)
+    impl_wrapper = impl_wrapper or wrapped_span
+    tracing_context.tracing_impl.set(impl_wrapper.__class__)
+    if wrapped_span is not None or (
+        span_instance is not None and impl_wrapper is not None
+    ):
+        span_instance = span_instance or wrapped_span.span_instance
+        impl_wrapper.set_current_span(span_instance)
+
+
+def get_parent(kwargs, *args):
+    # type: (Any) -> Tuple(Any, Any, Any)
+    """Returns the parent span that of the span that represents the function and the spans before that parent span"""
+    parent_span = kwargs.pop("parent_span", None)  # type: AbstractSpan
+    orig_wrapped_span = tracing_context.current_span.get()
+
+    # wrapper class get from tracing_context, settings or assume OpencensusWrapper if opencesus is installed
+    wrapper_class = (
+        tracing_context.tracing_impl.get()
+        or settings.tracing_implementation()
+        or _get_opencensus_wrapper()
+    )
+    if wrapper_class is None:
+        return None, orig_wrapped_span, None
+
+    # 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:
+        current_span = wrapper_class.get_current_span()
+        parent_span = (
+            wrapper_class(span=current_span)
+            if current_span
+            else wrapper_class(name="azure-sdk-for-python-first_parent_span")
+        )
+
+    return parent_span, orig_wrapped_span, wrapper_class.get_current_span()
+
+
+def should_use_trace(parent_span):
+    # type: (AbstractSpan, List[str], str)
+    """Given Parent Span Returns whether the function should be traced"""
+    only_propagate = settings.tracing_should_only_propagate()
+    return parent_span and not only_propagate

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

@@ -3,6 +3,7 @@
 # Licensed under the MIT License.
 # ------------------------------------
 import threading
+from azure.core.settings import settings
 
 try:
     from typing import TYPE_CHECKING
@@ -129,7 +130,7 @@ def with_current_context(self, func):
         :return: The target the pass in instead of the function
         """
         wrapped_span = tracing_context.current_span.get()
-        wrapper_class = self.tracing_impl.get()
+        wrapper_class = self.tracing_impl.get() or settings.tracing_implementation()
         if wrapper_class is not None:
             current_impl_span = wrapper_class.get_current_span()
             current_impl_tracer = wrapper_class.get_current_tracer()

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

@@ -0,0 +1,58 @@
+# --------------------------------------------------------------------------
+#
+# Copyright (c) Microsoft Corporation. All rights reserved.
+#
+# The MIT License (MIT)
+#
+# Permission is hereby granted, free of charge, to any person obtaining a copy
+# of this software and associated documentation files (the ""Software""), to deal
+# in the Software without restriction, including without limitation the rights
+# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+# copies of the Software, and to permit persons to whom the Software is
+# furnished to do so, subject to the following conditions:
+#
+# The above copyright notice and this permission notice shall be included in
+# all copies or substantial portions of the Software.
+#
+# THE SOFTWARE IS PROVIDED *AS IS*, WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
+# THE SOFTWARE.
+#
+# --------------------------------------------------------------------------
+import functools
+import azure.core.tracing.common as common
+
+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(
+            kwargs
+        )
+        ans = None
+        if common.should_use_trace(parent_span, func.__name__):
+            common.set_span_contexts(parent_span)
+            name = self.__class__.__name__ + "." + func.__name__
+            child = parent_span.span(name=name)
+            child.start()
+            common.set_span_contexts(child)
+            ans = func(self, *args, **kwargs)
+            child.finish()
+            common.set_span_contexts(parent_span)
+            if getattr(parent_span, "was_created_by_azure_sdk", False):
+                parent_span.finish()
+            common.set_span_contexts(
+                original_span_from_sdk_context,
+                span_instance=original_span_instance,
+                impl_wrapper=parent_span,
+            )
+        else:
+            ans = func(self, *args, **kwargs)
+        return ans
+
+    return wrapper_use_tracer

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

@@ -0,0 +1,58 @@
+# --------------------------------------------------------------------------
+#
+# Copyright (c) Microsoft Corporation. All rights reserved.
+#
+# The MIT License (MIT)
+#
+# Permission is hereby granted, free of charge, to any person obtaining a copy
+# of this software and associated documentation files (the ""Software""), to deal
+# in the Software without restriction, including without limitation the rights
+# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+# copies of the Software, and to permit persons to whom the Software is
+# furnished to do so, subject to the following conditions:
+#
+# The above copyright notice and this permission notice shall be included in
+# all copies or substantial portions of the Software.
+#
+# THE SOFTWARE IS PROVIDED *AS IS*, WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
+# THE SOFTWARE.
+#
+# --------------------------------------------------------------------------
+import functools
+import azure.core.tracing.common as common
+
+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(
+            kwargs
+        )
+        ans = None
+        if common.should_use_trace(parent_span, func.__name__):
+            common.set_span_contexts(parent_span)
+            name = self.__class__.__name__ + "." + func.__name__
+            child = parent_span.span(name=name)
+            child.start()
+            common.set_span_contexts(child)
+            ans = await func(self, *args, **kwargs)
+            child.finish()
+            common.set_span_contexts(parent_span)
+            if getattr(parent_span, "was_created_by_azure_sdk", False):
+                parent_span.finish()
+            common.set_span_contexts(
+                original_span_from_sdk_context,
+                span_instance=original_span_instance,
+                impl_wrapper=parent_span,
+            )
+        else:
+            ans = await func(self, *args, **kwargs)
+        return ans
+
+    return wrapper_use_tracer