diff --git a/aws_lambda_powertools/logging/formatter.py b/aws_lambda_powertools/logging/formatter.py index a15452da9fd..1b42df249ae 100644 --- a/aws_lambda_powertools/logging/formatter.py +++ b/aws_lambda_powertools/logging/formatter.py @@ -45,6 +45,11 @@ def append_keys(self, **additional_keys): def remove_keys(self, keys: Iterable[str]): raise NotImplementedError() + @abstractmethod + def clear_state(self): + """Removes any previously added logging keys""" + raise NotImplementedError() + class LambdaPowertoolsFormatter(BasePowertoolsFormatter): """AWS Lambda Powertools Logging formatter. @@ -180,6 +185,9 @@ def remove_keys(self, keys: Iterable[str]): for key in keys: self.log_format.pop(key, None) + def clear_state(self): + self.log_format = dict.fromkeys(self.log_record_order) + @staticmethod def _build_default_keys(): return { diff --git a/aws_lambda_powertools/logging/logger.py b/aws_lambda_powertools/logging/logger.py index 49321181b48..05fd6c98e04 100644 --- a/aws_lambda_powertools/logging/logger.py +++ b/aws_lambda_powertools/logging/logger.py @@ -385,14 +385,26 @@ def structure_logs(self, append: bool = False, **keys): append : bool, optional append keys provided to logger formatter, by default False """ + # There are 3 operational modes for this method + ## 1. Register a Powertools Formatter for the first time + ## 2. Append new keys to the current logger formatter; deprecated in favour of append_keys + ## 3. Add new keys and discard existing to the registered formatter - if append: - # Maintenance: Add deprecation warning for major version. Refer to append_keys() when docs are updated - self.append_keys(**keys) - else: - log_keys = {**self._default_log_keys, **keys} + # Mode 1 + log_keys = {**self._default_log_keys, **keys} + is_logger_preconfigured = getattr(self._logger, "init", False) + if not is_logger_preconfigured: formatter = self.logger_formatter or LambdaPowertoolsFormatter(**log_keys) # type: ignore - self.registered_handler.setFormatter(formatter) + return self.registered_handler.setFormatter(formatter) + + # Mode 2 (legacy) + if append: + # Maintenance: Add deprecation warning for major version + return self.append_keys(**keys) + + # Mode 3 + self.registered_formatter.clear_state() + self.registered_formatter.append_keys(**log_keys) def set_correlation_id(self, value: Optional[str]): """Sets the correlation_id in the logging json diff --git a/docs/core/logger.md b/docs/core/logger.md index 31f21147659..0edc4aa3ba7 100644 --- a/docs/core/logger.md +++ b/docs/core/logger.md @@ -862,7 +862,10 @@ logger.info("Collecting payment") By default, Logger uses [LambdaPowertoolsFormatter](#lambdapowertoolsformatter) that persists its custom structure between non-cold start invocations. There could be scenarios where the existing feature set isn't sufficient to your formatting needs. -For **minor changes like remapping keys** after all log record processing has completed, you can override `serialize` method from [LambdaPowertoolsFormatter](#lambdapowertoolsformatter): +???+ info + The most common use cases are remapping keys by bringing your existing schema, and redacting sensitive information you know upfront. + +For these, you can override the `serialize` method from [LambdaPowertoolsFormatter](#lambdapowertoolsformatter). === "custom_formatter.py" @@ -892,28 +895,39 @@ For **minor changes like remapping keys** after all log record processing has co } ``` -For **replacing the formatter entirely**, you can subclass `BasePowertoolsFormatter`, implement `append_keys` method, and override `format` standard logging method. This ensures the current feature set of Logger like [injecting Lambda context](#capturing-lambda-context-info) and [sampling](#sampling-debug-logs) will continue to work. +The `log` argument is the final log record containing [our standard keys](#standard-structured-keys), optionally [Lambda context keys](#capturing-lambda-context-info), and any custom key you might have added via [append_keys](#append_keys-method) or the [extra parameter](#extra-parameter). + +For exceptional cases where you want to completely replace our formatter logic, you can subclass `BasePowertoolsFormatter`. + +???+ warning + You will need to implement `append_keys`, `clear_state`, override `format`, and optionally `remove_keys` to keep the same feature set Powertools Logger provides. This also means keeping state of logging keys added. -???+ info - You might need to implement `remove_keys` method if you make use of the feature too. === "collect.py" - ```python hl_lines="2 4 7 12 16 27" + ```python hl_lines="5 7 9-10 13 17 21 24 35" + import logging + from typing import Iterable, List, Optional + from aws_lambda_powertools import Logger from aws_lambda_powertools.logging.formatter import BasePowertoolsFormatter class CustomFormatter(BasePowertoolsFormatter): - custom_format = {} # arbitrary dict to hold our structured keys + def __init__(self, log_record_order: Optional[List[str]], *args, **kwargs): + self.log_record_order = log_record_order or ["level", "location", "message", "timestamp"] + self.log_format = dict.fromkeys(self.log_record_order) + super().__init__(*args, **kwargs) def append_keys(self, **additional_keys): # also used by `inject_lambda_context` decorator - self.custom_format.update(additional_keys) + self.log_format.update(additional_keys) - # Optional unless you make use of this Logger feature def remove_keys(self, keys: Iterable[str]): for key in keys: - self.custom_format.pop(key, None) + self.log_format.pop(key, None) + + def clear_state(self): + self.log_format = dict.fromkeys(self.log_record_order) def format(self, record: logging.LogRecord) -> str: # noqa: A003 """Format logging record as structured JSON str""" @@ -922,7 +936,7 @@ For **replacing the formatter entirely**, you can subclass `BasePowertoolsFormat "event": super().format(record), "timestamp": self.formatTime(record), "my_default_key": "test", - **self.custom_format, + **self.log_format, } ) diff --git a/tests/functional/test_logger.py b/tests/functional/test_logger.py index 20b0a74fc64..9bbc4452d98 100644 --- a/tests/functional/test_logger.py +++ b/tests/functional/test_logger.py @@ -5,16 +5,17 @@ import random import re import string +from ast import Dict from collections import namedtuple from datetime import datetime, timezone -from typing import Iterable +from typing import Any, Callable, Iterable, List, Optional, Union import pytest from aws_lambda_powertools import Logger, Tracer from aws_lambda_powertools.logging import correlation_paths from aws_lambda_powertools.logging.exceptions import InvalidLoggerSamplingRateError -from aws_lambda_powertools.logging.formatter import BasePowertoolsFormatter +from aws_lambda_powertools.logging.formatter import BasePowertoolsFormatter, LambdaPowertoolsFormatter from aws_lambda_powertools.logging.logger import set_package_logger from aws_lambda_powertools.shared import constants from aws_lambda_powertools.utilities.data_classes import S3Event, event_source @@ -524,6 +525,9 @@ def remove_keys(self, keys: Iterable[str]): for key in keys: self.custom_format.pop(key, None) + def clear_state(self): + self.custom_format.clear() + def format(self, record: logging.LogRecord) -> str: # noqa: A003 return json.dumps( { @@ -564,6 +568,63 @@ def handler(event, context): assert logger.get_correlation_id() is None +def test_logger_custom_powertools_formatter_clear_state(stdout, service_name, lambda_context): + class CustomFormatter(LambdaPowertoolsFormatter): + def __init__( + self, + json_serializer: Optional[Callable[[Dict], str]] = None, + json_deserializer: Optional[Callable[[Union[Dict, str, bool, int, float]], str]] = None, + json_default: Optional[Callable[[Any], Any]] = None, + datefmt: Optional[str] = None, + use_datetime_directive: bool = False, + log_record_order: Optional[List[str]] = None, + utc: bool = False, + **kwargs, + ): + super().__init__( + json_serializer, + json_deserializer, + json_default, + datefmt, + use_datetime_directive, + log_record_order, + utc, + **kwargs, + ) + + custom_formatter = CustomFormatter() + + # GIVEN a Logger is initialized with a custom formatter + logger = Logger(service=service_name, stream=stdout, logger_formatter=custom_formatter) + + # WHEN a lambda function is decorated with logger + # and state is to be cleared in the next invocation + @logger.inject_lambda_context(clear_state=True) + def handler(event, context): + if event.get("add_key"): + logger.append_keys(my_key="value") + logger.info("Hello") + + handler({"add_key": True}, lambda_context) + handler({}, lambda_context) + + lambda_context_keys = ( + "function_name", + "function_memory_size", + "function_arn", + "function_request_id", + ) + + first_log, second_log = capture_multiple_logging_statements_output(stdout) + + # THEN my_key should only present once + # and lambda contextual info should also be in both logs + assert "my_key" in first_log + assert "my_key" not in second_log + assert all(k in first_log for k in lambda_context_keys) + assert all(k in second_log for k in lambda_context_keys) + + def test_logger_custom_handler(lambda_context, service_name, tmp_path): # GIVEN a Logger is initialized with a FileHandler log_file = tmp_path / "log.json"