From 8071f5decc7d5a0ef21097df72c95ed6d870f5aa Mon Sep 17 00:00:00 2001 From: Ana Falcao Date: Fri, 17 Jan 2025 17:40:42 -0300 Subject: [PATCH 1/4] add context manager to logger --- aws_lambda_powertools/logging/logger.py | 37 +++++-- docs/core/logger.md | 9 ++ .../required_dependencies/test_logger.py | 103 ++++++++++++++++++ 3 files changed, 139 insertions(+), 10 deletions(-) diff --git a/aws_lambda_powertools/logging/logger.py b/aws_lambda_powertools/logging/logger.py index 75a14c6ea2b..ea8ff5aded7 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, @@ -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, diff --git a/docs/core/logger.md b/docs/core/logger.md index 2a45ff08280..0b809133dd2 100644 --- a/docs/core/logger.md +++ b/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`. diff --git a/tests/functional/logger/required_dependencies/test_logger.py b/tests/functional/logger/required_dependencies/test_logger.py index e86dba27eb6..75a8e096a3a 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 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 == {} From edcbb9920e015c62bf70a441146dca9c3d1e19e0 Mon Sep 17 00:00:00 2001 From: Leandro Damascena Date: Mon, 20 Jan 2025 18:28:01 +0000 Subject: [PATCH 2/4] Passing the method implementation to the formatter class --- aws_lambda_powertools/logging/formatter.py | 32 ++++++++++++++- aws_lambda_powertools/logging/logger.py | 45 ++++++++++------------ 2 files changed, 51 insertions(+), 26 deletions(-) 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 3b6c81a5a6e..c242f5c9bd4 100644 --- a/aws_lambda_powertools/logging/logger.py +++ b/aws_lambda_powertools/logging/logger.py @@ -330,31 +330,6 @@ 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, @@ -606,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: From a83fcc5654b0a20d3152219ee953e11abf0d41c0 Mon Sep 17 00:00:00 2001 From: Ana Falcao Date: Tue, 21 Jan 2025 10:26:52 -0300 Subject: [PATCH 3/4] modify logger tests --- .../required_dependencies/test_logger.py | 18 +++++++++--------- 1 file changed, 9 insertions(+), 9 deletions(-) diff --git a/tests/functional/logger/required_dependencies/test_logger.py b/tests/functional/logger/required_dependencies/test_logger.py index 75a8e096a3a..70f08f1bbdd 100644 --- a/tests/functional/logger/required_dependencies/test_logger.py +++ b/tests/functional/logger/required_dependencies/test_logger.py @@ -1129,8 +1129,9 @@ def test_append_context_keys_adds_and_removes_keys(stdout, service_name): # 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 + 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): @@ -1149,21 +1150,20 @@ def test_append_context_keys_handles_empty_dict(stdout, service_name): 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"} + test_user_id = "128" # WHEN an exception occurs within the context + exception_raised = False try: - with logger.append_context_keys(**test_keys): + 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 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 + # 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): From 66e16f05e2271a2310fbdaf1e3ee0305dcdd466d Mon Sep 17 00:00:00 2001 From: Ana Falcao Date: Tue, 21 Jan 2025 10:42:17 -0300 Subject: [PATCH 4/4] add examples to doc --- docs/core/logger.md | 12 +++++++++++- examples/logger/src/append_context_keys.json | 18 ++++++++++++++++++ examples/logger/src/append_context_keys.py | 13 +++++++++++++ 3 files changed, 42 insertions(+), 1 deletion(-) create mode 100644 examples/logger/src/append_context_keys.json create mode 100644 examples/logger/src/append_context_keys.py diff --git a/docs/core/logger.md b/docs/core/logger.md index 9a7aaca8180..9915f7cc4b4 100644 --- a/docs/core/logger.md +++ b/docs/core/logger.md @@ -194,7 +194,17 @@ You can append your own keys to your existing Logger via `append_keys(**addition 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 +=== "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 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"