address ruben comments

Author: roger-zhangg

aws_lambda_powertools/tracing/provider/__init__.py

@@ -1,5 +0,0 @@
-from .datadog.dd_tracer import DDTraceProvider
-from .otel.otel_tracer import OtelProvider
-from .xray.xray_tracer import XrayProvider
-
-__all__ = ["DDTraceProvider", "OtelProvider", "XrayProvider"]

aws_lambda_powertools/tracing/provider/base.py

@@ -1,78 +1,105 @@
 import abc
-import numbers
 from contextlib import asynccontextmanager, contextmanager
-from typing import AsyncGenerator, Generator, Sequence, Union
+from typing import Any, AsyncGenerator, Generator, Sequence
 
 
 class BaseSpan(abc.ABC):
-    """Holds common properties and methods on span."""
+    """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: Union[str, numbers.Number, bool], **kwargs) -> None:
-        """set attribute for span with a key-value pair.
+    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: Union[str, numbers.Number, bool]
+        value: Any
             Attribute value
+        kwargs: Optional[dict]
+            Optional parameters
         """
 
     @abc.abstractmethod
     def record_exception(self, exception: BaseException, **kwargs):
-        """Add an exception to trace entities.
+        """Records an exception to this Span.
 
         Parameters
         ----------
         exception: Exception
-            Caught exception
-            Output from `traceback.extract_stack()`.
+            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]:
-        """Return a span context manger.
+        """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 span
+            Optional parameters to be propagated to the span
         """
 
     @abc.abstractmethod
     @asynccontextmanager
     def trace_async(self, name: str, **kwargs) -> AsyncGenerator[BaseSpan, None]:
-        """Return a async span context manger.
+        """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 span
+            Optional parameters to be propagated to the span
         """
 
     @abc.abstractmethod
-    def set_attribute(self, key: str, value: Union[str, numbers.Number, bool], **kwargs) -> None:
-        """set attribute on current active trace entity with a key-value pair.
+    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: Union[str, numbers.Number, bool]
+        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 supported libraries
+        """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
         ----------

aws_lambda_powertools/tracing/provider/otel/otel_tracer.py

@@ -15,7 +15,7 @@ class OtelSpan(BaseSpan):
     def __init__(self, otel_span=otel_trace.Span):
         self.otel_span = otel_span
 
-    def set_attribute(self, key: str, value: str | float | bool, **kwargs) -> None:  # type: ignore[override]
+    def set_attribute(self, key: str, value: str | float | bool, **kwargs) -> None:
         if not isinstance(value, (str, bool, int, float)):
             # convert value to str if value is not a supported structur
             value = str(value)
@@ -69,7 +69,7 @@ async def trace_async(self, name: str, **kwargs) -> AsyncGenerator[OtelSpan, Non
 
     in_subsegment_async = trace_async
 
-    def set_attribute(self, key: str, value: str | float | bool, **kwargs) -> None:  # type: ignore[override]
+    def set_attribute(self, key: str, value: str | float | bool, **kwargs) -> None:
         active_span = otel_trace.get_current_span()
         active_span.set_attribute(key=key, value=value)
 

aws_lambda_powertools/tracing/provider/xray/xray_tracer.py

@@ -2,7 +2,7 @@
 
 from contextlib import asynccontextmanager, contextmanager
 from numbers import Number
-from typing import Any, AsyncGenerator, Generator, Sequence
+from typing import Any, AsyncGenerator, Generator, Sequence, Union
 
 from ....shared import constants
 from ....shared.lazy_import import LazyLoader
@@ -21,7 +21,7 @@ def __init__(self, subsegment):
         self.add_exception = self.subsegment.add_exception
         self.close = self.subsegment.close
 
-    def set_attribute(self, key: str, value: str | Number | bool, **kwargs) -> None:
+    def set_attribute(self, key: str, value: Any, **kwargs) -> None:
         if kwargs.get("namespace", "") != "":
             self.put_metadata(key=key, value=value, namespace=kwargs["namespace"])
         else:
@@ -42,21 +42,23 @@ def __init__(self, xray_recorder=None):
 
     @contextmanager
     def trace(self, name: str, **kwargs) -> Generator[XraySpan, None, None]:
-        with self.in_subsegment(name=name) as sub_segment:
+        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) as subsegment:
+        async with self.in_subsegment_async(name=name, **kwargs) as subsegment:
             yield XraySpan(subsegment=subsegment)
 
-    def set_attribute(self, key: str, value: str | Number | bool, **kwargs) -> None:
-        if kwargs.get("namespace", "") != "":
-            self.put_metadata(key=key, value=value, namespace=kwargs["namespace"])
-        else:
+    def set_attribute(self, key: str, value: Any, **kwargs) -> None:
+        # for x_ray, put annotation support str, Number, bool
+        # we use put_metadata if any unsupported values are provided
+        if isinstance(value, (str, Number, bool)):
             self.put_annotation(key=key, value=value)
+        else:
+            self.put_metadata(key=key, value=value, namespace=kwargs["namespace"])
 
-    def put_annotation(self, key: str, value: str | Number | bool) -> None:
+    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:

aws_lambda_powertools/tracing/tracer.py

@@ -15,8 +15,8 @@
 )
 from aws_lambda_powertools.shared.lazy_import import LazyLoader
 from aws_lambda_powertools.shared.types import AnyCallableT
-from aws_lambda_powertools.tracing.provider import XrayProvider
 from aws_lambda_powertools.tracing.provider.base import BaseProvider, BaseSpan
+from aws_lambda_powertools.tracing.provider.xray.xray_tracer import XrayProvider
 
 is_cold_start = True
 logger = logging.getLogger(__name__)
@@ -178,15 +178,17 @@ def __init__(
         if self._is_xray_provider():
             self._disable_xray_trace_batching()
 
-    def set_attribute(self, key: str, value: Union[str, numbers.Number, bool]):
+    def set_attribute(self, key: str, value: Any, **kwargs):
         """Set attribute on current active trace entity with a key-value pair.
 
         Parameters
         ----------
         key : str
             attribute key
-        value : Union[str, numbers.Number, bool]
+        value : Any
             Value for attribute
+        kwargs: Optional[dict]
+            Optional parameters to be passed to provider.set_attributes
 
         Example
         -------
@@ -201,7 +203,10 @@ def set_attribute(self, key: str, value: Union[str, numbers.Number, bool]):
 
         logger.debug(f"setting attribute on key '{key}' with '{value}'")
 
-        self.provider.set_attribute(key=key, value=value)
+        namespace = kwargs.get("namespace") or self.service
+        kwargs.update({"namespace": namespace})
+
+        self.provider.set_attribute(key=key, value=value, **kwargs)
 
     def put_annotation(self, key: str, value: Union[str, numbers.Number, bool]):
         """Adds annotation to existing segment or subsegment