Skip to content

feat(logger): add clear_state method #5956

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

Merged
merged 4 commits into from
Jan 31, 2025
Merged
Show file tree
Hide file tree
Changes from 1 commit
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
10 changes: 10 additions & 0 deletions 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}
self.child = child
self.logger_formatter = logger_formatter
self._stream = stream or sys.stdout
Expand Down Expand Up @@ -605,6 +606,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)

# 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
27 changes: 26 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.

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,29 @@ 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

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

This allows for more granular control over the logger's state within a single function execution, enabling you to reset the logger to its initial state before specific logging operations or at the end of certain processes within the same Lambda run.

=== "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"
29 changes: 29 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,32 @@ 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():
# 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):
logger = Logger(service=service_name, stream=stdout)
logger.info("message for the user")
logger.clear_state()

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
Loading