From d7ab602e500ef5714ebfdabace33f85c0cc8bc66 Mon Sep 17 00:00:00 2001 From: heitorlessa Date: Tue, 30 Nov 2021 15:00:11 +0100 Subject: [PATCH 1/4] fix: lazy loading core attr, custom provider patch --- Makefile | 3 ++ aws_lambda_powertools/tracing/tracer.py | 41 +++++++++++++++++++------ tests/unit/test_tracing.py | 19 ++++++++++++ 3 files changed, 54 insertions(+), 9 deletions(-) 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..eb3b342366c 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: @@ -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__ diff --git a/tests/unit/test_tracing.py b/tests/unit/test_tracing.py index 7c8b6244f01..80191b8d307 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 @@ -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") From 67a57eb8b64ae18c99a9771a6dad91c9f56be4d9 Mon Sep 17 00:00:00 2001 From: heitorlessa Date: Tue, 30 Nov 2021 15:39:54 +0100 Subject: [PATCH 2/4] feat: add Service annotation when service is present --- aws_lambda_powertools/tracing/tracer.py | 5 ++++- tests/unit/test_tracing.py | 23 +++++++++++++++++++++-- 2 files changed, 25 insertions(+), 3 deletions(-) diff --git a/aws_lambda_powertools/tracing/tracer.py b/aws_lambda_powertools/tracing/tracer.py index eb3b342366c..2b110693823 100644 --- a/aws_lambda_powertools/tracing/tracer.py +++ b/aws_lambda_powertools/tracing/tracer.py @@ -136,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, @@ -317,6 +317,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) diff --git a/tests/unit/test_tracing.py b/tests/unit/test_tracing.py index 80191b8d307..55273b072c6 100644 --- a/tests/unit/test_tracing.py +++ b/tests/unit/test_tracing.py @@ -589,10 +589,10 @@ 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 == mocker.call(key="ColdStart", value=False) + 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): @@ -609,3 +609,22 @@ def handler(event, context): # 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) From ec939c5ad89c68cd2f01dff737b2f49b2b05f56e Mon Sep 17 00:00:00 2001 From: heitorlessa Date: Tue, 30 Nov 2021 15:42:21 +0100 Subject: [PATCH 3/4] docs: mention new Service annotation --- docs/core/tracer.md | 1 + 1 file changed, 1 insertion(+) 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 From b0c018829537f4b236fdcc6b1adc3478ec94b863 Mon Sep 17 00:00:00 2001 From: heitorlessa Date: Tue, 30 Nov 2021 15:58:38 +0100 Subject: [PATCH 4/4] chore: fix test coverage --- aws_lambda_powertools/tracing/tracer.py | 22 +++++++++++----------- 1 file changed, 11 insertions(+), 11 deletions(-) diff --git a/aws_lambda_powertools/tracing/tracer.py b/aws_lambda_powertools/tracing/tracer.py index 2b110693823..70580663e7b 100644 --- a/aws_lambda_powertools/tracing/tracer.py +++ b/aws_lambda_powertools/tracing/tracer.py @@ -166,19 +166,9 @@ def __init__( if self.auto_patch: self.patch(modules=patch_modules) - if "aws_xray_sdk" in self.provider.__module__: + if self._is_xray_provider(): 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 @@ -776,5 +766,15 @@ def _patch_xray_provider(self): 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__