feat(tracer): add service annotation when service is set (#861)

Author: heitorlessa

Makefile

@@ -20,6 +20,9 @@ test:
 	poetry run pytest -m "not perf" --cov=aws_lambda_powertools --cov-report=xml
 	poetry run pytest --cache-clear tests/performance
 
+unit-test:
+	poetry run pytest tests/unit
+
 coverage-html:
 	poetry run pytest -m "not perf" --cov=aws_lambda_powertools --cov-report=html
 

aws_lambda_powertools/tracing/tracer.py

@@ -17,7 +17,6 @@
 logger = logging.getLogger(__name__)
 
 aws_xray_sdk = LazyLoader(constants.XRAY_SDK_MODULE, globals(), constants.XRAY_SDK_MODULE)
-aws_xray_sdk.core = LazyLoader(constants.XRAY_SDK_CORE_MODULE, globals(), constants.XRAY_SDK_CORE_MODULE)  # type: ignore # noqa: E501
 
 
 class Tracer:
@@ -137,7 +136,7 @@ def handler(event: dict, context: Any) -> Dict:
     """
 
     _default_config: Dict[str, Any] = {
-        "service": "service_undefined",
+        "service": "",
         "disabled": False,
         "auto_patch": True,
         "patch_modules": None,
@@ -156,7 +155,7 @@ def __init__(
         self.__build_config(
             service=service, disabled=disabled, auto_patch=auto_patch, patch_modules=patch_modules, provider=provider
         )
-        self.provider: BaseProvider = self._config["provider"]
+        self.provider = self._config["provider"]
         self.disabled = self._config["disabled"]
         self.service = self._config["service"]
         self.auto_patch = self._config["auto_patch"]
@@ -167,10 +166,8 @@ def __init__(
         if self.auto_patch:
             self.patch(modules=patch_modules)
 
-        # Set the streaming threshold to 0 on the default recorder to force sending
-        # subsegments individually, rather than batching them.
-        # See https://github.com/awslabs/aws-lambda-powertools-python/issues/283
-        aws_xray_sdk.core.xray_recorder.configure(streaming_threshold=0)  # noqa: E800
+        if self._is_xray_provider():
+            self._disable_xray_trace_batching()
 
     def put_annotation(self, key: str, value: Union[str, numbers.Number, bool]):
         """Adds annotation to existing segment or subsegment
@@ -239,9 +236,9 @@ def patch(self, modules: Optional[Sequence[str]] = None):
             return
 
         if modules is None:
-            aws_xray_sdk.core.patch_all()
+            self.provider.patch_all()
         else:
-            aws_xray_sdk.core.patch(modules)
+            self.provider.patch(modules)
 
     def capture_lambda_handler(
         self,
@@ -310,6 +307,9 @@ def decorate(event, context, **kwargs):
                 if is_cold_start:
                     is_cold_start = False
 
+                if self.service:
+                    subsegment.put_annotation(key="Service", value=self.service)
+
                 try:
                     logger.debug("Calling lambda handler")
                     response = lambda_handler(event, context, **kwargs)
@@ -743,7 +743,8 @@ def __build_config(
         is_disabled = disabled if disabled is not None else self._is_tracer_disabled()
         is_service = resolve_env_var_choice(choice=service, env=os.getenv(constants.SERVICE_NAME_ENV))
 
-        self._config["provider"] = provider or self._config["provider"] or aws_xray_sdk.core.xray_recorder
+        # Logic: Choose overridden option first, previously cached config, or default if available
+        self._config["provider"] = provider or self._config["provider"] or self._patch_xray_provider()
         self._config["auto_patch"] = auto_patch if auto_patch is not None else self._config["auto_patch"]
         self._config["service"] = is_service or self._config["service"]
         self._config["disabled"] = is_disabled or self._config["disabled"]
@@ -752,3 +753,28 @@ def __build_config(
     @classmethod
     def _reset_config(cls):
         cls._config = copy.copy(cls._default_config)
+
+    def _patch_xray_provider(self):
+        # Due to Lazy Import, we need to activate `core` attrib via import
+        # we also need to include `patch`, `patch_all` methods
+        # to ensure patch calls are done via the provider
+        from aws_xray_sdk.core import xray_recorder
+
+        provider = xray_recorder
+        provider.patch = aws_xray_sdk.core.patch
+        provider.patch_all = aws_xray_sdk.core.patch_all
+
+        return provider
+
+    def _disable_xray_trace_batching(self):
+        """Configure X-Ray SDK to send subsegment individually over batching
+        Known issue: https://github.com/awslabs/aws-lambda-powertools-python/issues/283
+        """
+        if self.disabled:
+            logger.debug("Tracing has been disabled, aborting streaming override")
+            return
+
+        aws_xray_sdk.core.xray_recorder.configure(streaming_threshold=0)
+
+    def _is_xray_provider(self):
+        return "aws_xray_sdk" in self.provider.__module__

docs/core/tracer.md

@@ -58,6 +58,7 @@ You can quickly start by importing the `Tracer` class, initialize it outside the
 When using this `capture_lambda_handler` decorator, Tracer performs these additional tasks to ease operations:
 
 * Creates a `ColdStart` annotation to easily filter traces that have had an initialization overhead
+* Creates a `Service` annotation if `service` parameter or `POWERTOOLS_SERVICE_NAME` is set
 * Captures any response, or full exceptions generated by the handler, and include as tracing metadata
 
 ### Annotations & Metadata

tests/unit/test_tracing.py

@@ -44,6 +44,9 @@ def in_subsegment(self, *args, **kwargs):
         def patch(self, *args, **kwargs):
             return self.patch_mock(*args, **kwargs)
 
+        def patch_all(self):
+            ...
+
     return CustomProvider
 
 
@@ -586,7 +589,42 @@ def handler(event, context):
     handler({}, mocker.MagicMock())
 
     # THEN
-    assert in_subsegment_mock.put_annotation.call_args == mocker.call(key="ColdStart", value=True)
+    assert in_subsegment_mock.put_annotation.call_args_list[0] == mocker.call(key="ColdStart", value=True)
+
+    handler({}, mocker.MagicMock())
+    assert in_subsegment_mock.put_annotation.call_args_list[2] == mocker.call(key="ColdStart", value=False)
+
+
+def test_tracer_lambda_handler_add_service_annotation(mocker, dummy_response, provider_stub, in_subsegment_mock):
+    # GIVEN
+    provider = provider_stub(in_subsegment=in_subsegment_mock.in_subsegment)
+    tracer = Tracer(provider=provider, service="booking")
+
+    # WHEN
+    @tracer.capture_lambda_handler
+    def handler(event, context):
+        return dummy_response
 
     handler({}, mocker.MagicMock())
-    assert in_subsegment_mock.put_annotation.call_args == mocker.call(key="ColdStart", value=False)
+
+    # THEN
+    assert in_subsegment_mock.put_annotation.call_args == mocker.call(key="Service", value="booking")
+
+
+def test_tracer_lambda_handler_do_not_add_service_annotation_when_missing(
+    mocker, dummy_response, provider_stub, in_subsegment_mock
+):
+    # GIVEN
+    provider = provider_stub(in_subsegment=in_subsegment_mock.in_subsegment)
+    tracer = Tracer(provider=provider)
+
+    # WHEN
+    @tracer.capture_lambda_handler
+    def handler(event, context):
+        return dummy_response
+
+    handler({}, mocker.MagicMock())
+
+    # THEN
+    assert in_subsegment_mock.put_annotation.call_count == 1
+    assert in_subsegment_mock.put_annotation.call_args == mocker.call(key="ColdStart", value=True)