From 902f8c31ed31988a8b7116f9cc3c720820c49a6d Mon Sep 17 00:00:00 2001 From: Michael Brewer Date: Fri, 9 Jul 2021 12:10:33 -0700 Subject: [PATCH 1/4] feat(logger): add option to remove correlation_id Call remove_keys when the correlation_id is set to None --- aws_lambda_powertools/logging/logger.py | 11 +++++++---- tests/functional/test_logger.py | 14 ++++++++++++++ 2 files changed, 21 insertions(+), 4 deletions(-) diff --git a/aws_lambda_powertools/logging/logger.py b/aws_lambda_powertools/logging/logger.py index 689409d9813..b370b52643f 100644 --- a/aws_lambda_powertools/logging/logger.py +++ b/aws_lambda_powertools/logging/logger.py @@ -387,15 +387,18 @@ def structure_logs(self, append: bool = False, **keys): formatter = self.logger_formatter or LambdaPowertoolsFormatter(**log_keys) self.registered_handler.setFormatter(formatter) - def set_correlation_id(self, value: str): + def set_correlation_id(self, value: Optional[str]): """Sets the correlation_id in the logging json Parameters ---------- - value : str - Value for the correlation id + value : str, optional + Value for the correlation id. None will remove the correlation_id """ - self.append_keys(correlation_id=value) + if value is None: + self.remove_keys(["correlation_id"]) + else: + self.append_keys(correlation_id=value) @staticmethod def _get_log_level(level: Union[str, int, None]) -> Union[str, int]: diff --git a/tests/functional/test_logger.py b/tests/functional/test_logger.py index 44249af6250..9494d39344b 100644 --- a/tests/functional/test_logger.py +++ b/tests/functional/test_logger.py @@ -460,6 +460,20 @@ def handler(event, _): assert request_id == log["correlation_id"] +def test_logger_set_correlation_id_to_none(lambda_context, stdout, service_name): + # GIVEN a logger with a correlation_id set + logger = Logger(service=service_name, stream=stdout) + logger.set_correlation_id("foo") + + # WHEN calling set_correlation_id with None + logger.set_correlation_id(None) + + # THEN there should be no correlation_id + logger.info("Foo") + log = capture_logging_output(stdout) + assert "correlation_id" not in log + + def test_logger_set_correlation_id_path(lambda_context, stdout, service_name): # GIVEN logger = Logger(service=service_name, stream=stdout) From db6870d96cbadabf8b9b6fc1e32b7fc21d2b2f01 Mon Sep 17 00:00:00 2001 From: Michael Brewer Date: Sat, 10 Jul 2021 09:03:44 -0700 Subject: [PATCH 2/4] feat(logging): add get_correlation_id --- aws_lambda_powertools/logging/logger.py | 10 ++++++++++ tests/functional/test_logger.py | 12 ++++++++++++ 2 files changed, 22 insertions(+) diff --git a/aws_lambda_powertools/logging/logger.py b/aws_lambda_powertools/logging/logger.py index b370b52643f..1152480f17e 100644 --- a/aws_lambda_powertools/logging/logger.py +++ b/aws_lambda_powertools/logging/logger.py @@ -400,6 +400,16 @@ def set_correlation_id(self, value: Optional[str]): else: self.append_keys(correlation_id=value) + def get_correlation_id(self) -> Optional[str]: + """Gets the correlation_id in the logging json + + Returns + ------- + str, optional + Value for the correlation id + """ + return self.registered_formatter.log_format.get("correlation_id") + @staticmethod def _get_log_level(level: Union[str, int, None]) -> Union[str, int]: """Returns preferred log level set by the customer in upper case""" diff --git a/tests/functional/test_logger.py b/tests/functional/test_logger.py index 9494d39344b..c067c683576 100644 --- a/tests/functional/test_logger.py +++ b/tests/functional/test_logger.py @@ -474,6 +474,18 @@ def test_logger_set_correlation_id_to_none(lambda_context, stdout, service_name) assert "correlation_id" not in log +def test_logger_get_correlation_id(lambda_context, stdout, service_name): + # GIVEN a logger with a correlation_id set + logger = Logger(service=service_name, stream=stdout) + logger.set_correlation_id("foo") + + # WHEN calling get_correlation_id + correlation_id = logger.get_correlation_id() + + # THEN it should return the correlation_id + assert "foo" == correlation_id + + def test_logger_set_correlation_id_path(lambda_context, stdout, service_name): # GIVEN logger = Logger(service=service_name, stream=stdout) From 20614da645e49c0d745a29ad72166fc7611317b8 Mon Sep 17 00:00:00 2001 From: Michael Brewer Date: Mon, 19 Jul 2021 08:26:52 -0700 Subject: [PATCH 3/4] chore: setting correlation is none is enough --- aws_lambda_powertools/logging/logger.py | 5 +---- 1 file changed, 1 insertion(+), 4 deletions(-) diff --git a/aws_lambda_powertools/logging/logger.py b/aws_lambda_powertools/logging/logger.py index bc845a28a5d..35054f86137 100644 --- a/aws_lambda_powertools/logging/logger.py +++ b/aws_lambda_powertools/logging/logger.py @@ -395,10 +395,7 @@ def set_correlation_id(self, value: Optional[str]): value : str, optional Value for the correlation id. None will remove the correlation_id """ - if value is None: - self.remove_keys(["correlation_id"]) - else: - self.append_keys(correlation_id=value) + self.append_keys(correlation_id=value) def get_correlation_id(self) -> Optional[str]: """Gets the correlation_id in the logging json From e620fd195348e611592e568c728e82e510a5532d Mon Sep 17 00:00:00 2001 From: Michael Brewer Date: Mon, 19 Jul 2021 08:29:42 -0700 Subject: [PATCH 4/4] chore: remove duplicate test --- tests/functional/test_logger.py | 14 -------------- 1 file changed, 14 deletions(-) diff --git a/tests/functional/test_logger.py b/tests/functional/test_logger.py index c067c683576..a8d92c05257 100644 --- a/tests/functional/test_logger.py +++ b/tests/functional/test_logger.py @@ -460,20 +460,6 @@ def handler(event, _): assert request_id == log["correlation_id"] -def test_logger_set_correlation_id_to_none(lambda_context, stdout, service_name): - # GIVEN a logger with a correlation_id set - logger = Logger(service=service_name, stream=stdout) - logger.set_correlation_id("foo") - - # WHEN calling set_correlation_id with None - logger.set_correlation_id(None) - - # THEN there should be no correlation_id - logger.info("Foo") - log = capture_logging_output(stdout) - assert "correlation_id" not in log - - def test_logger_get_correlation_id(lambda_context, stdout, service_name): # GIVEN a logger with a correlation_id set logger = Logger(service=service_name, stream=stdout)