add context manager to logger

Author: anafalcao

aws_lambda_powertools/logging/logger.py

@@ -7,16 +7,8 @@
 import random
 import sys
 import warnings
-from typing import (
-    IO,
-    TYPE_CHECKING,
-    Any,
-    Callable,
-    Iterable,
-    Mapping,
-    TypeVar,
-    overload,
-)
+from contextlib import contextmanager
+from typing import IO, TYPE_CHECKING, Any, Callable, Generator, Iterable, Mapping, TypeVar, overload
 
 from aws_lambda_powertools.logging.constants import (
     LOGGER_ATTRIBUTE_PRECONFIGURED,
@@ -338,6 +330,31 @@ def _configure_sampling(self) -> None:
                 ),
             )
 
+    @contextmanager
+    def append_context_keys(self, **keys: Any) -> Generator[Any, Any, Any]:
+        """
+        Context manager to temporarily add logging keys.
+
+        Parameters:
+        -----------
+        **keys: Any
+            Key-value pairs to include in the log context during the lifespan of the context manager.
+
+        Example:
+        --------
+        >>> logger = Logger(service="example_service")
+        >>> with logger.append_context_keys(user_id="123", operation="process"):
+        >>>     logger.info("Log with context")
+        >>> logger.info("Log without context")
+        """
+        # Add keys to the context
+        self.append_keys(**keys)
+        try:
+            yield
+        finally:
+            # Remove the keys after exiting the context
+            self.remove_keys(keys.keys())
+
     @overload
     def inject_lambda_context(
         self,

docs/core/logger.md

@@ -186,6 +186,15 @@ You can append your own keys to your existing Logger via `append_keys(**addition
 
     This example will add `order_id` if its value is not empty, and in subsequent invocations where `order_id` might not be present it'll remove it from the Logger.
 
+#### append_context_keys method
+
+???+ warning
+	`append_context_keys` is not thread-safe.
+
+The append_context_keys method allows temporary modification of a Logger instance's context without creating a new logger. It's useful for adding context keys to specific workflows while maintaining the logger's overall state and simplicity.
+
+* Add examples
+
 #### ephemeral metadata
 
 You can pass an arbitrary number of keyword arguments (kwargs) to all log level's methods, e.g. `logger.info, logger.warning`.

tests/functional/logger/required_dependencies/test_logger.py

@@ -1114,3 +1114,106 @@ def test_logger_json_unicode(stdout, service_name):
 
     assert log["message"] == non_ascii_chars
     assert log[japanese_field] == japanese_string
+
+
+def test_append_context_keys_adds_and_removes_keys(stdout, service_name):
+    # GIVEN a Logger is initialized
+    logger = Logger(service=service_name, stream=stdout)
+    test_keys = {"user_id": "123", "operation": "test"}
+
+    # WHEN context keys are added
+    with logger.append_context_keys(**test_keys):
+        logger.info("message with context keys")
+    logger.info("message without context keys")
+
+    # THEN context keys should only be present in the first log statement
+    with_context_log, without_context_log = capture_multiple_logging_statements_output(stdout)
+
+    assert test_keys.items() <= with_context_log.items()
+    assert (test_keys.items() <= without_context_log.items()) is False
+
+
+def test_append_context_keys_handles_empty_dict(stdout, service_name):
+    # GIVEN a Logger is initialized
+    logger = Logger(service=service_name, stream=stdout)
+
+    # WHEN context is added with no keys
+    with logger.append_context_keys():
+        logger.info("message with empty context")
+
+    # THEN log should contain only default keys
+    log_output = capture_logging_output(stdout)
+    assert set(log_output.keys()) == {"service", "timestamp", "level", "message", "location"}
+
+
+def test_append_context_keys_handles_exception(stdout, service_name):
+    # GIVEN a Logger is initialized
+    logger = Logger(service=service_name, stream=stdout)
+    test_keys = {"user_id": "123"}
+
+    # WHEN an exception occurs within the context
+    try:
+        with logger.append_context_keys(**test_keys):
+            logger.info("message before exception")
+            raise ValueError("Test exception")
+    except ValueError:
+        logger.info("message after exception")
+
+    # THEN context keys should only be present in the first log statement
+    before_exception, after_exception = capture_multiple_logging_statements_output(stdout)
+
+    assert test_keys.items() <= before_exception.items()
+    assert (test_keys.items() <= after_exception.items()) is False
+
+
+def test_append_context_keys_nested_contexts(stdout, service_name):
+    # GIVEN a Logger is initialized
+    logger = Logger(service=service_name, stream=stdout)
+
+    # WHEN nested contexts are used
+    with logger.append_context_keys(level1="outer"):
+        logger.info("outer context message")
+        with logger.append_context_keys(level2="inner"):
+            logger.info("nested context message")
+        logger.info("back to outer context message")
+    logger.info("no context message")
+
+    # THEN logs should contain appropriate context keys
+    outer, nested, back_outer, no_context = capture_multiple_logging_statements_output(stdout)
+
+    assert outer["level1"] == "outer"
+    assert "level2" not in outer
+
+    assert nested["level1"] == "outer"
+    assert nested["level2"] == "inner"
+
+    assert back_outer["level1"] == "outer"
+    assert "level2" not in back_outer
+
+    assert "level1" not in no_context
+    assert "level2" not in no_context
+
+
+def test_append_context_keys_with_formatter(stdout, service_name):
+    # GIVEN a Logger is initialized with a custom formatter
+    class CustomFormatter(BasePowertoolsFormatter):
+        def append_keys(self, **additional_keys):
+            pass
+
+        def clear_state(self) -> None:
+            pass
+
+        def remove_keys(self, keys: Iterable[str]) -> None:
+            pass
+
+    custom_formatter = CustomFormatter()
+    logger = Logger(service=service_name, stream=stdout, logger_formatter=custom_formatter)
+    test_keys = {"request_id": "id", "context": "value"}
+
+    # WHEN context keys are added
+    with logger.append_context_keys(**test_keys):
+        logger.info("message with context")
+
+    # THEN the context keys should not persist
+    current_keys = logger.get_current_keys()
+    assert current_keys == {}