add LambdaLoggerHandlerWithFrameType

Author: Artem Krivonos

awslambdaric/bootstrap.py

@@ -299,6 +299,16 @@ def __init__(self, log_sink):
         logging.Handler.__init__(self)
         self.log_sink = log_sink
 
+    def emit(self, record):
+        msg = self.format(record)
+        self.log_sink.log(msg)
+
+
+class LambdaLoggerHandlerWithFrameType(logging.Handler):
+    def __init__(self, log_sink):
+        super().__init__()
+        self.log_sink = log_sink
+
     def emit(self, record):
         self.log_sink.log(
             self.format(record),
@@ -425,7 +435,12 @@ def create_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 == LogFormat.JSON or log_level:
+        logger_handler = LambdaLoggerHandlerWithFrameType(log_sink)
+    else:
+        logger_handler = LambdaLoggerHandler(log_sink)
+
     if log_format == LogFormat.JSON:
         logger_handler.setFormatter(JsonFormatter())
     else:

awslambdaric/lambda_runtime_log_utils.py

@@ -71,7 +71,7 @@ def _get_log_level_from_env_var(log_level):
 _encode_json = _json_encoder.encode
 
 
-def _format_log_level(record: logging.LogRecord):
+def _format_log_level(record: logging.LogRecord) -> int:
     return min(50, max(0, record.levelno)) // 10 * 10
 
 

tests/test_bootstrap.py

@@ -1205,6 +1205,52 @@ def test_log_level(self) -> None:
                 self.assertEqual(expected_level, logging.getLogger().level)
 
 
+class TestLambdaLoggerHandlerSetup(unittest.TestCase):
+    @classmethod
+    def tearDownClass(cls):
+        importlib.reload(bootstrap)
+        logging.getLogger().handlers.clear()
+        logging.getLogger().level = logging.NOTSET
+
+    def test_handler_setup(self, *_):
+        test_cases = [
+            (62, 0xA55A0003, 46, {}),
+            (133, 0xa55a001a, 117, {"AWS_LAMBDA_LOG_FORMAT": "JSON"}),
+            (62, 0xa55a001b, 46, {"AWS_LAMBDA_LOG_LEVEL": "INFO"}),
+        ]
+
+        for total_length, header, message_length, env_vars in test_cases:
+            with patch.dict(os.environ, env_vars, clear=True), NamedTemporaryFile() as temp_file:
+                importlib.reload(bootstrap)
+                logging.getLogger().handlers.clear()
+                logging.getLogger().level = logging.NOTSET
+
+                before = int(time.time_ns() / 1000)
+                with bootstrap.FramedTelemetryLogSink(
+                    os.open(temp_file.name, os.O_CREAT | os.O_RDWR)
+                ) as ls:
+                    bootstrap._setup_logging(bootstrap._AWS_LAMBDA_LOG_FORMAT, bootstrap._AWS_LAMBDA_LOG_LEVEL, ls)
+                    logger = logging.getLogger()
+                    logger.critical("critical")
+                after = int(time.time_ns() / 1000)
+
+                content = open(temp_file.name, "rb").read()
+                self.assertEqual(len(content), total_length)
+
+                pos = 0
+                frame_type = int.from_bytes(content[pos : pos + 4], "big")
+                self.assertEqual(frame_type, header)
+                pos += 4
+
+                length = int.from_bytes(content[pos : pos + 4], "big")
+                self.assertEqual(length, message_length)
+                pos += 4
+
+                timestamp = int.from_bytes(content[pos : pos + 8], "big")
+                self.assertTrue(before <= timestamp)
+                self.assertTrue(timestamp <= after)
+
+
 class TestLogging(unittest.TestCase):
     @classmethod
     def setUpClass(cls) -> None: