diff --git a/Makefile b/Makefile index 6b9d6ef0963..5b8e9b0d689 100644 --- a/Makefile +++ b/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 diff --git a/aws_lambda_powertools/tracing/tracer.py b/aws_lambda_powertools/tracing/tracer.py index 2beab0483be..70580663e7b 100644 --- a/aws_lambda_powertools/tracing/tracer.py +++ b/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__ diff --git a/docs/core/tracer.md b/docs/core/tracer.md index e2e2df52e18..9e94d2549d9 100644 --- a/docs/core/tracer.md +++ b/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 diff --git a/tests/unit/test_tracing.py b/tests/unit/test_tracing.py index 7c8b6244f01..55273b072c6 100644 --- a/tests/unit/test_tracing.py +++ b/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)