refactor tracer provider, add a xray provider

Author: roger-zhangg

aws_lambda_powertools/tracing/provider/__init__.py

@@ -0,0 +1,3 @@
+from .xray_provider import Xray_provider
+
+__all__ = ["Xray_provider"]

aws_lambda_powertools/tracing/provider/base.py

@@ -2,7 +2,7 @@
 import numbers
 import traceback
 from contextlib import contextmanager
-from typing import Any, Generator, List, Optional, Sequence, Union
+from typing import Any, Callable, Generator, List, Optional, Sequence, Union
 
 
 class BaseSegment(abc.ABC):
@@ -72,11 +72,30 @@ def add_exception(self, exception: BaseException, stack: List[traceback.StackSum
             Whether it's a client error (False) or downstream service error (True), by default False
         """
 
+    # new function below
+    @abc.abstractmethod
+    def ignore_endpoint(self, hostname: Optional[str] = None, urls: Optional[List[str]] = None):
+        """To ignore the endpoints you don't want requests to be traced,
+        perhaps due to the volume of calls or sensitive URLs."""
+
+    @abc.abstractmethod
+    def inject_context(self, context):
+        """To inject missing context/information like service name"""
+
+    @abc.abstractmethod
+    def capture_method_async(
+        self,
+        method: Callable,
+        capture_response: Optional[Union[bool, str]] = None,
+        capture_error: Optional[Union[bool, str]] = None,
+    ):
+        """To capture async method"""
+
 
 class BaseProvider(abc.ABC):
     @abc.abstractmethod
     @contextmanager
-    def in_subsegment(self, name=None, **kwargs) -> Generator[BaseSegment, None, None]:
+    def trace(self, name=None, **kwargs) -> Generator[BaseSegment, None, None]:
         """Return a subsegment context manger.
 
         Parameters
@@ -89,7 +108,7 @@ def in_subsegment(self, name=None, **kwargs) -> Generator[BaseSegment, None, Non
 
     @abc.abstractmethod
     @contextmanager
-    def in_subsegment_async(self, name=None, **kwargs) -> Generator[BaseSegment, None, None]:
+    def trace_async(self, name=None, **kwargs) -> Generator[BaseSegment, None, None]:
         """Return a subsegment async context manger.
 
         Parameters
@@ -100,6 +119,18 @@ def in_subsegment_async(self, name=None, **kwargs) -> Generator[BaseSegment, Non
             Optional parameters to be propagated to segment
         """
 
+    @abc.abstractmethod
+    def start_span(self, name=None) -> Generator[BaseSegment, None, None]:
+        """This method is proposed as a solution as it exists for other providers
+        This method is responsible for starting the trace. This might involve initializing some data structures,
+        connecting to an external service, or performing some other setup work"""
+
+    @abc.abstractmethod
+    def end_span(self, span: BaseSegment):
+        """This method is proposed as a solution as it exists for other providers.
+        This method is responsible for ending the tracing of a span. This might involve finalizing data structures,
+        sending data to an external service, or performing some other cleanup work"""
+
     @abc.abstractmethod
     def put_annotation(self, key: str, value: Union[str, numbers.Number, bool]) -> None:
         """Annotate current active trace entity with a key-value pair.
@@ -130,6 +161,10 @@ def put_metadata(self, key: str, value: Any, namespace: str = "default") -> None
             Metadata namespace, by default 'default'
         """
 
+    @abc.abstractmethod
+    def put_exception(self, *args, **kwargs) -> None:
+        """Add exception to the current active trace entity."""
+
     @abc.abstractmethod
     def patch(self, modules: Sequence[str]) -> None:
         """Instrument a set of supported libraries

aws_lambda_powertools/tracing/provider/xray_provider.py

@@ -0,0 +1,70 @@
+from contextlib import contextmanager
+from typing import Generator, Optional
+
+from ...shared import constants
+from ...shared.lazy_import import LazyLoader
+from .base import BaseProvider, BaseSegment
+
+aws_xray_sdk = LazyLoader(constants.XRAY_SDK_MODULE, globals(), constants.XRAY_SDK_MODULE)
+
+
+class Xray_provider(BaseProvider):
+    def __init__(self):
+        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.put_annotation = self.recorder.put_annotation
+        self.put_metadata = self.recorder.put_metadata
+
+    @contextmanager
+    def trace(self, name=None, **kwargs) -> Generator[BaseSegment, None, None]:
+        """Return a subsegment context manger.
+
+        Parameters
+        ----------
+        name: str
+            Subsegment name
+        kwargs: Optional[dict]
+            Optional parameters to be propagated to segment
+        """
+        return self.recorder.in_subsegment(name)
+
+    @contextmanager
+    def trace_async(self, name=None, **kwargs) -> Generator[BaseSegment, None, None]:
+        """Return a subsegment async context manger.
+
+        Parameters
+        ----------
+        name: str
+            Subsegment name
+        kwargs: Optional[dict]
+            Optional parameters to be propagated to segment
+        """
+        return self.recorder.in_subsegment_async(name)
+
+    # we don't need to start,end explicitly in xray
+    def start_span(self, name=None) -> Generator[BaseSegment, None, None]:
+        """This method is proposed as a solution as it exists for other providers
+        This method is responsible for starting the trace. This might involve initializing some data structures,
+        connecting to an external service, or performing some other setup work"""
+        raise Exception("Not implemented")
+
+    def end_span(self, span: BaseSegment):
+        """This method is proposed as a solution as it exists for other providers.
+        This method is responsible for ending the tracing of a span. This might involve finalizing data structures,
+        sending data to an external service, or performing some other cleanup work"""
+        raise Exception("Not implemented")
+
+    def put_exception(
+        self,
+        method_name: str,
+        error: Exception,
+        subsegment: BaseSegment,
+        capture_error: Optional[bool] = None,
+        namespace: str = None,
+    ):
+        if not capture_error:
+            return
+        self.put_metadata(key=f"{method_name} error", value=error, namespace=namespace)