feat: bring your own formatter

Author: heitorlessa

aws_lambda_powertools/logging/formatter.py

@@ -2,6 +2,7 @@
 import logging
 import os
 import time
+from abc import ABCMeta, abstractmethod
 from functools import partial
 from typing import Any, Callable, Dict, Iterable, List, Optional, Tuple, Union
 
@@ -35,7 +36,17 @@
 )
 
 
-class LambdaPowertoolsFormatter(logging.Formatter):
+class BasePowertoolsFormatter(logging.Formatter, metaclass=ABCMeta):
+    @abstractmethod
+    def append_keys(self, **additional_keys):
+        raise NotImplementedError()
+
+    @abstractmethod
+    def remove_keys(self, keys: Iterable[str]):
+        raise NotImplementedError()
+
+
+class LambdaPowertoolsFormatter(BasePowertoolsFormatter):
     """AWS Lambda Powertools Logging formatter.
 
     Formats the log message as a JSON encoded string. If the message is a

aws_lambda_powertools/logging/logger.py

@@ -12,7 +12,7 @@
 from ..shared.functions import resolve_env_var_choice, resolve_truthy_env_var_choice
 from .exceptions import InvalidLoggerSamplingRateError
 from .filters import SuppressFilter
-from .formatter import LambdaPowertoolsFormatter
+from .formatter import BasePowertoolsFormatter, LambdaPowertoolsFormatter
 from .lambda_context import build_lambda_context_model
 
 logger = logging.getLogger(__name__)
@@ -124,6 +124,7 @@ def __init__(
         child: bool = False,
         sampling_rate: float = None,
         stream: sys.stdout = None,
+        logger_formatter: Optional[BasePowertoolsFormatter] = None,
         **kwargs,
     ):
         self.service = resolve_env_var_choice(
@@ -132,11 +133,12 @@ def __init__(
         self.sampling_rate = resolve_env_var_choice(
             choice=sampling_rate, env=os.getenv(constants.LOGGER_LOG_SAMPLING_RATE)
         )
+        self.child = child
+        self.logger_formatter = logger_formatter
+        self.log_level = self._get_log_level(level)
         self._is_deduplication_disabled = resolve_truthy_env_var_choice(
             env=os.getenv(constants.LOGGER_LOG_DEDUPLICATION_ENV, "false")
         )
-        self.log_level = self._get_log_level(level)
-        self.child = child
         self._handler = logging.StreamHandler(stream) or logging.StreamHandler(sys.stdout)
         self._default_log_keys = {"service": self.service, "sampling_rate": self.sampling_rate}
         self._logger = self._get_logger()
@@ -292,13 +294,13 @@ def remove_keys(self, keys: Iterable[str]):
 
     @property
     def registered_handler(self) -> logging.Handler:
-        """Registered Logger handler"""
+        """Convenience property to access logger handler"""
         handlers = self._logger.parent.handlers if self.child else self._logger.handlers
         return handlers[0]
 
     @property
-    def registered_formatter(self) -> Optional[LambdaPowertoolsFormatter]:
-        """Registered Logger formatter"""
+    def registered_formatter(self) -> Optional[BasePowertoolsFormatter]:
+        """Convenience property to access logger formatter"""
         return self.registered_handler.formatter
 
     def structure_logs(self, append: bool = False, **keys):
@@ -312,15 +314,16 @@ def structure_logs(self, append: bool = False, **keys):
         Parameters
         ----------
         append : bool, optional
-            [description], by default False
+            append keys provided to logger formatter, by default False
         """
 
         if append:
             # Maintenance: Add deprecation warning for major version. Refer to append_keys() when docs are updated
             self.append_keys(**keys)
         else:
-            # Set a new formatter for a logger handler
-            self.registered_handler.setFormatter(LambdaPowertoolsFormatter(**self._default_log_keys, **keys))
+            log_keys = {**self._default_log_keys, **keys}
+            formatter = self.logger_formatter or LambdaPowertoolsFormatter(**log_keys)
+            self.registered_handler.setFormatter(formatter)
 
     def set_correlation_id(self, value: str):
         """Sets the correlation_id in the logging json

tests/functional/test_logger.py

@@ -5,12 +5,14 @@
 import random
 import string
 from collections import namedtuple
+from typing import Iterable
 
 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.logger import set_package_logger
 from aws_lambda_powertools.shared import constants
 
@@ -494,3 +496,51 @@ def test_logger_append_remove_keys(stdout, service_name):
 
     assert extra_keys.items() <= extra_keys_log.items()
     assert (extra_keys.items() <= keys_removed_log.items()) is False
+
+
+def test_logger_custom_formatter(stdout, service_name, lambda_context):
+    class CustomFormatter(BasePowertoolsFormatter):
+        custom_format = {}
+
+        def append_keys(self, **additional_keys):
+            self.custom_format.update(additional_keys)
+
+        def remove_keys(self, keys: Iterable[str]):
+            for key in keys:
+                self.custom_format.pop(key, None)
+
+        def format(self, record: logging.LogRecord) -> str:  # noqa: A003
+            return json.dumps(
+                {
+                    "message": super().format(record),
+                    "timestamp": self.formatTime(record),
+                    "my_default_key": "test",
+                    **self.custom_format,
+                }
+            )
+
+    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
+    @logger.inject_lambda_context
+    def handler(event, context):
+        logger.info("Hello")
+
+    handler({}, lambda_context)
+
+    lambda_context_keys = (
+        "function_name",
+        "function_memory_size",
+        "function_arn",
+        "function_request_id",
+    )
+
+    log = capture_logging_output(stdout)
+
+    # THEN custom key should always be present
+    # and lambda contextual info should also be in the logs
+    assert "my_default_key" in log
+    assert all(k in log for k in lambda_context_keys)