Add decorator

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

@@ -31,7 +31,19 @@
 from collections import namedtuple
 import logging
 import os
-from typing import Any, Union
+import six
+import sys
+
+try:
+    from typing import TYPE_CHECKING
+except ImportError:
+    TYPE_CHECKING = False
+
+if TYPE_CHECKING:
+    from typing import Any, Union
+
+
+from azure.core.tracing import AbstractSpan
 
 
 __all__ = ("settings",)
@@ -61,9 +73,9 @@ def convert_bool(value):
         return value  # type: ignore
 
     val = value.lower()  # type: ignore
-    if val in ["yes", "1", "on"]:
+    if val in ["yes", "1", "on", "true", "True"]:
         return True
-    if val in ["no", "0", "off"]:
+    if val in ["no", "0", "off", "false", "False"]:
         return False
     raise ValueError("Cannot convert {} to boolean value".format(value))
 
@@ -102,14 +114,63 @@ def convert_logging(value):
     val = value.upper()  # type: ignore
     level = _levels.get(val)
     if not level:
-        raise ValueError(
-            "Cannot convert {} to log level, valid values are: {}".format(
-                value, ", ".join(_levels)
-            )
-        )
+        raise ValueError("Cannot convert {} to log level, valid values are: {}".format(value, ", ".join(_levels)))
     return level
 
 
+def get_opencensus_span():
+    # type: () -> OpenCensusSpan
+    """Returns the OpenCensusSpan if opencensus is installed else returns None"""
+    try:
+        from azure.core.tracing.ext.opencensus_span import OpenCensusSpan
+
+        return OpenCensusSpan
+    except ImportError:
+        return None
+
+
+def get_opencensus_span_if_opencensus_is_imported():
+    if "opencensus" not in sys.modules:
+        return None
+    return get_opencensus_span()
+
+
+_tracing_implementation_dict = {"opencensus": get_opencensus_span}
+
+
+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 value is None:
+        return get_opencensus_span_if_opencensus_is_imported()
+
+    wrapper_class = value
+    if isinstance(value, six.string_types):
+        value = value.lower()
+        get_wrapper_class = _tracing_implementation_dict.get(value, lambda: _Unset)
+        wrapper_class = get_wrapper_class()
+        if wrapper_class is _Unset:
+            raise ValueError(
+                "Cannot convert {} to AbstractSpan, valid values are: {}".format(
+                    value, ", ".join(_tracing_implementation_dict)
+                )
+            )
+
+    return wrapper_class
+
+
 class PrioritizedSetting(object):
     """Return a value for a global setting according to configuration precedence.
 
@@ -138,9 +199,7 @@ class PrioritizedSetting(object):
 
     """
 
-    def __init__(
-        self, name, env_var=None, system_hook=None, default=_Unset, convert=None
-    ):
+    def __init__(self, name, env_var=None, system_hook=None, default=_Unset, convert=None):
 
         self._name = name
         self._env_var = env_var
@@ -311,11 +370,7 @@ def defaults(self):
         """ Return implicit default values for all settings, ignoring environment and system.
 
         """
-        props = {
-            k: v.default
-            for (k, v) in self.__class__.__dict__.items()
-            if isinstance(v, PrioritizedSetting)
-        }
+        props = {k: v.default for (k, v) in self.__class__.__dict__.items() if isinstance(v, PrioritizedSetting)}
         return self._config(props)
 
     @property
@@ -335,11 +390,7 @@ def config(self, **kwargs):
         settings.config(log_level=logging.DEBUG)
 
         """
-        props = {
-            k: v()
-            for (k, v) in self.__class__.__dict__.items()
-            if isinstance(v, PrioritizedSetting)
-        }
+        props = {k: v() for (k, v) in self.__class__.__dict__.items() if isinstance(v, PrioritizedSetting)}
         props.update(kwargs)
         return self._config(props)
 
@@ -348,17 +399,19 @@ def _config(self, props):  # pylint: disable=no-self-use
         return Config(**props)
 
     log_level = PrioritizedSetting(
-        "log_level",
-        env_var="AZURE_LOG_LEVEL",
-        convert=convert_logging,
-        default=logging.INFO,
+        "log_level", env_var="AZURE_LOG_LEVEL", convert=convert_logging, default=logging.INFO
     )
 
     tracing_enabled = PrioritizedSetting(
-        "tracing_enbled",
-        env_var="AZURE_TRACING_ENABLED",
-        convert=convert_bool,
-        default=False,
+        "tracing_enbled", env_var="AZURE_TRACING_ENABLED", convert=convert_bool, 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
     )
 
 

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

@@ -3,7 +3,7 @@
 # Licensed under the MIT License.
 # ------------------------------------
 from azure.core.tracing.abstract_span import AbstractSpan
-from azure.core.tracing.context import tracing_context
-
-__all__ = ["tracing_context", "AbstractSpan"]
 
+__all__ = [
+    "AbstractSpan",
+]

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

@@ -2,6 +2,8 @@
 # Copyright (c) Microsoft Corporation.
 # Licensed under the MIT License.
 # ------------------------------------
+"""Protocol that defines what functions wrappers of tracing libraries should implement."""
+
 try:
     from typing import TYPE_CHECKING
 except ImportError:
@@ -19,7 +21,7 @@
 class AbstractSpan(Protocol):
     """Wraps a span from a distributed tracing implementation."""
 
-    def __init__(self, span=None, name=None):
+    def __init__(self, span=None, name=None): # pylint: disable=super-init-not-called
         # type: (Optional[Any], Optional[str]) -> None
         """
         If a span is given wraps the span. Else a new span is created.

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

@@ -0,0 +1,92 @@
+# --------------------------------------------------------------------------
+#
+# 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.
+#
+# --------------------------------------------------------------------------
+"""Common functions shared by both the sync and the async decorators."""
+
+from azure.core.tracing.context import tracing_context
+from azure.core.tracing.abstract_span import AbstractSpan
+from azure.core.settings import settings
+
+
+try:
+    from typing import TYPE_CHECKING
+except ImportError:
+    TYPE_CHECKING = False
+
+if TYPE_CHECKING:
+    from typing import Any, Optional
+
+
+def set_span_contexts(wrapped_span, span_instance=None):
+    # type: (AbstractSpan, Optional[AbstractSpan]) -> None
+    """
+    Set the sdk context and the implementation context. `span_instance` will be used to set the implementation context
+    if passed in else will use `wrapped_span.span_instance`.
+
+    :param wrapped_span: The `AbstractSpan` to set as the sdk context
+    :type wrapped_span: `azure.core.tracing.abstract_span.AbstractSpan`
+    :param span_instance: The span to set as the current span for the implementation context
+    """
+    tracing_context.current_span.set(wrapped_span)
+    impl_wrapper = settings.tracing_implementation()
+    if wrapped_span is not None:
+        span_instance = wrapped_span.span_instance
+    if impl_wrapper is not None:
+        impl_wrapper.set_current_span(span_instance)
+
+
+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.
+
+    :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`
+    """
+    wrapper_class = settings.tracing_implementation()
+    if wrapper_class is None:
+        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:
+        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
+
+
+def should_use_trace(parent_span):
+    # type: (AbstractSpan) -> bool
+    """Given Parent Span Returns whether the function should be traced"""
+    only_propagate = settings.tracing_should_only_propagate()
+    return bool(parent_span and not only_propagate)

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

@@ -2,7 +2,10 @@
 # Copyright (c) Microsoft Corporation.
 # Licensed under the MIT License.
 # ------------------------------------
+"""The context for the azure.core.tracing. Holds global variables in a thread and async safe way."""
+
 import threading
+from azure.core.settings import settings
 
 try:
     from typing import TYPE_CHECKING
@@ -119,7 +122,6 @@ class TracingContext:
     def __init__(self):
         # type: () -> None
         self.current_span = TracingContext._get_context_class("current_span", None)
-        self.tracing_impl = TracingContext._get_context_class("tracing_impl", None)
 
     def with_current_context(self, func):
         # type: (Callable[[Any], Any]) -> Any
@@ -129,7 +131,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 = 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()
@@ -140,7 +142,6 @@ def call_with_current_context(*args, **kwargs):
                 wrapper_class.set_current_tracer(current_impl_tracer)
                 current_span = wrapped_span or wrapper_class(current_impl_span)
                 self.current_span.set(current_span)
-                self.tracing_impl.set(wrapper_class)
             return func(*args, **kwargs)
 
         return call_with_current_context

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

@@ -0,0 +1,64 @@
+# --------------------------------------------------------------------------
+#
+# 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.
+#
+# --------------------------------------------------------------------------
+"""The decorator to apply if you want the given function traced."""
+
+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_trace(func):
+    # type: (Callable[[Any], Any]) -> Callable[[Any], Any]
+    @functools.wraps(func)
+    def wrapper_use_tracer(self, *args, **kwargs):
+        # type: (Any) -> Any
+        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)
+            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 orig_wrapped_span is None and passed_in_parent is None:
+                parent_span.finish()
+            common.set_span_contexts(orig_wrapped_span, span_instance=original_span_instance)
+        else:
+            ans = func(self, *args, **kwargs)
+        return ans
+
+    return wrapper_use_tracer