fix: lazy loading core attr, custom provider patch

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:
@@ -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,18 @@ 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 "aws_xray_sdk" in self.provider.__module__:
+            self._disable_xray_trace_batching()
+
+    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 put_annotation(self, key: str, value: Union[str, numbers.Number, bool]):
         """Adds annotation to existing segment or subsegment
@@ -239,9 +246,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,
@@ -743,7 +750,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 +760,18 @@ 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 _is_xray_provider(self):
+        return "aws_xray_sdk" in self.provider.__module__

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
 
 
@@ -590,3 +593,19 @@ def handler(event, context):
 
     handler({}, mocker.MagicMock())
     assert in_subsegment_mock.put_annotation.call_args == 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())
+
+    # THEN
+    assert in_subsegment_mock.put_annotation.call_args == mocker.call(key="Service", value="booking")