add new base and xray provider

Author: roger-zhangg

aws_lambda_powertools/tracing/provider/base.py

@@ -0,0 +1,90 @@
+import abc
+import numbers
+from contextlib import contextmanager
+from typing import Generator, Sequence, Union
+
+
+## TO-Discuss how to refactor this one. Segment seems X-ray exclusive concept
+class BaseSpan(abc.ABC):
+    """Holds common properties and methods on segment and subsegment."""
+
+    @abc.abstractmethod
+    def set_attribute(self, key: str, value: Union[str, numbers.Number, bool], **kwargs) -> None:
+        """Annotate segment or subsegment with a key-value pair.
+
+        Note: Annotations will be indexed for later search query.
+
+        Parameters
+        ----------
+        key: str
+            Metadata key
+        value: Union[str, numbers.Number, bool]
+            Annotation value
+        """
+
+    @abc.abstractmethod
+    def record_exception(self, exception: BaseException, **kwargs):
+        """Add an exception to trace entities.
+
+        Parameters
+        ----------
+        exception: Exception
+            Caught exception
+            Output from `traceback.extract_stack()`.
+        """
+
+
+class BaseProvider(abc.ABC):
+    @abc.abstractmethod
+    @contextmanager
+    def trace(self, name=None, **kwargs) -> Generator[BaseSpan, None, None]:
+        """Return a span context manger.
+
+        Parameters
+        ----------
+        name: str
+            Span name
+        kwargs: Optional[dict]
+            Optional parameters to be propagated to span
+        """
+
+    @abc.abstractmethod
+    @contextmanager
+    def trace_async(self, name=None, **kwargs) -> Generator[BaseSpan, None, None]:
+        """Return a async span context manger.
+
+        Parameters
+        ----------
+        name: str
+            Span name
+        kwargs: Optional[dict]
+            Optional parameters to be propagated to span
+        """
+
+    @abc.abstractmethod
+    def set_attribute(self, key: str, value: Union[str, numbers.Number, bool], **kwargs) -> None:
+        """Annotate current active trace entity with a key-value pair.
+
+        Note: Annotations will be indexed for later search query.
+
+        Parameters
+        ----------
+        key: str
+            Metadata key
+        value: Union[str, numbers.Number, bool]
+            Annotation value
+        """
+
+    @abc.abstractmethod
+    def patch(self, modules: Sequence[str]) -> None:
+        """Instrument a set of supported libraries
+
+        Parameters
+        ----------
+        modules: Set[str]
+            Set of modules to be patched
+        """
+
+    @abc.abstractmethod
+    def patch_all(self) -> None:
+        """Instrument all supported libraries"""

aws_lambda_powertools/tracing/provider/otel_tracer.py

@@ -13,6 +13,11 @@
 logger = logging.getLogger(__name__)
 
 
+# optl terminology first
+# 1. Provider based on OTel terminology
+# 2. X-Ray provider on top of the new BaseProvider
+# 3. Datadog provider on top of the new BaseProvider
+# access xray sdk
 class OtelSpan(BaseSegment):
     def __init__(self, otel_span=otel_trace.Span):
         self.otel_span = otel_span

aws_lambda_powertools/tracing/provider/xray_tracer.py

@@ -0,0 +1,61 @@
+from __future__ import annotations
+
+from contextlib import contextmanager
+from numbers import Number
+from typing import Generator
+
+from ...shared import constants
+from ...shared.lazy_import import LazyLoader
+from .base import BaseProvider, BaseSpan
+
+aws_xray_sdk = LazyLoader(constants.XRAY_SDK_MODULE, globals(), constants.XRAY_SDK_MODULE)
+
+
+class XraySpan(BaseSpan):
+    def __init__(self, subsegment):
+        self.subsegment = subsegment
+        self.add_subsegment = self.subsegment.add_subsegment
+        self.remove_subsegment = self.subsegment.remove_subsegment
+        self.put_annotation = self.subsegment.put_annotation
+        self.put_metadata = self.subsegment.put_metadata
+        self.add_exception = self.subsegment.add_exception
+        self.close = self.subsegment.close
+
+    def set_attribute(self, key: str, value: str | Number | bool, **kwargs) -> None:
+        if kwargs.get("namespace", "") != "":
+            self.put_metadata(key=key, value=value, namespace=kwargs["namespace"])
+        else:
+            self.put_annotation(key=key, value=value)
+
+    def record_exception(self, exception: BaseException, **kwargs):
+        stack = aws_xray_sdk.core.utils.stacktrace.get_stacktrace(limit=self._max_trace_back)
+        self.add_exception(exception=exception, stack=stack)
+
+
+class XrayProvider(BaseProvider):
+    def __init__(self, xray_recorder=None):
+        if not xray_recorder:
+            from aws_xray_sdk.core import xray_recorder
+        self.recorder = xray_recorder
+        self.patch = aws_xray_sdk.core.patch
+        self.patch_all = aws_xray_sdk.core.patch_all
+        self.in_subsegment = self.recorder.in_subsegment
+        self.in_subsegment_async = self.recorder.in_subsegment_async
+        self.put_annotation = self.recorder.put_annotation
+        self.put_metadata = self.recorder.put_metadata
+
+    @contextmanager
+    def trace(self, name: str) -> Generator[XraySpan, None, None]:
+        with self.in_subsegment(name=name) as sub_segment:
+            yield XraySpan(subsegment=sub_segment)
+
+    @contextmanager
+    def trace_async(self, name: str) -> Generator[XraySpan, None, None]:
+        with self.in_subsegment_async(name=name) as sub_segment:
+            yield XraySpan(subsegment=sub_segment)
+
+    def set_attribute(self, key: str, value: str | Number | bool, **kwargs) -> None:
+        if kwargs.get("namespace", "") != "":
+            self.put_metadata(key=key, value=value, namespace=kwargs["namespace"])
+        else:
+            self.put_annotation(key=key, value=value)