diff --git a/aws_lambda_powertools/logging/formatter.py b/aws_lambda_powertools/logging/formatter.py index 0a497e8244a..824c5c0ef16 100644 --- a/aws_lambda_powertools/logging/formatter.py +++ b/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: diff --git a/aws_lambda_powertools/logging/logger.py b/aws_lambda_powertools/logging/logger.py index acefe9757c6..c242f5c9bd4 100644 --- a/aws_lambda_powertools/logging/logger.py +++ b/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: diff --git a/docs/core/logger.md b/docs/core/logger.md index 818d5a6589b..9915f7cc4b4 100644 --- a/docs/core/logger.md +++ b/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`. diff --git a/examples/logger/src/append_context_keys.json b/examples/logger/src/append_context_keys.json new file mode 100644 index 00000000000..97770a657fa --- /dev/null +++ b/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" + } +] \ No newline at end of file diff --git a/examples/logger/src/append_context_keys.py b/examples/logger/src/append_context_keys.py new file mode 100644 index 00000000000..704735eeb9a --- /dev/null +++ b/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" diff --git a/tests/functional/logger/required_dependencies/test_logger.py b/tests/functional/logger/required_dependencies/test_logger.py index e86dba27eb6..70f08f1bbdd 100644 --- a/tests/functional/logger/required_dependencies/test_logger.py +++ b/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 == {}