Skip to content

Commit

Permalink
fix(logger): preserve std keys when using custom formatters (#1264)
Browse files Browse the repository at this point in the history
  • Loading branch information
heitorlessa authored Jul 4, 2022
1 parent 928f2bb commit 6c1a4d2
Show file tree
Hide file tree
Showing 3 changed files with 51 additions and 4 deletions.
5 changes: 5 additions & 0 deletions aws_lambda_powertools/logging/formatter.py
Original file line number Diff line number Diff line change
@@ -1,3 +1,4 @@
import inspect
import json
import logging
import os
Expand Down Expand Up @@ -286,3 +287,7 @@ def _strip_none_records(records: Dict[str, Any]) -> Dict[str, Any]:


JsonFormatter = LambdaPowertoolsFormatter # alias to previous formatter


# Fetch current and future parameters from PowertoolsFormatter that should be reserved
RESERVED_FORMATTER_CUSTOM_KEYS: List[str] = inspect.getfullargspec(LambdaPowertoolsFormatter).args[1:]
16 changes: 12 additions & 4 deletions aws_lambda_powertools/logging/logger.py
Original file line number Diff line number Diff line change
Expand Up @@ -12,7 +12,7 @@
from ..shared.functions import resolve_env_var_choice, resolve_truthy_env_var_choice
from .exceptions import InvalidLoggerSamplingRateError
from .filters import SuppressFilter
from .formatter import BasePowertoolsFormatter, LambdaPowertoolsFormatter
from .formatter import RESERVED_FORMATTER_CUSTOM_KEYS, BasePowertoolsFormatter, LambdaPowertoolsFormatter
from .lambda_context import build_lambda_context_model

logger = logging.getLogger(__name__)
Expand Down Expand Up @@ -82,7 +82,7 @@ class Logger(logging.Logger): # lgtm [py/missing-call-to-init]
datefmt: str, optional
String directives (strftime) to format log timestamp using `time`, by default it uses RFC
3339.
use_datetime_directive: str, optional
use_datetime_directive: bool, optional
Interpret `datefmt` as a format string for `datetime.datetime.strftime`, rather than
`time.strftime`.
Expand Down Expand Up @@ -368,7 +368,7 @@ def registered_handler(self) -> logging.Handler:
return handlers[0]

@property
def registered_formatter(self) -> PowertoolsFormatter:
def registered_formatter(self) -> BasePowertoolsFormatter:
"""Convenience property to access logger formatter"""
return self.registered_handler.formatter # type: ignore

Expand All @@ -395,7 +395,15 @@ def structure_logs(self, append: bool = False, **keys):
is_logger_preconfigured = getattr(self._logger, "init", False)
if not is_logger_preconfigured:
formatter = self.logger_formatter or LambdaPowertoolsFormatter(**log_keys) # type: ignore
return self.registered_handler.setFormatter(formatter)
self.registered_handler.setFormatter(formatter)

# when using a custom Lambda Powertools Formatter
# standard and custom keys that are not Powertools Formatter parameters should be appended
# and custom keys that might happen to be Powertools Formatter parameters should be discarded
# this prevents adding them as custom keys, for example, `json_default=<callable>`
# see https://github.com/awslabs/aws-lambda-powertools-python/issues/1263
custom_keys = {k: v for k, v in log_keys.items() if k not in RESERVED_FORMATTER_CUSTOM_KEYS}
return self.registered_formatter.append_keys(**custom_keys)

# Mode 2 (legacy)
if append:
Expand Down
34 changes: 34 additions & 0 deletions tests/functional/test_logger.py
Original file line number Diff line number Diff line change
Expand Up @@ -625,6 +625,40 @@ def handler(event, context):
assert all(k in second_log for k in lambda_context_keys)


def test_logger_custom_formatter_has_standard_and_custom_keys(stdout, service_name, lambda_context):
class CustomFormatter(LambdaPowertoolsFormatter):
...

# GIVEN a Logger is initialized with a custom formatter
logger = Logger(service=service_name, stream=stdout, logger_formatter=CustomFormatter(), my_key="value")

# WHEN a lambda function is decorated with logger
@logger.inject_lambda_context
def handler(event, context):
logger.info("Hello")

handler({}, lambda_context)

standard_keys = (
"level",
"location",
"message",
"timestamp",
"service",
"cold_start",
"function_name",
"function_memory_size",
"function_arn",
"function_request_id",
)

log = capture_logging_output(stdout)

# THEN all standard keys should be available
assert all(k in log for k in standard_keys)
assert "my_key" in log


def test_logger_custom_handler(lambda_context, service_name, tmp_path):
# GIVEN a Logger is initialized with a FileHandler
log_file = tmp_path / "log.json"
Expand Down

0 comments on commit 6c1a4d2

Please sign in to comment.