fix(logger): correctly pick powertools or custom handler in custom environments (#6083)

* Fixing Powertools logger handler

* Fixing Powertools logger handler

* Make code check tool happy

Author: leandrodamascena

Diff for: aws_lambda_powertools/logging/exceptions.py

@@ -4,3 +4,11 @@ class InvalidLoggerSamplingRateError(Exception):
     """
 
     pass
+
+
+class OrphanedChildLoggerError(Exception):
+    """
+    Orphaned Child logger exception
+    """
+
+    pass

Diff for: aws_lambda_powertools/logging/logger.py

@@ -14,12 +14,14 @@
 import sys
 import warnings
 from contextlib import contextmanager
-from typing import IO, TYPE_CHECKING, Any, Callable, Generator, Iterable, Mapping, TypeVar, overload
+from typing import IO, TYPE_CHECKING, Any, Callable, Generator, Iterable, Mapping, TypeVar, cast, overload
 
 from aws_lambda_powertools.logging.constants import (
+    LOGGER_ATTRIBUTE_HANDLER,
+    LOGGER_ATTRIBUTE_POWERTOOLS_HANDLER,
     LOGGER_ATTRIBUTE_PRECONFIGURED,
 )
-from aws_lambda_powertools.logging.exceptions import InvalidLoggerSamplingRateError
+from aws_lambda_powertools.logging.exceptions import InvalidLoggerSamplingRateError, OrphanedChildLoggerError
 from aws_lambda_powertools.logging.filters import SuppressFilter
 from aws_lambda_powertools.logging.formatter import (
     RESERVED_FORMATTER_CUSTOM_KEYS,
@@ -230,13 +232,14 @@ def __init__(
         self.child = child
         self.logger_formatter = logger_formatter
         self._stream = stream or sys.stdout
-        self.logger_handler = logger_handler or logging.StreamHandler(self._stream)
+
         self.log_uncaught_exceptions = log_uncaught_exceptions
 
         self._is_deduplication_disabled = resolve_truthy_env_var_choice(
             env=os.getenv(constants.LOGGER_LOG_DEDUPLICATION_ENV, "false"),
         )
         self._logger = self._get_logger()
+        self.logger_handler = logger_handler or self._get_handler()
 
         # NOTE: This is primarily to improve UX, so IDEs can autocomplete LambdaPowertoolsFormatter options
         # previously, we masked all of them as kwargs thus limiting feature discovery
@@ -275,6 +278,23 @@ def _get_logger(self) -> logging.Logger:
 
         return logging.getLogger(logger_name)
 
+    def _get_handler(self) -> logging.Handler:
+        # is a logger handler already configured?
+        if getattr(self, LOGGER_ATTRIBUTE_HANDLER, None):
+            return self.logger_handler
+
+        # Detect Powertools logger by checking for unique handler
+        # Retrieve the first handler if it's a Powertools instance
+        if getattr(self._logger, "powertools_handler", None):
+            return self._logger.handlers[0]
+
+        # for children, use parent's handler
+        if self.child:
+            return getattr(self._logger.parent, LOGGER_ATTRIBUTE_POWERTOOLS_HANDLER, None)  # type: ignore[return-value] # always checked in formatting
+
+        # otherwise, create a new stream handler (first time init)
+        return logging.StreamHandler(self._stream)
+
     def _init_logger(
         self,
         formatter_options: dict | None = None,
@@ -317,6 +337,7 @@ def _init_logger(
         # std logging will return the same Logger with our attribute if name is reused
         logger.debug(f"Marking logger {self.service} as preconfigured")
         self._logger.init = True  # type: ignore[attr-defined]
+        self._logger.powertools_handler = self.logger_handler  # type: ignore[attr-defined]
 
     def _configure_sampling(self) -> None:
         """Dynamically set log level based on sampling rate
@@ -723,13 +744,20 @@ def registered_handler(self) -> logging.Handler:
         """Convenience property to access the first logger handler"""
         # We ignore mypy here because self.child encodes whether or not self._logger.parent is
         # None, mypy can't see this from context but we can
-        handlers = self._logger.parent.handlers if self.child else self._logger.handlers  # type: ignore[union-attr]
-        return handlers[0]
+        return self._get_handler()
 
     @property
     def registered_formatter(self) -> BasePowertoolsFormatter:
         """Convenience property to access the first logger formatter"""
-        return self.registered_handler.formatter  # type: ignore[return-value]
+        handler = self.registered_handler
+        if handler is None:
+            raise OrphanedChildLoggerError(
+                "Orphan child loggers cannot append nor remove keys until a parent is initialized first. "
+                "To solve this issue, you can A) make sure a parent logger is initialized first, or B) move append/remove keys operations to a later stage."  # noqa: E501
+                "Reference: https://docs.powertools.aws.dev/lambda/python/latest/core/logger/#reusing-logger-across-your-code",
+            )
+
+        return cast(BasePowertoolsFormatter, handler.formatter)
 
     @property
     def log_level(self) -> int:

Diff for: docs/core/logger.md

@@ -722,31 +722,9 @@ stateDiagram-v2
 ```
 </center>
 
-> Python Logging hierarchy happens via the dot notation: `service`, `service.child`, `service.child_2`
-For inheritance, Logger uses a `child=True` parameter along with `service` being the same value across Loggers.
+For inheritance, Logger uses `child` parameter to ensure we don't compete with its parents config. We name child Loggers following Python's convention: _`{service}`.`{filename}`_.
 
-For child Loggers, we introspect the name of your module where `Logger(child=True, service="name")` is called, and we name your Logger as **{service}.{filename}**.
-
-???+ danger
-    A common issue when migrating from other Loggers is that `service` might be defined in the parent Logger (no child param), and not defined in the child Logger:
-
-=== "logging_inheritance_bad.py"
-
-    ```python hl_lines="1 9"
-    --8<-- "examples/logger/src/logging_inheritance_bad.py"
-    ```
-
-=== "logging_inheritance_module.py"
-    ```python hl_lines="1 9"
-    --8<-- "examples/logger/src/logging_inheritance_module.py"
-    ```
-
-In this case, Logger will register a Logger named `payment`, and a Logger named `service_undefined`. The latter isn't inheriting from the parent, and will have no handler, resulting in no message being logged to standard output.
-
-???+ tip
-    This can be fixed by either ensuring both has the `service` value as `payment`, or simply use the environment variable `POWERTOOLS_SERVICE_NAME` to ensure service value will be the same across all Loggers when not explicitly set.
-
-Do this instead:
+Changes are bidirectional between parents and loggers. That is, appending a key in a child or parent will ensure both have them. This means, having the same `service` name is important when instantiating them.
 
 === "logging_inheritance_good.py"
 
@@ -773,7 +751,6 @@ There are two important side effects when using child loggers:
     ```
 
 === "logging_inheritance_module.py"
-
     ```python hl_lines="1 9"
     --8<-- "examples/logger/src/logging_inheritance_module.py"
     ```

Diff for: tests/e2e/logger/handlers/multiple_logger_instances.py

@@ -0,0 +1,15 @@
+from aws_lambda_powertools import Logger
+
+# Instance 1
+logger = Logger()
+
+# Simulating importing from another file
+logger = Logger()
+
+
+@logger.inject_lambda_context
+def lambda_handler(event, context):
+    message, append_keys = event.get("message", ""), event.get("append_keys", {})
+    logger.append_keys(**append_keys)
+    logger.info(message)
+    return "success"

Diff for: tests/e2e/logger/test_logger.py

@@ -29,6 +29,16 @@ def tz_handler_fn_arn(infrastructure: dict) -> str:
     return infrastructure.get("TzHandlerArn", "")
 
 
+@pytest.fixture
+def multiple_logger_instances_fn(infrastructure: dict) -> str:
+    return infrastructure.get("MultipleLoggerInstances", "")
+
+
+@pytest.fixture
+def multiple_logger_instances_arn(infrastructure: dict) -> str:
+    return infrastructure.get("MultipleLoggerInstancesArn", "")
+
+
 @pytest.mark.xdist_group(name="logger")
 def test_basic_lambda_logs_visible(basic_handler_fn, basic_handler_fn_arn):
     # GIVEN
@@ -50,6 +60,31 @@ def test_basic_lambda_logs_visible(basic_handler_fn, basic_handler_fn_arn):
     assert logs.have_keys(*LOGGER_LAMBDA_CONTEXT_KEYS) is True
 
 
+@pytest.mark.xdist_group(name="logger")
+def test_multiple_logger_instances(multiple_logger_instances_fn, multiple_logger_instances_arn):
+    # GIVEN
+    message = "logs should be visible with default settings"
+    custom_key = "order_id"
+    additional_keys = {custom_key: f"{uuid4()}"}
+    payload = json.dumps({"message": message, "append_keys": additional_keys})
+
+    # WHEN
+    _, execution_time = data_fetcher.get_lambda_response(lambda_arn=multiple_logger_instances_arn, payload=payload)
+    data_fetcher.get_lambda_response(lambda_arn=multiple_logger_instances_arn, payload=payload)
+
+    # THEN
+    logs = data_fetcher.get_logs(
+        function_name=multiple_logger_instances_fn,
+        start_time=execution_time,
+        minimum_log_entries=2,
+    )
+
+    assert len(logs) == 2
+    assert len(logs.get_cold_start_log()) == 1
+    assert len(logs.get_log(key=custom_key)) == 2
+    assert logs.have_keys(*LOGGER_LAMBDA_CONTEXT_KEYS) is True
+
+
 @pytest.mark.xdist_group(name="logger")
 @pytest.mark.parametrize("tz", ["US/Eastern", "UTC", "Asia/Shanghai"])
 @pytest.mark.parametrize("datefmt", ["%z", None])

Diff for: tests/functional/logger/required_dependencies/test_logger.py

@@ -16,7 +16,7 @@
 
 from aws_lambda_powertools import Logger
 from aws_lambda_powertools.logging import correlation_paths
-from aws_lambda_powertools.logging.exceptions import InvalidLoggerSamplingRateError
+from aws_lambda_powertools.logging.exceptions import InvalidLoggerSamplingRateError, OrphanedChildLoggerError
 from aws_lambda_powertools.logging.formatter import (
     BasePowertoolsFormatter,
     LambdaPowertoolsFormatter,
@@ -1286,3 +1286,84 @@ def test_clear_state_log_output(stdout, service_name):
     assert "custom_key" in first_log
     assert first_log["custom_key"] == "test_value"
     assert "custom_key" not in second_log
+
+
+def test_logger_registered_handler_is_custom_handler(service_name):
+    # GIVEN a library or environment pre-setup a logger for us using the same name (see #4277)
+    class ForeignHandler(logging.StreamHandler): ...
+
+    foreign_handler = ForeignHandler()
+    logging.getLogger(service_name).addHandler(foreign_handler)
+
+    # WHEN Logger init with a custom handler
+    custom_handler = logging.StreamHandler()
+    logger = Logger(service=service_name, logger_handler=custom_handler)
+
+    # THEN registered handler should always return what we provided
+    assert logger.registered_handler is not foreign_handler
+    assert logger.registered_handler is custom_handler
+    assert logger.logger_handler is custom_handler
+    assert logger.handlers == [foreign_handler, custom_handler]
+
+
+def test_child_logger_registered_handler_is_custom_handler(service_name):
+    # GIVEN
+    class ForeignHandler(logging.StreamHandler): ...
+
+    foreign_handler = ForeignHandler()
+    logging.getLogger(service_name).addHandler(foreign_handler)
+
+    custom_handler = logging.StreamHandler()
+    custom_handler.name = "CUSTOM HANDLER"
+    parent = Logger(service=service_name, logger_handler=custom_handler)
+
+    # WHEN a child Logger init
+    child = Logger(service=service_name, child=True)
+
+    # THEN child registered handler should always return what we provided in the parent
+    assert child.registered_handler is not foreign_handler
+    assert child.registered_handler is custom_handler
+    assert child.registered_handler is parent.registered_handler
+
+
+def test_logger_handler_is_created_despite_env_pre_setup(service_name):
+    # GIVEN a library or environment pre-setup a logger for us using the same name
+    environment_handler = logging.StreamHandler()
+    logging.getLogger(service_name).addHandler(environment_handler)
+
+    # WHEN Logger init without a custom handler
+    logger = Logger(service=service_name)
+
+    # THEN registered handler should be Powertools default handler, not env
+    assert logger.registered_handler is not environment_handler
+
+
+def test_child_logger_append_keys_before_parent(stdout, service_name):
+    # GIVEN a child Logger is initialized before its/without parent
+    child = Logger(stream=stdout, service=service_name, child=True)
+
+    # WHEN a child Logger appends a key
+    # THEN it will raise an AttributeError
+    with pytest.raises(OrphanedChildLoggerError):
+        child.append_keys(customer_id="value")
+
+
+def test_powertools_logger_handler_is_created_only_once_and_propagated(lambda_context, stdout, service_name):
+    # GIVEN an instance of Logger
+    logger = Logger(service=service_name, stream=stdout)
+    request_id = "xxx-111-222"
+    mock_event = {"requestContext": {"requestId": request_id}}
+
+    # GIVEN another instance of Logger to mimic importing from another file
+    logger = Logger(service=service_name, stream=stdout)
+
+    # WHEN we use inject_lambda_context
+    @logger.inject_lambda_context(correlation_id_path=correlation_paths.API_GATEWAY_REST)
+    def handler(event, context):
+        logger.info("Foo")
+
+    handler(mock_event, lambda_context)
+
+    # THEN we must be able to inject context
+    log = capture_logging_output(stdout)
+    assert request_id == log["correlation_id"]