Adding decorator

Author: leandrodamascena

aws_lambda_powertools/logging/logger.py

@@ -400,6 +400,7 @@ def inject_lambda_context(
         log_event: bool | None = None,
         correlation_id_path: str | None = None,
         clear_state: bool | None = False,
+        flush_buffer_on_uncaught_error: bool = False,
     ) -> AnyCallableT: ...
 
     @overload
@@ -409,6 +410,7 @@ def inject_lambda_context(
         log_event: bool | None = None,
         correlation_id_path: str | None = None,
         clear_state: bool | None = False,
+        flush_buffer_on_uncaught_error: bool = False,
     ) -> Callable[[AnyCallableT], AnyCallableT]: ...
 
     def inject_lambda_context(
@@ -417,6 +419,7 @@ def inject_lambda_context(
         log_event: bool | None = None,
         correlation_id_path: str | None = None,
         clear_state: bool | None = False,
+        flush_buffer_on_uncaught_error: bool = False,
     ) -> Any:
         """Decorator to capture Lambda contextual info and inject into logger
 
@@ -473,6 +476,7 @@ def handler(event, context):
                 log_event=log_event,
                 correlation_id_path=correlation_id_path,
                 clear_state=clear_state,
+                flush_buffer_on_uncaught_error=flush_buffer_on_uncaught_error,
             )
 
         log_event = resolve_truthy_env_var_choice(
@@ -505,7 +509,18 @@ def decorate(event, context, *args, **kwargs):
             if self.sampling_rate and not cold_start:
                 self.refresh_sample_rate_calculation()
 
-            return lambda_handler(event, context, *args, **kwargs)
+            try:
+                # Execute the Lambda handler with provided event and context
+                return lambda_handler(event, context, *args, **kwargs)
+            except:
+                # Re-raise any exceptions that occur during handler execution
+                raise
+            finally:
+                # Flush the log buffer if configured to do so on uncaught errors
+                # Ensures logging state is cleaned up even if an exception is raised
+                if flush_buffer_on_uncaught_error:
+                    logger.debug("An error happened, flushing the buffer")
+                    self.flush_buffer()
 
         return decorate
 

tests/functional/logger/required_dependencies/test_powertools_logger_buffer.py

@@ -273,3 +273,51 @@ def test_logger_buffer_log_output_for_levels_above_minimum(log_level, stdout, se
     log = capture_multiple_logging_statements_output(stdout)
     assert len(log) == 1
     assert log[0]["message"] == msg
+
+
+def test_logger_buffer_flush_on_uncaught_exception(stdout, service_name, monkeypatch, lambda_context):
+    monkeypatch.setenv(constants.XRAY_TRACE_ID_ENV, "1234")
+
+    # GIVEN: A logger configured with a large buffer and error-based flushing
+    logger_buffer_config = LoggerBufferConfig(max_size=10240, minimum_log_level="DEBUG", flush_on_error=True)
+    logger = Logger(level="DEBUG", service=service_name, stream=stdout, logger_buffer=logger_buffer_config)
+
+    @logger.inject_lambda_context(flush_buffer_on_uncaught_error=True)
+    def handler(event, context):
+        # Log messages that should be flushed when an exception occurs
+        logger.debug("this log line will be flushed after error - 1")
+        logger.debug("this log line will be flushed after error - 2")
+        raise ValueError("Test error")
+
+    # WHEN Invoking the handler and expecting a ValueError
+    with pytest.raises(ValueError):
+        handler({}, lambda_context)
+
+    # THEN Verify that buffered log messages are flushed before the exception
+    log = capture_multiple_logging_statements_output(stdout)
+    assert len(log) == 2, "Expected two log messages to be flushed"
+    assert log[0]["message"] == "this log line will be flushed after error - 1"
+    assert log[1]["message"] == "this log line will be flushed after error - 2"
+
+
+def test_logger_buffer_not_flush_on_uncaught_exception(stdout, service_name, monkeypatch, lambda_context):
+    monkeypatch.setenv(constants.XRAY_TRACE_ID_ENV, "1234")
+
+    # GIVEN: A logger configured with a large buffer and error-based flushing
+    logger_buffer_config = LoggerBufferConfig(max_size=10240, minimum_log_level="DEBUG", flush_on_error=True)
+    logger = Logger(level="DEBUG", service=service_name, stream=stdout, logger_buffer=logger_buffer_config)
+
+    @logger.inject_lambda_context(flush_buffer_on_uncaught_error=False)
+    def handler(event, context):
+        # Log messages that should be flushed when an exception occurs
+        logger.debug("this log line will be flushed after error - 1")
+        logger.debug("this log line will be flushed after error - 2")
+        raise ValueError("Test error")
+
+    # WHEN Invoking the handler and expecting a ValueError
+    with pytest.raises(ValueError):
+        handler({}, lambda_context)
+
+    # THEN Verify that buffered log messages are flushed before the exception
+    log = capture_multiple_logging_statements_output(stdout)
+    assert len(log) == 0