add AWS_LAMBDA_LOG_LEVEL, return stack trace as a list, allow any value in extras

Author: Artem Krivonos

awslambdaric/bootstrap.py

@@ -369,16 +369,18 @@ def create_log_sink():
 
 
 AWS_LAMBDA_LOG_FORMAT = os.environ.get("AWS_LAMBDA_LOG_FORMAT", "TEXT")
+AWS_LAMBDA_LOG_LEVEL = os.environ.get("AWS_LAMBDA_LOG_LEVEL", "INFO")
 _GLOBAL_AWS_REQUEST_ID = None
 
 
-def setup_logging(log_format, log_sink):
+def setup_logging(log_format, log_level, log_sink):
     logging.Formatter.converter = time.gmtime
     logger = logging.getLogger()
     logger_handler = LambdaLoggerHandler(log_sink)
     if log_format == "JSON":
         logger_handler.setFormatter(JsonFormatter())
         logging._levelToName[logging.CRITICAL] = "FATAL"
+        logger.setLevel(log_level)
     else:
         logger_handler.setFormatter(
             logging.Formatter(
@@ -399,7 +401,7 @@ def run(app_root, handler, lambda_runtime_api_addr):
         lambda_runtime_client = LambdaRuntimeClient(lambda_runtime_api_addr)
 
         try:
-            setup_logging(AWS_LAMBDA_LOG_FORMAT, log_sink)
+            setup_logging(AWS_LAMBDA_LOG_FORMAT, AWS_LAMBDA_LOG_LEVEL, log_sink)
             global _GLOBAL_AWS_REQUEST_ID
 
             request_handler = _get_handler(handler)

awslambdaric/lambda_runtime_log_utils.py

@@ -3,6 +3,7 @@
 """
 
 import logging
+import traceback
 
 from .lambda_runtime_marshaller import to_json
 
@@ -35,10 +36,11 @@ class JsonFormatter(logging.Formatter):
     def __init__(self):
         super().__init__(datefmt="%Y-%m-%dT%H:%M:%SZ", validate=False)
 
-    def format_stacktrace(self, exc_info):
+    @staticmethod
+    def format_stacktrace(exc_info):
         if not exc_info:
             return None
-        return self.formatException(exc_info)
+        return traceback.format_tb(exc_info[2])
 
     @staticmethod
     def format_exception_name(exc_info):
@@ -90,9 +92,9 @@ def format(self, record: logging.LogRecord) -> str:
 INVOCATION_LOGGING_CONTEXT = dict()
 
 
-def lambda_instance_logs_update(key: str, value: str):
-    INSTANCE_LOGGING_CONTEXT[str(key)] = str(value)
+def lambda_instance_logs_update(key: str, value):
+    INSTANCE_LOGGING_CONTEXT[str(key)] = value
 
 
-def lambda_invocation_logs_update(key: str, value: str):
-    INVOCATION_LOGGING_CONTEXT[str(key)] = str(value)
+def lambda_invocation_logs_update(key: str, value):
+    INVOCATION_LOGGING_CONTEXT[str(key)] = value

tests/test_bootstrap.py

@@ -1184,7 +1184,7 @@ def test_multiple_frame(self):
 class TestLogging(unittest.TestCase):
     @classmethod
     def setUpClass(cls) -> None:
-        bootstrap.setup_logging("JSON", bootstrap.StandardLogSink())
+        bootstrap.setup_logging("JSON", "INFO", bootstrap.StandardLogSink())
 
     def tearDown(self) -> None:
         INVOCATION_LOGGING_CONTEXT.clear()
@@ -1252,7 +1252,9 @@ def test_exception(self, mock_stderr, mock_stdout):
         self.assertIn("stackTrace", exception_log)
         exception_log.pop("timestamp")
         exception_log.pop("location")
-        exception_log.pop("stackTrace")
+        stack_trace = exception_log.pop("stackTrace")
+
+        self.assertEqual(len(stack_trace), 1)
 
         self.assertEqual(
             exception_log,
@@ -1273,6 +1275,9 @@ def test_exception(self, mock_stderr, mock_stdout):
     def test_log_with_extra_params(self, mock_stderr, mock_stdout):
         lambda_instance_logs_update("instance_key", "instance_value")
         lambda_invocation_logs_update("invocation_key", "invocation_value")
+        lambda_invocation_logs_update("int_param", 42)
+        lambda_invocation_logs_update("list_param", ["1", 2, {}])
+        lambda_invocation_logs_update("dict_param", {"a": "b"})
 
         logging.getLogger("test.logger").error("test extra params")
 
@@ -1284,6 +1289,9 @@ def test_log_with_extra_params(self, mock_stderr, mock_stdout):
             {
                 "instance_key": "instance_value",
                 "invocation_key": "invocation_value",
+                "int_param": 42,
+                "list_param": ["1", 2, {}],
+                "dict_param": {"a": "b"},
                 "level": "ERROR",
                 "logger": "test.logger",
                 "message": "test extra params",
@@ -1292,6 +1300,29 @@ def test_log_with_extra_params(self, mock_stderr, mock_stdout):
         )
         self.assertEqual(mock_stderr.getvalue(), "")
 
+    @patch("sys.stdout", new_callable=StringIO)
+    @patch("sys.stderr", new_callable=StringIO)
+    def test_log_level(self, mock_stderr, mock_stdout):
+        logger = logging.getLogger("test.logger")
+
+        logger.debug("debug message")
+        logger.info("info message")
+
+        # print("some ", mock_stdout.getvalue())
+        data = json.loads(mock_stdout.getvalue())
+        data.pop("timestamp")
+
+        self.assertEqual(
+            data,
+            {
+                "level": "INFO",
+                "logger": "test.logger",
+                "message": "info message",
+                "requestId": "",
+            },
+        )
+        self.assertEqual(mock_stderr.getvalue(), "")
+
 
 class TestBootstrapModule(unittest.TestCase):
     @patch("awslambdaric.bootstrap.handle_event_request")