diff --git a/aws_lambda_powertools/logging/formatter.py b/aws_lambda_powertools/logging/formatter.py index 90799b84ed1..becfc9de85c 100644 --- a/aws_lambda_powertools/logging/formatter.py +++ b/aws_lambda_powertools/logging/formatter.py @@ -127,8 +127,8 @@ def __init__( super(LambdaPowertoolsFormatter, self).__init__(datefmt=self.datefmt) - keys_combined = {**self._build_default_keys(), **kwargs} - self.log_format.update(**keys_combined) + self.keys_combined = {**self._build_default_keys(), **kwargs} + self.log_format.update(**self.keys_combined) def serialize(self, log: Dict) -> str: """Serialize structured log dict to JSON str""" @@ -187,7 +187,7 @@ def remove_keys(self, keys: Iterable[str]): def clear_state(self): self.log_format = dict.fromkeys(self.log_record_order) - self.log_format.update(**self._build_default_keys()) + self.log_format.update(**self.keys_combined) @staticmethod def _build_default_keys(): diff --git a/tests/functional/test_logger.py b/tests/functional/test_logger.py index de9de42601f..f4a814e4913 100644 --- a/tests/functional/test_logger.py +++ b/tests/functional/test_logger.py @@ -683,6 +683,26 @@ def handler(event, context): assert key in second_log +def test_clear_state_keeps_custom_keys(lambda_context, stdout, service_name): + # GIVEN + location_format = "%(module)s.%(funcName)s:clear_state" + logger = Logger(service=service_name, stream=stdout, location=location_format, custom_key="foo") + + # WHEN clear_state is set + @logger.inject_lambda_context(clear_state=True) + def handler(event, context): + logger.info("Foo") + + # THEN all standard keys should be available as usual + handler({}, lambda_context) + handler({}, lambda_context) + + first_log, second_log = capture_multiple_logging_statements_output(stdout) + for log in (first_log, second_log): + assert "foo" == log["custom_key"] + assert "test_logger.handler:clear_state" == log["location"] + + def test_clear_state_keeps_exception_keys(lambda_context, stdout, service_name): # GIVEN logger = Logger(service=service_name, stream=stdout)