feat(tracer): Support for external observability providers - Tracer

aws_lambda_powertools/tracing/base.py

@@ -5,17 +5,18 @@
 from typing import Any, Generator, List, Optional, Sequence, Union
 
 
+## MAINTENANCE: deprecated, kept for backwards compatibility until v3 is released
 class BaseSegment(abc.ABC):
     """Holds common properties and methods on segment and subsegment."""
 
     @abc.abstractmethod
-    def close(self, end_time: Optional[int] = None):
+    def close(self, end_time: Optional[float] = None):
         """Close the trace entity by setting `end_time`
         and flip the in progress flag to False.
 
         Parameters
         ----------
-        end_time: int
+        end_time: float
             Time in epoch seconds, by default current time will be used.
         """
 

aws_lambda_powertools/tracing/provider/base.py

@@ -0,0 +1,112 @@
+import abc
+from contextlib import asynccontextmanager, contextmanager
+from typing import Any, AsyncGenerator, Generator, Sequence
+
+
+class BaseSpan(abc.ABC):
+    """A span represents a unit of work or operation within a trace.
+    Spans are the building blocks of Traces."""
+
+    @abc.abstractmethod
+    def set_attribute(self, key: str, value: Any, **kwargs) -> None:
+        """Set an attribute for a span with a key-value pair.
+
+        Parameters
+        ----------
+        key: str
+            Attribute key
+        value: Any
+            Attribute value
+        kwargs: Optional[dict]
+            Optional parameters
+        """
+
+    @abc.abstractmethod
+    def record_exception(self, exception: BaseException, **kwargs):
+        """Records an exception to this Span.
+
+        Parameters
+        ----------
+        exception: Exception
+            Caught exception during the exectution of this Span
+        kwargs: Optional[dict]
+            Optional parameters
+        """
+
+
+class BaseProvider(abc.ABC):
+    """BaseProvider is an abstract base class that defines the expected behavior for tracing providers
+    used by Tracer. Inheriting classes must implement this interface to be compatible with Tracer.
+    """
+
+    @abc.abstractmethod
+    @contextmanager
+    def trace(self, name: str, **kwargs) -> Generator[BaseSpan, None, None]:
+        """Context manager for creating a new span and set it
+        as the current span in this tracer's context.
+
+        Exiting the context manager will call the span's end method,
+        as well as return the current span to its previous value by
+        returning to the previous context.
+
+        Parameters
+        ----------
+        name: str
+            Span name
+        kwargs: Optional[dict]
+            Optional parameters to be propagated to the span
+        """
+
+    @abc.abstractmethod
+    @asynccontextmanager
+    def trace_async(self, name: str, **kwargs) -> AsyncGenerator[BaseSpan, None]:
+        """Async Context manager for creating a new span and set it
+        as the current span in this tracer's context.
+
+        Exiting the context manager will call the span's end method,
+        as well as return the current span to its previous value by
+        returning to the previous context.
+
+        Parameters
+        ----------
+        name: str
+            Span name
+        kwargs: Optional[dict]
+            Optional parameters to be propagated to the span
+        """
+
+    @abc.abstractmethod
+    def set_attribute(self, key: str, value: Any, **kwargs) -> None:
+        """set attribute on current active span with a key-value pair.
+
+        Parameters
+        ----------
+        key: str
+            attribute key
+        value: Any
+            attribute value
+        kwargs: Optional[dict]
+            Optional parameters to be propagated to the span
+        """
+
+    @abc.abstractmethod
+    def patch(self, modules: Sequence[str]) -> None:
+        """Instrument a set of given libraries if supported by provider
+        See specific provider for more detail
+
+        Exmaple
+        -------
+        tracer = Tracer(service="payment")
+        libraries = (['aioboto3',mysql])
+        # provider.patch will be called by tracer.patch
+        tracer.patch(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/xray/xray_tracer.py

@@ -0,0 +1,141 @@
+from __future__ import annotations
+
+from contextlib import asynccontextmanager, contextmanager
+from numbers import Number
+from typing import Any, AsyncGenerator, Generator, Literal, Sequence, Union
+
+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: Any,
+        category: Literal["Annotation", "Metadata", "Auto"] = "Auto",
+        **kwargs,
+    ) -> None:
+        """
+        Set 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 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.in_subsegment = self.recorder.in_subsegment
+        self.in_subsegment_async = self.recorder.in_subsegment_async
+
+    @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 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: Union[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()