refactor: add append/remove keys, common properties

Author: heitorlessa

aws_lambda_powertools/logging/formatter.py

@@ -51,7 +51,7 @@ def __init__(
         json_deserializer: Optional[Callable[[Any], Any]] = None,
         json_default: Optional[Callable[[Any], Any]] = None,
         datefmt: str = None,
-        log_record_order: List = None,
+        log_record_order: List[str] = None,
         utc: bool = False,
         **kwargs
     ):
@@ -99,6 +99,7 @@ def __init__(
         self.utc = utc
         self.log_record_order = log_record_order or ["level", "location", "message", "timestamp"]
         self.log_format = dict.fromkeys(self.log_record_order)  # Set the insertion order for the log messages
+        self.update_formatter = self.append_keys  # alias to previous method
 
         if self.utc:
             self.converter = time.gmtime
@@ -108,6 +109,13 @@ def __init__(
         keys_combined = {**self._build_default_keys(), **kwargs}
         self.log_format.update(**keys_combined)
 
+    def append_keys(self, **additional_keys):
+        self.log_format.update(additional_keys)
+
+    def remove_keys(self, keys: Iterable[str]):
+        for key in keys:
+            self.log_format.pop(key, None)
+
     @staticmethod
     def _build_default_keys():
         return {
@@ -121,9 +129,6 @@ def _get_latest_trace_id():
         xray_trace_id = os.getenv(constants.XRAY_TRACE_ID_ENV)
         return xray_trace_id.split(";")[0].replace("Root=", "") if xray_trace_id else None
 
-    def update_formatter(self, **kwargs):
-        self.log_format.update(kwargs)
-
     def _extract_log_message(self, log_record: logging.LogRecord) -> Union[Dict[str, Any], str, bool, Iterable]:
         """Extract message from log record and attempt to JSON decode it if str
 

aws_lambda_powertools/logging/logger.py

@@ -4,7 +4,7 @@
 import os
 import random
 import sys
-from typing import Any, Callable, Dict, Optional, Union
+from typing import Any, Callable, Dict, Iterable, Optional, Union
 
 import jmespath
 
@@ -42,6 +42,7 @@ def _is_cold_start() -> bool:
 # so we need to return to subclassing removed in #97
 # All methods/properties continue to be proxied to inner logger
 # https://github.com/awslabs/aws-lambda-powertools-python/issues/107
+# noinspection PyRedeclaration
 class Logger(logging.Logger):  # lgtm [py/missing-call-to-init]
     """Creates and setups a logger to format statements in JSON.
 
@@ -96,7 +97,7 @@ class Logger(logging.Logger):  # lgtm [py/missing-call-to-init]
         >>> logger = Logger(service="payment")
         >>>
         >>> def handler(event, context):
-                logger.structure_logs(append=True, payment_id=event["payment_id"])
+                logger.append_keys(payment_id=event["payment_id"])
                 logger.info("Hello")
 
     **Create child Logger using logging inheritance via child param**
@@ -139,6 +140,7 @@ def __init__(
         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()
+
         self._init_logger(**kwargs)
 
     def __getattr__(self, name):
@@ -282,12 +284,28 @@ def decorate(event, context):
 
         return decorate
 
-    def structure_logs(self, append: bool = False, **kwargs):
+    def append_keys(self, **additional_keys):
+        self.registered_formatter.append_keys(**additional_keys)
+
+    def remove_keys(self, keys: Iterable[str]):
+        self.registered_formatter.remove_keys(keys)
+
+    @property
+    def registered_handler(self) -> logging.Handler:
+        """Registered 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"""
+        return self.registered_handler.formatter
+
+    def structure_logs(self, append: bool = False, **keys):
         """Sets logging formatting to JSON.
 
         Optionally, it can append keyword arguments
-        to an existing logger so it is available
-        across future log statements.
+        to an existing logger so it is available across future log statements.
 
         Last keyword argument and value wins if duplicated.
 
@@ -297,15 +315,12 @@ def structure_logs(self, append: bool = False, **kwargs):
             [description], by default False
         """
 
-        # Child loggers don't have handlers attached, use its parent handlers
-        handlers = self._logger.parent.handlers if self.child else self._logger.handlers
-        for handler in handlers:
-            if append:
-                # Update existing formatter in an existing logger handler
-                handler.formatter.update_formatter(**kwargs)
-            else:
-                # Set a new formatter for a logger handler
-                handler.setFormatter(LambdaPowertoolsFormatter(**self._default_log_keys, **kwargs))
+        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))
 
     def set_correlation_id(self, value: str):
         """Sets the correlation_id in the logging json

tests/functional/test_logger.py

@@ -474,3 +474,23 @@ def handler(event, context):
     # THEN
     log = capture_logging_output(stdout)
     assert request_id == log["correlation_id"]
+
+
+def test_logger_append_remove_keys(stdout, service_name):
+    # GIVEN a Logger is initialized
+    logger = Logger(service=service_name, stream=stdout)
+    extra_keys = {"request_id": "id", "context": "value"}
+
+    # WHEN keys are updated
+    logger.append_keys(**extra_keys)
+    logger.info("message with new keys")
+
+    # And removed
+    logger.remove_keys(extra_keys.keys())
+    logger.info("message after keys being removed")
+
+    # THEN additional keys should only be present in the first log statement
+    extra_keys_log, keys_removed_log = capture_multiple_logging_statements_output(stdout)
+
+    assert extra_keys.items() <= extra_keys_log.items()
+    assert (extra_keys.items() <= keys_removed_log.items()) is False