feat(tracer): support for external observability providers

aws_lambda_powertools/tracing/base.py

@@ -1,6 +1,6 @@
 from __future__ import annotations
 
-import abc
+from abc import ABC, abstractmethod
 from contextlib import contextmanager
 from typing import TYPE_CHECKING, Any, Generator, Sequence
 
@@ -9,10 +9,10 @@
     import traceback
 
 
-class BaseSegment(abc.ABC):
+class BaseSegment(ABC):
     """Holds common properties and methods on segment and subsegment."""
 
-    @abc.abstractmethod
+    @abstractmethod
     def close(self, end_time: int | None = None):
         """Close the trace entity by setting `end_time`
         and flip the in progress flag to False.
@@ -23,15 +23,15 @@ def close(self, end_time: int | None = None):
             Time in epoch seconds, by default current time will be used.
         """
 
-    @abc.abstractmethod
+    @abstractmethod
     def add_subsegment(self, subsegment: Any):
         """Add input subsegment as a child subsegment."""
 
-    @abc.abstractmethod
+    @abstractmethod
     def remove_subsegment(self, subsegment: Any):
         """Remove input subsegment from child subsegments."""
 
-    @abc.abstractmethod
+    @abstractmethod
     def put_annotation(self, key: str, value: str | numbers.Number | bool) -> None:
         """Annotate segment or subsegment with a key-value pair.
 
@@ -45,7 +45,7 @@ def put_annotation(self, key: str, value: str | numbers.Number | bool) -> None:
             Annotation value
         """
 
-    @abc.abstractmethod
+    @abstractmethod
     def put_metadata(self, key: str, value: Any, namespace: str = "default") -> None:
         """Add metadata to segment or subsegment. Metadata is not indexed
         but can be later retrieved by BatchGetTraces API.
@@ -60,7 +60,7 @@ def put_metadata(self, key: str, value: Any, namespace: str = "default") -> None
             Metadata namespace, by default 'default'
         """
 
-    @abc.abstractmethod
+    @abstractmethod
     def add_exception(self, exception: BaseException, stack: list[traceback.StackSummary], remote: bool = False):
         """Add an exception to trace entities.
 
@@ -77,8 +77,8 @@ def add_exception(self, exception: BaseException, stack: list[traceback.StackSum
         """
 
 
-class BaseProvider(abc.ABC):
-    @abc.abstractmethod
+class BaseProvider(ABC):
+    @abstractmethod
     @contextmanager
     def in_subsegment(self, name=None, **kwargs) -> Generator[BaseSegment, None, None]:
         """Return a subsegment context manger.
@@ -91,7 +91,7 @@ def in_subsegment(self, name=None, **kwargs) -> Generator[BaseSegment, None, Non
             Optional parameters to be propagated to segment
         """
 
-    @abc.abstractmethod
+    @abstractmethod
     @contextmanager
     def in_subsegment_async(self, name=None, **kwargs) -> Generator[BaseSegment, None, None]:
         """Return a subsegment async context manger.
@@ -104,7 +104,7 @@ def in_subsegment_async(self, name=None, **kwargs) -> Generator[BaseSegment, Non
             Optional parameters to be propagated to segment
         """
 
-    @abc.abstractmethod
+    @abstractmethod
     def put_annotation(self, key: str, value: str | numbers.Number | bool) -> None:
         """Annotate current active trace entity with a key-value pair.
 
@@ -118,7 +118,7 @@ def put_annotation(self, key: str, value: str | numbers.Number | bool) -> None:
             Annotation value
         """
 
-    @abc.abstractmethod
+    @abstractmethod
     def put_metadata(self, key: str, value: Any, namespace: str = "default") -> None:
         """Add metadata to the current active trace entity.
 
@@ -134,7 +134,7 @@ def put_metadata(self, key: str, value: Any, namespace: str = "default") -> None
             Metadata namespace, by default 'default'
         """
 
-    @abc.abstractmethod
+    @abstractmethod
     def patch(self, modules: Sequence[str]) -> None:
         """Instrument a set of supported libraries
 
@@ -144,6 +144,6 @@ def patch(self, modules: Sequence[str]) -> None:
             Set of modules to be patched
         """
 
-    @abc.abstractmethod
+    @abstractmethod
     def patch_all(self) -> None:
         """Instrument all supported libraries"""

aws_lambda_powertools/tracing/provider/aws_xray/aws_xray_tracer.py

@@ -0,0 +1,154 @@
+from __future__ import annotations
+
+from contextlib import asynccontextmanager, contextmanager
+from numbers import Number
+from typing import Any, AsyncGenerator, Generator, Literal, Sequence
+
+from aws_lambda_powertools.shared import constants
+from aws_lambda_powertools.shared.lazy_import import LazyLoader
+from aws_lambda_powertools.tracing.provider.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: Any,
+        category: Literal["Annotation", "Metadata", "Auto"] = "Auto",
+        **kwargs,
+    ) -> None:
+        """
+        Set an attribute on this span with a key-value pair.
+
+        Parameters
+        ----------
+        key: str
+            attribute key
+        value: Any
+            Value for attribute
+        category: Literal["Annotation","Metadata","Auto"] = "Auto"
+            This parameter specifies the category of attribute to set.
+            - **"Annotation"**: Sets the attribute as an Annotation.
+            - **"Metadata"**: Sets the attribute as Metadata.
+            - **"Auto" (default)**: Automatically determines the attribute
+            type based on its value.
+
+        kwargs: Optional[dict]
+            Optional parameters to be passed to provider.set_attributes
+        """
+        if category == "Annotation":
+            self.put_annotation(key=key, value=value)
+            return
+
+        if category == "Metadata":
+            self.put_metadata(key=key, value=value, namespace=kwargs.get("namespace", "dafault"))
+            return
+
+        # Auto
+        if isinstance(value, (str, Number, bool)):
+            self.put_annotation(key=key, value=value)
+            return
+
+        # Auto & not in (str, Number, bool)
+        self.put_metadata(key=key, value=value, namespace=kwargs.get("namespace", "dafault"))
+
+    def record_exception(self, exception: BaseException, **kwargs):
+        stack = aws_xray_sdk.core.utils.stacktrace.get_stacktrace()
+        self.add_exception(exception=exception, stack=stack)
+
+
+class AwsXrayProvider(BaseProvider):
+
+    def __init__(
+        self,
+        service: str = "",
+        disabled: bool | None = None,
+        auto_patch: bool | None = None,
+        patch_modules: Sequence[str] | None = None,
+    ):
+        from aws_xray_sdk.core import xray_recorder  # type: ignore
+
+        self.recorder = xray_recorder
+        self.in_subsegment = self.recorder.in_subsegment
+        self.in_subsegment_async = self.recorder.in_subsegment_async
+
+        self.service = service
+
+        super().__init__(
+            service=self.service,
+        )
+
+    @contextmanager
+    def trace(self, name: str, **kwargs) -> Generator[XraySpan, None, None]:
+        with self.in_subsegment(name=name, **kwargs) as sub_segment:
+            yield XraySpan(subsegment=sub_segment)
+
+    @asynccontextmanager
+    async def trace_async(self, name: str, **kwargs) -> AsyncGenerator[XraySpan, None]:
+        async with self.in_subsegment_async(name=name, **kwargs) as subsegment:
+            yield XraySpan(subsegment=subsegment)
+
+    def set_attribute(
+        self,
+        key: str,
+        value: Any,
+        category: Literal["Annotation", "Metadata", "Auto"] = "Auto",
+        **kwargs,
+    ) -> None:
+        """
+        Set an attribute on the current active span with a key-value pair.
+
+        Parameters
+        ----------
+        key: str
+            attribute key
+        value: Any
+            Value for attribute
+        category: Literal["Annotation","Metadata","Auto"] = "Auto"
+            This parameter specifies the type of attribute to set.
+            - **"Annotation"**: Sets the attribute as an Annotation.
+            - **"Metadata"**: Sets the attribute as Metadata.
+            - **"Auto" (default)**: Automatically determines the attribute
+            type based on its value.
+
+        kwargs: Optional[dict]
+            Optional parameters to be passed to provider.set_attributes
+        """
+        if category == "Annotation":
+            self.put_annotation(key=key, value=value)
+            return
+
+        if category == "Metadata":
+            self.put_metadata(key=key, value=value, namespace=kwargs.get("namespace", "dafault"))
+            return
+
+        # Auto
+        if isinstance(value, (str, Number, bool)):
+            self.put_annotation(key=key, value=value)
+            return
+
+        # Auto & not in (str, Number, bool)
+        self.put_metadata(key=key, value=value, namespace=kwargs.get("namespace", "dafault"))
+
+    def put_annotation(self, key: str, value: str | Number | bool) -> None:
+        return self.recorder.put_annotation(key=key, value=value)
+
+    def put_metadata(self, key: str, value: Any, namespace: str = "default") -> None:
+        return self.recorder.put_metadata(key=key, value=value, namespace=namespace)
+
+    def patch(self, modules: Sequence[str]) -> None:
+        return aws_xray_sdk.core.patch(modules)
+
+    def patch_all(self) -> None:
+        return aws_xray_sdk.core.patch_all()