add set_attribute in tracer

Author: roger-zhangg

aws_lambda_powertools/tracing/provider/otel/otel_tracer.py

@@ -80,4 +80,5 @@ def patch(self, modules: Sequence[str]) -> None:
         pass
 
     def patch_all(self) -> None:
+        # OTEL sdk doesn't have patch
         pass

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 AsyncGenerator, Generator, Sequence
+from typing import Any, AsyncGenerator, Generator, Sequence
 
 from ....shared import constants
 from ....shared.lazy_import import LazyLoader
@@ -37,12 +37,8 @@ def __init__(self, xray_recorder=None):
         if not xray_recorder:
             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.in_subsegment = self.recorder.in_subsegment
         self.in_subsegment_async = self.recorder.in_subsegment_async
-        self.put_annotation = self.recorder.put_annotation
-        self.put_metadata = self.recorder.put_metadata
 
     @contextmanager
     def trace(self, name: str, **kwargs) -> Generator[XraySpan, None, None]:
@@ -60,9 +56,14 @@ def set_attribute(self, key: str, value: str | Number | bool, **kwargs) -> None:
         else:
             self.put_annotation(key=key, value=value)
 
+    def put_annotation(self, key: str, value: 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:
-        # defined in init
-        pass
+        return aws_xray_sdk.core.patch(modules)
 
     def patch_all(self) -> None:
-        pass
+        return aws_xray_sdk.core.patch_all()

aws_lambda_powertools/tracing/tracer.py

@@ -11,8 +11,8 @@
 from aws_lambda_powertools.shared.functions import resolve_env_var_choice, resolve_truthy_env_var_choice
 from aws_lambda_powertools.shared.lazy_import import LazyLoader
 from aws_lambda_powertools.shared.types import AnyCallableT
-from aws_lambda_powertools.tracing.base import BaseProvider, BaseSegment
 from aws_lambda_powertools.tracing.provider import XrayProvider
+from aws_lambda_powertools.tracing.provider.base import BaseProvider, BaseSpan
 
 is_cold_start = True
 logger = logging.getLogger(__name__)
@@ -172,9 +172,17 @@ def __init__(
             self.patch(modules=patch_modules)
 
         if self._is_xray_provider():
-            # TO why?
             self._disable_xray_trace_batching()
 
+    def set_attribute(self, key: str, value: Union[str, numbers.Number, bool]):
+        if self.disabled:
+            logger.debug("Tracing has been disabled, aborting put_annotation")
+            return
+
+        logger.debug(f"setting attribute on key '{key}' with '{value}'")
+
+        self.provider.set_attribute(key=key, value=value)
+
     def put_annotation(self, key: str, value: Union[str, numbers.Number, bool]):
         """Adds annotation to existing segment or subsegment
 
@@ -197,7 +205,11 @@ def put_annotation(self, key: str, value: Union[str, numbers.Number, bool]):
             return
 
         logger.debug(f"Annotating on key '{key}' with '{value}'")
-        self.provider.put_annotation(key=key, value=value)
+
+        if not self._is_custom_provider():
+            self.provider.put_annotation(key=key, value=value)  # type: ignore
+        else:
+            self.provider.set_attribute(key=key, value=value)
 
     def put_metadata(self, key: str, value: Any, namespace: Optional[str] = None):
         """Adds metadata to existing segment or subsegment
@@ -225,7 +237,10 @@ def put_metadata(self, key: str, value: Any, namespace: Optional[str] = None):
 
         namespace = namespace or self.service
         logger.debug(f"Adding metadata on key '{key}' with '{value}' at namespace '{namespace}'")
-        self.provider.put_metadata(key=key, value=value, namespace=namespace)
+        if not self._is_custom_provider():
+            self.provider.put_metadata(key=key, value=value, namespace=namespace)  # type: ignore
+        else:
+            self.provider.set_attribute(key=f"{namespace}.{key}", value=value)
 
     def patch(self, modules: Optional[Sequence[str]] = None):
         """Patch modules for instrumentation.
@@ -700,7 +715,7 @@ def _add_response_as_metadata(
         self,
         method_name: Optional[str] = None,
         data: Optional[Any] = None,
-        subsegment: Optional[BaseSegment] = None,
+        subsegment: Optional[BaseSpan] = None,
         capture_response: Optional[Union[bool, str]] = None,
     ):
         """Add response as metadata for given subsegment
@@ -718,14 +733,16 @@ def _add_response_as_metadata(
         """
         if data is None or not capture_response or subsegment is None:
             return
-
-        subsegment.put_metadata(key=f"{method_name} response", value=data, namespace=self.service)
+        if not self._is_custom_provider():
+            subsegment.put_metadata(key=f"{method_name} response", value=data, namespace=self.service)  # type: ignore
+        else:
+            subsegment.set_attribute(key=f"{method_name} response", value=data, namespace=self.service)
 
     def _add_full_exception_as_metadata(
         self,
         method_name: str,
         error: Exception,
-        subsegment: BaseSegment,
+        subsegment: BaseSpan,
         capture_error: Optional[bool] = None,
     ):
         """Add full exception object as metadata for given subsegment
@@ -744,7 +761,14 @@ def _add_full_exception_as_metadata(
         if not capture_error:
             return
 
-        subsegment.put_metadata(key=f"{method_name} error", value=error, namespace=self.service)
+        if not self._is_custom_provider():
+            subsegment.put_metadata(key=f"{method_name} error", value=error, namespace=self.service)  # type: ignore
+        else:
+            subsegment.set_attribute(
+                key=f"{self.service}.{method_name} error",
+                value=str(error),
+                namespace=self.service,
+            )
 
     @staticmethod
     def _disable_tracer_provider():
@@ -827,6 +851,9 @@ def _disable_xray_trace_batching(self):
     def _is_xray_provider(self):
         return isinstance(self.provider, XrayProvider)
 
+    def _is_custom_provider(self):
+        return not self._is_xray_provider and isinstance(self.provider, BaseProvider)
+
     def ignore_endpoint(self, hostname: Optional[str] = None, urls: Optional[List[str]] = None):
         """If you want to ignore certain httplib requests you can do so based on the hostname or URL that is being
         requested.

examples/tracer/src/sdk_escape_hatch.py

@@ -11,9 +11,9 @@ def collect_payment(charge_id: str) -> str:
 @tracer.capture_lambda_handler
 def lambda_handler(event: dict, context: LambdaContext) -> str:
     charge_id = event.get("charge_id", "")
-    with tracer.provider.in_subsegment("## collect_payment") as subsegment:
-        subsegment.put_annotation(key="PaymentId", value=charge_id)
+    with tracer.provider.trace("## collect_payment") as span:
+        span.set_attribute(key="PaymentId", value=charge_id)
         ret = collect_payment(charge_id=charge_id)
-        subsegment.put_metadata(key="payment_response", value=ret)
+        span.set_attribute(key="payment_response", value=ret)
 
     return ret