Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

feat(logger): add clear_state method #5956

Merged
merged 4 commits into from
Jan 31, 2025
Merged
Show file tree
Hide file tree
Changes from 2 commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
11 changes: 10 additions & 1 deletion aws_lambda_powertools/logging/logger.py
Original file line number Diff line number Diff line change
Expand Up @@ -222,6 +222,7 @@ def __init__(
choice=sampling_rate,
env=os.getenv(constants.LOGGER_LOG_SAMPLING_RATE),
)
self._default_log_keys: dict[str, Any] = {"service": self.service, "sampling_rate": self.sampling_rate}
anafalcao marked this conversation as resolved.
Show resolved Hide resolved
self.child = child
self.logger_formatter = logger_formatter
self._stream = stream or sys.stdout
Expand All @@ -231,7 +232,6 @@ def __init__(
self._is_deduplication_disabled = resolve_truthy_env_var_choice(
env=os.getenv(constants.LOGGER_LOG_DEDUPLICATION_ENV, "false"),
)
self._default_log_keys = {"service": self.service, "sampling_rate": self.sampling_rate}
self._logger = self._get_logger()

# NOTE: This is primarily to improve UX, so IDEs can autocomplete LambdaPowertoolsFormatter options
Expand Down Expand Up @@ -605,6 +605,15 @@ def append_context_keys(self, **additional_keys: Any) -> Generator[None, None, N
with self.registered_formatter.append_context_keys(**additional_keys):
yield

def clear_state(self) -> None:
"""Removes all custom keys that were appended to the Logger."""
# Clear all custom keys from the formatter
self.registered_formatter.clear_state()

# Reset to default keys
default_keys: dict[Any, Any] = dict(self._default_log_keys)
self.structure_logs(**default_keys)
anafalcao marked this conversation as resolved.
Show resolved Hide resolved

# 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:
Expand Down
25 changes: 24 additions & 1 deletion docs/core/logger.md
Original file line number Diff line number Diff line change
Expand Up @@ -274,13 +274,15 @@ You can remove any additional key from Logger state using `remove_keys`.

#### Clearing all state

##### Decorator with clear_state

Logger is commonly initialized in the global scope. Due to [Lambda Execution Context reuse](https://docs.aws.amazon.com/lambda/latest/dg/runtimes-context.html){target="_blank"}, this means that custom keys can be persisted across invocations. If you want all custom keys to be deleted, you can use `clear_state=True` param in `inject_lambda_context` decorator.

???+ tip "Tip: When is this useful?"
It is useful when you add multiple custom keys conditionally, instead of setting a default `None` value if not present. Any key with `None` value is automatically removed by Logger.

???+ danger "Danger: This can have unintended side effects if you use Layers"
Lambda Layers code is imported before the Lambda handler.
Lambda Layers code is imported before the Lambda handler. When a Lambda function starts, it first imports and executes all code in the Layers (including any global scope code) before proceeding to the function's own code.
anafalcao marked this conversation as resolved.
Show resolved Hide resolved

This means that `clear_state=True` will instruct Logger to remove any keys previously added before Lambda handler execution proceeds.

Expand All @@ -304,6 +306,27 @@ Logger is commonly initialized in the global scope. Due to [Lambda Execution Con
--8<-- "examples/logger/src/clear_state_event_two.json"
```

##### clear_state method

You can call `clear_state()` as a method explicitly within your code to clear appended keys at any point during the execution of your Lambda invocation.

=== "clear_state_method.py"

```python hl_lines="12"
--8<-- "examples/logger/src/clear_state_method.py"
```
=== "Output before clear_state()"

```json hl_lines="9 17"
--8<-- "examples/logger/src/before_clear_state.json"
```

=== "Output after clear_state()"

```json hl_lines="4"
--8<-- "examples/logger/src/after_clear_state.json"
```

### Accessing currently configured keys

You can view all currently configured keys from the Logger state using the `get_current_keys()` method. This method is useful when you need to avoid overwriting keys that are already configured.
Expand Down
7 changes: 7 additions & 0 deletions examples/logger/src/after_clear_state.json
Original file line number Diff line number Diff line change
@@ -0,0 +1,7 @@
{
"level": "INFO",
"location": "lambda_handler:126",
"message": "State after clearing - only show default keys",
"timestamp": "2025-01-30 13:56:03,158-0300",
"service": "payment"
}
20 changes: 20 additions & 0 deletions examples/logger/src/before_clear_state.json
Original file line number Diff line number Diff line change
@@ -0,0 +1,20 @@
{
"logs": [
{
"level": "INFO",
"location": "lambda_handler:122",
"message": "Starting order processing",
"timestamp": "2025-01-30 13:56:03,157-0300",
"service": "payment",
"order_id": "12345"
},
{
"level": "INFO",
"location": "lambda_handler:124",
"message": "Final state before clearing",
"timestamp": "2025-01-30 13:56:03,157-0300",
"service": "payment",
"order_id": "12345"
}
]
}
15 changes: 15 additions & 0 deletions examples/logger/src/clear_state_method.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,15 @@
from aws_lambda_powertools import Logger
from aws_lambda_powertools.utilities.typing import LambdaContext

logger = Logger(service="payment", level="DEBUG")


def lambda_handler(event: dict, context: LambdaContext) -> str:
try:
logger.append_keys(order_id="12345")
logger.info("Starting order processing")
finally:
logger.info("Final state before clearing")
logger.clear_state()
logger.info("State after clearing - only show default keys")
return "Completed"
54 changes: 54 additions & 0 deletions tests/functional/logger/required_dependencies/test_logger.py
Original file line number Diff line number Diff line change
Expand Up @@ -1232,3 +1232,57 @@ def test_logger_change_level_child_logger(stdout, service_name):
logs = list(stdout.getvalue().strip().split("\n"))
assert len(logs) == 1
assert "service" in logs[0]


def test_clear_state_with_append_keys():
anafalcao marked this conversation as resolved.
Show resolved Hide resolved
# GIVEN a Logger is initialized
logger = Logger(service="service_name", stream=stdout)

# WHEN append keys are added
logger.append_keys(custom_key="custom_key")
logger.info("message with appended keys")
logger.clear_state()

# THEN context keys should be cleared
assert "custom_key" not in logger.get_current_keys()


def test_clear_state(stdout, service_name):
anafalcao marked this conversation as resolved.
Show resolved Hide resolved
# GIVEN a Logger is initialized
logger = Logger(service=service_name, stream=stdout)
logger.info("message for the user")

# WHEN the clear_state method is called
logger.clear_state()

# THEN the logger's current keys should be reset to their default values
expected_keys = {
"level": "%(levelname)s",
"location": "%(funcName)s:%(lineno)d",
"message": None,
"timestamp": "%(asctime)s",
"service": service_name,
"sampling_rate": None,
}
assert logger.get_current_keys() == expected_keys


def test_clear_state_log_output(stdout, service_name):
# GIVEN a Logger is initialized
logger = Logger(service=service_name, stream=stdout)

# WHEN we append a custom key and log
logger.append_keys(custom_key="test_value")
logger.info("first message")

# AND we clear the state and log again
logger.clear_state()
logger.info("second message")

# THEN the first log should contain the custom key
# AND the second log should not contain the custom key
first_log, second_log = capture_multiple_logging_statements_output(stdout)

assert "custom_key" in first_log
assert first_log["custom_key"] == "test_value"
assert "custom_key" not in second_log
Loading