feat: add tracer support for async, escape hatch, and patch given modules

python/aws_lambda_powertools/metrics/exceptions.py

@@ -1,14 +1,22 @@
 class MetricUnitError(Exception):
+    """When metric unit is not supported by CloudWatch"""
+
     pass
 
 
 class SchemaValidationError(Exception):
+    """When serialization fail schema validation"""
+
     pass
 
 
 class MetricValueError(Exception):
+    """When metric value isn't a valid number"""
+
     pass
 
 
 class UniqueNamespaceError(Exception):
+    """When an additional namespace is set"""
+
     pass

python/aws_lambda_powertools/middleware_factory/exceptions.py

@@ -1,2 +1,4 @@
 class MiddlewareInvalidArgumentError(Exception):
+    """When middleware receives non keyword=arguments"""
+
     pass

python/aws_lambda_powertools/tracing/base.py

@@ -0,0 +1,233 @@
+from __future__ import annotations
+
+import abc
+import logging
+from typing import Any, List
+
+import aws_xray_sdk
+import aws_xray_sdk.core
+
+is_cold_start = True
+logger = logging.getLogger(__name__)
+
+
+class TracerProvider(metaclass=abc.ABCMeta):
+    """Tracer provider abstract class
+
+    Providers should be initialized independently. This
+    allows providers to control their config/initialization,
+    and only pass a class instance to
+    `aws_lambda_powertools.tracing.tracer.Tracer`.
+
+    It also allows custom providers to keep lean while Tracer provide:
+
+    * a simplified UX
+    * decorators for Lambda handler and methods
+    * auto-patching, patch all modules by default or a subset
+    * disabling all tracing operations with a single parameter or env var
+
+    Trace providers should implement the following methods:
+
+    * **patch**
+    * **create_subsegment**
+    * **end_subsegment**
+    * **put_metadata**
+    * **put_annotation**
+    * **disable_tracing_provider**
+
+    These methods will be called by
+    `aws_lambda_powertools.tracing.tracer.Tracer` -
+    See `aws_lambda_powertools.tracing.base.XrayProvider`
+    for a reference implementation.
+
+    `aws_lambda_powertools.tracing.tracer.Tracer` decorators
+    for Lambda and methods use the following provider methods:
+
+    * create_subsegment
+    * put_metadata
+    * end_subsegment
+
+    Example
+    -------
+    **Client using a custom tracing provider**
+
+        from aws_lambda_powertools.tracing import Tracer
+        ... import ... ProviderX
+        custom_provider = ProviderX()
+        tracer = Tracer(service="greeting", provider=custom_provider)
+    """
+
+    @abc.abstractmethod
+    def patch(self, modules: List[str] = None):
+        """Patch modules for instrumentation
+
+        If modules are None, it should patch
+        all supported modules by the provider.
+
+        Parameters
+        ----------
+        modules : List[str], optional
+            List of modules to be pathced, by default None
+            e.g. `['boto3', 'requests']`
+        """
+        raise NotImplementedError
+
+    @abc.abstractmethod
+    def create_subsegment(self, name: str):
+        """Creates subsegment/span with a given name
+
+        Parameters
+        ----------
+        name : str
+            Subsegment/span name
+        """
+        raise NotImplementedError
+
+    @abc.abstractmethod
+    def end_subsegment(self):
+        """Ends an existing subsegment"""
+        raise NotImplementedError
+
+    @abc.abstractmethod
+    def put_metadata(self, key: str, value: Any, namespace: str = None):
+        """Adds metadata to existing segment/span or subsegment
+
+        Parameters
+        ----------
+        key : str
+            Metadata key
+        value : Any
+            Metadata value
+        namespace : str, optional
+            Metadata namespace, by default None
+        """
+        raise NotImplementedError
+
+    @abc.abstractmethod
+    def put_annotation(self, key: str, value: Any):
+        """Adds annotation/label to existing segment/span or subsegment
+
+        Parameters
+        ----------
+        key : str
+            Annotation/label key
+        value : Any
+            Annotation/label value
+        """
+        raise NotImplementedError
+
+    @abc.abstractmethod
+    def disable_tracing_provider(self):
+        """Forcefully disables tracing provider"""
+        raise NotImplementedError
+
+
+class XrayProvider(TracerProvider):
+    """X-Ray Tracer provider
+
+    It implements all basic ``aws_lambda_powertools.tracing.base.TracerProvider` methods,
+    and automatically annotates cold start on first subsegment created.
+
+    Parameters
+    ----------
+    client : aws_xray_sdk.core.xray_recorder
+        X-Ray recorder client
+    """
+
+    def __init__(self, client: aws_xray_sdk.core.xray_recorder = aws_xray_sdk.core.xray_recorder):
+        self.client = client
+
+    def create_subsegment(self, name: str) -> aws_xray_sdk.core.models.subsegment:
+        """Creates subsegment/span with a given name
+
+        Parameters
+        ----------
+        name : str
+            Subsegment name
+
+        Example
+        -------
+
+        **Creates a subsegment**
+
+            self.create_subsegment(name="a meaningful name")
+
+        Returns
+        -------
+        aws_xray_sdk.core.models.subsegment
+            AWS X-Ray Subsegment
+        """
+        # Will no longer be needed once #155 is resolved
+        # https://github.com/aws/aws-xray-sdk-python/issues/155
+        subsegment = self.client.begin_subsegment(name=name)
+        global is_cold_start
+        if is_cold_start:
+            logger.debug("Annotating cold start")
+            subsegment.put_annotation(key="ColdStart", value=True)
+            is_cold_start = False
+
+        return subsegment
+
+    def end_subsegment(self):
+        """Ends an existing subsegment"""
+        self.client.end_subsegment()
+
+    def put_annotation(self, key, value):
+        """Adds annotation to existing segment or subsegment
+
+        Example
+        -------
+        Custom annotation for a pseudo service named payment
+
+            tracer = Tracer(service="payment")
+            tracer.put_annotation("PaymentStatus", "CONFIRMED")
+
+        Parameters
+        ----------
+        key : str
+            Annotation key (e.g. PaymentStatus)
+        value : Any
+            Value for annotation (e.g. "CONFIRMED")
+        """
+        self.client.put_annotation(key=key, value=value)
+
+    def put_metadata(self, key, value, namespace=None):
+        """Adds metadata to existing segment or subsegment
+
+        Parameters
+        ----------
+        key : str
+            Metadata key
+        value : object
+            Value for metadata
+        namespace : str, optional
+            Namespace that metadata will lie under, by default None
+
+        Example
+        -------
+        Custom metadata for a pseudo service named payment
+
+            tracer = Tracer(service="payment")
+            response = collect_payment()
+            tracer.put_metadata("Payment collection", response)
+        """
+        self.client.put_metadata(key=key, value=value, namespace=namespace)
+
+    def patch(self, modules: List[str] = None):
+        """Patch modules for instrumentation.
+
+        Patches all supported modules by default if none are given.
+
+        Parameters
+        ----------
+        modules : List[str]
+            List of modules to be patched, optional by default
+        """
+        if modules is None:
+            aws_xray_sdk.core.patch_all()
+        else:
+            aws_xray_sdk.core.patch(modules)
+
+    def disable_tracing_provider(self):
+        """Forcefully disables X-Ray tracing globally"""
+        aws_xray_sdk.global_sdk_config.set_sdk_enabled(False)

python/aws_lambda_powertools/tracing/exceptions.py

@@ -0,0 +1,10 @@
+class InvalidTracerProviderError(Exception):
+    """When given provider doesn't implement `aws_lambda_powertools.tracing.base.TracerProvider`"""
+
+    pass
+
+
+class TracerProviderNotInitializedError(Exception):
+    """When given provider isn't initialized/bound"""
+
+    pass