feat(logger): add context manager for logger keys (#5883)

* add context manager to logger

* Passing the method implementation to the formatter class

* modify logger tests

* add examples to doc

---------

Co-authored-by: Leandro Damascena <lcdama@amazon.pt>

aws_lambda_powertools/logging/formatter.py

@@ -7,10 +7,11 @@
 import time
 import traceback
 from abc import ABCMeta, abstractmethod
+from contextlib import contextmanager
 from contextvars import ContextVar
 from datetime import datetime, timezone
 from functools import partial
-from typing import TYPE_CHECKING, Any, Callable, Iterable
+from typing import TYPE_CHECKING, Any, Callable, Generator, Iterable
 
 from aws_lambda_powertools.shared import constants
 from aws_lambda_powertools.shared.functions import powertools_dev_is_set
@@ -62,6 +63,10 @@ def clear_state(self) -> None:
         """Removes any previously added logging keys"""
         raise NotImplementedError()
 
+    @contextmanager
+    def append_context_keys(self, **additional_keys: Any) -> Generator[None, None, None]:
+        yield
+
     # These specific thread-safe methods are necessary to manage shared context in concurrent environments.
     # They prevent race conditions and ensure data consistency across multiple threads.
     def thread_safe_append_keys(self, **additional_keys) -> None:
@@ -263,6 +268,31 @@ def clear_state(self) -> None:
         self.log_format = dict.fromkeys(self.log_record_order)
         self.log_format.update(**self.keys_combined)
 
+    @contextmanager
+    def append_context_keys(self, **additional_keys: Any) -> Generator[None, None, None]:
+        """
+        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(**additional_keys)
+        try:
+            yield
+        finally:
+            # Remove the keys after exiting the context
+            self.remove_keys(additional_keys.keys())
+
     # These specific thread-safe methods are necessary to manage shared context in concurrent environments.
     # They prevent race conditions and ensure data consistency across multiple threads.
     def thread_safe_append_keys(self, **additional_keys) -> None:

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,
@@ -589,6 +581,26 @@ def get_current_keys(self) -> dict[str, Any]:
     def remove_keys(self, keys: Iterable[str]) -> None:
         self.registered_formatter.remove_keys(keys)
 
+    @contextmanager
+    def append_context_keys(self, **additional_keys: Any) -> Generator[None, None, None]:
+        """
+        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")
+        """
+        with self.registered_formatter.append_context_keys(**additional_keys):
+            yield
+
     # These specific thread-safe methods are necessary to manage shared context in concurrent environments.
     # They prevent race conditions and ensure data consistency across multiple threads.
     def thread_safe_append_keys(self, **additional_keys: object) -> None:

docs/core/logger.md

@@ -187,6 +187,25 @@ 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.
+
+=== "append_context_keys.py"
+
+    ```python hl_lines="7 8"
+    --8<-- "examples/logger/src/append_context_keys.py"
+    ```
+
+=== "append_context_keys_output.json"
+
+    ```json hl_lines="8 9"
+    --8<-- "examples/logger/src/append_context_keys.json"
+    ```
+
 #### ephemeral metadata
 
 You can pass an arbitrary number of keyword arguments (kwargs) to all log level's methods, e.g. `logger.info, logger.warning`.

examples/logger/src/append_context_keys.json

@@ -0,0 +1,18 @@
+[
+    {
+        "level": "INFO",
+        "location": "lambda_handler:8",
+        "message": "Log with context",
+        "timestamp": "2024-03-21T10:30:00.123Z",
+        "service": "example_service",
+        "user_id": "123",
+        "operation": "process"
+    },
+    {
+        "level": "INFO",
+        "location": "lambda_handler:10",
+        "message": "Log without context",
+        "timestamp": "2024-03-21T10:30:00.124Z",
+        "service": "example_service"
+    }
+]

examples/logger/src/append_context_keys.py

@@ -0,0 +1,13 @@
+from aws_lambda_powertools import Logger
+from aws_lambda_powertools.utilities.typing import LambdaContext
+
+logger = Logger(service="example_service")
+
+
+def lambda_handler(event: dict, context: LambdaContext) -> str:
+    with logger.append_context_keys(user_id="123", operation="process"):
+        logger.info("Log with context")
+
+    logger.info("Log without context")
+
+    return "hello world"

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 "user_id" in with_context_log
+    assert test_keys["user_id"] == with_context_log["user_id"]
+    assert "user_id" not in without_context_log
+
+
+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_user_id = "128"
+
+    # WHEN an exception occurs within the context
+    exception_raised = False
+    try:
+        with logger.append_context_keys(user_id=test_user_id):
+            logger.info("message before exception")
+            raise ValueError("Test exception")
+    except ValueError:
+        exception_raised = True
+        logger.info("message after exception")
+
+    # THEN verify the exception was raised and handled
+    assert exception_raised, "Expected ValueError to be raised"
+
+
+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 == {}