Adding tracer to v3

Author: leandrodamascena

aws_lambda_powertools/tracing/base.py

@@ -1,14 +1,14 @@
-import abc
 import numbers
 import traceback
+from abc import ABC, abstractmethod
 from contextlib import contextmanager
 from typing import Any, Generator, List, Optional, Sequence, Union
 
 
-class BaseSegment(abc.ABC):
+class BaseSegment(ABC):
     """Holds common properties and methods on segment and subsegment."""
 
-    @abc.abstractmethod
+    @abstractmethod
     def close(self, end_time: Optional[int] = None):
         """Close the trace entity by setting `end_time`
         and flip the in progress flag to False.
@@ -19,15 +19,15 @@ def close(self, end_time: Optional[int] = 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: Union[str, numbers.Number, bool]) -> None:
         """Annotate segment or subsegment with a key-value pair.
 
@@ -41,7 +41,7 @@ def put_annotation(self, key: str, value: Union[str, numbers.Number, bool]) -> N
             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.
@@ -56,7 +56,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.
 
@@ -73,8 +73,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.
@@ -87,7 +87,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.
@@ -100,7 +100,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: Union[str, numbers.Number, bool]) -> None:
         """Annotate current active trace entity with a key-value pair.
 
@@ -114,7 +114,7 @@ def put_annotation(self, key: str, value: Union[str, numbers.Number, bool]) -> N
             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.
 
@@ -130,7 +130,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
 
@@ -140,6 +140,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/__init__.py

@@ -0,0 +1,142 @@
+from __future__ import annotations
+
+from contextlib import asynccontextmanager, contextmanager
+from numbers import Number
+from typing import Any, AsyncGenerator, Generator, Literal, Sequence, Union
+
+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, 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 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: 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()

aws_lambda_powertools/tracing/provider/aws_xray/aws_xray_tracer.py

@@ -0,0 +1,112 @@
+from abc import ABC, abstractmethod
+from contextlib import asynccontextmanager, contextmanager
+from typing import Any, AsyncGenerator, Generator, Sequence
+
+
+class BaseSpan(ABC):
+    """A span represents a unit of work or operation within a trace.
+    Spans are the building blocks of Traces."""
+
+    @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
+        """
+
+    @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):
+    """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.
+    """
+
+    @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
+        """
+
+    @abstractmethod
+    @asynccontextmanager
+    def trace_async(self, name: str, **kwargs) -> AsyncGenerator[BaseSpan, None]:
+        """Async Context manager for creating a new span async 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
+        """
+
+    @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
+        """
+
+    @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
+        """
+
+    @abstractmethod
+    def patch_all(self) -> None:
+        """Instrument all supported libraries"""