move frame_type to JsonFormatter, change make_error, fix tests

Author: Artem Krivonos

awslambdaric/bootstrap.py

@@ -17,14 +17,14 @@
     JsonFormatter,
     _DATETIME_FORMAT,
     _JSON_FORMAT,
-    _TEXT_FORMAT,
+    _FRAME_TYPES,
+    _DEFAULT_FRAME_TYPE,
     _get_log_format_from_str,
 )
 from .lambda_runtime_marshaller import to_json
 
 ERROR_LOG_LINE_TERMINATE = "\r"
 ERROR_LOG_IDENT = "\u00a0"  # NO-BREAK SPACE U+00A0
-RUNTIME_ERROR_LOGGER_NAME = "system"
 _AWS_LAMBDA_LOG_FORMAT = _get_log_format_from_str(
     os.environ.get("AWS_LAMBDA_LOG_FORMAT", "TEXT").upper()
 )
@@ -86,14 +86,12 @@ def result(*args):
 
 
 def make_error(
-    error_message, error_type, stack_trace, invoke_id=None, level=logging.ERROR
+    error_message,
+    error_type,
+    stack_trace,
+    invoke_id=None,
 ):
     result = {
-        "timestamp": time.strftime(
-            _DATETIME_FORMAT, logging.Formatter.converter(time.time())
-        ),
-        "log_level": level,
-        "logger": RUNTIME_ERROR_LOGGER_NAME,
         "errorMessage": error_message if error_message else "",
         "errorType": error_type if error_type else "",
         "requestId": invoke_id if invoke_id is not None else "",
@@ -111,51 +109,48 @@ def replace_line_indentation(line, indent_char, new_indent_char):
     return (new_indent_char * ident_chars_count) + line[ident_chars_count:]
 
 
-def log_error_json(error_result, log_sink):
-    log_level = error_result.pop("log_level", logging.ERROR)
-    error_result["level"] = logging.getLevelName(log_level)
-    log_sink.log_error(
-        [to_json(error_result, ensure_ascii=False)],
-        log_level=log_level,
-    )
+if _AWS_LAMBDA_LOG_FORMAT == _JSON_FORMAT:
+    def log_error(error_result, log_sink):
+        error_result = {
+            "timestamp": time.strftime(
+                _DATETIME_FORMAT, logging.Formatter.converter(time.time())
+            ),
+            "log_level": "ERROR",
+            **error_result,
+        }
+        log_sink.log_error(
+            [to_json(error_result)],
+        )
 
+else:
+    def log_error(error_result, log_sink):
+        error_description = "[ERROR]"
 
-def log_error_text(error_result, log_sink):
-    error_description = "[ERROR]"
+        error_result_type = error_result.get("errorType")
+        if error_result_type:
+            error_description += " " + error_result_type
 
-    error_result_type = error_result.get("errorType")
-    if error_result_type:
-        error_description += " " + error_result_type
+        error_result_message = error_result.get("errorMessage")
+        if error_result_message:
+            if error_result_type:
+                error_description += ":"
+            error_description += " " + error_result_message
 
-    error_result_message = error_result.get("errorMessage")
-    if error_result_message:
-        if error_result_type:
-            error_description += ":"
-        error_description += " " + error_result_message
-
-    error_message_lines = [error_description]
-
-    stack_trace = error_result.get("stackTrace")
-    if stack_trace is not None:
-        error_message_lines += ["Traceback (most recent call last):"]
-        for trace_element in stack_trace:
-            if trace_element == "":
-                error_message_lines += [""]
-            else:
-                for trace_line in trace_element.splitlines():
-                    error_message_lines += [
-                        replace_line_indentation(trace_line, " ", ERROR_LOG_IDENT)
-                    ]
-
-    log_sink.log_error(
-        error_message_lines, log_level=error_result.get("log_level", logging.ERROR)
-    )
+        error_message_lines = [error_description]
 
+        stack_trace = error_result.get("stackTrace")
+        if stack_trace is not None:
+            error_message_lines += ["Traceback (most recent call last):"]
+            for trace_element in stack_trace:
+                if trace_element == "":
+                    error_message_lines += [""]
+                else:
+                    for trace_line in trace_element.splitlines():
+                        error_message_lines += [
+                            replace_line_indentation(trace_line, " ", ERROR_LOG_IDENT)
+                        ]
 
-if _AWS_LAMBDA_LOG_FORMAT == _JSON_FORMAT:
-    log_error = log_error_json
-else:
-    log_error = log_error_text
+        log_sink.log_error(error_message_lines)
 
 
 def handle_event_request(
@@ -189,7 +184,10 @@ def handle_event_request(
     except FaultException as e:
         xray_fault = make_xray_fault("LambdaValidationError", e.msg, os.getcwd(), [])
         error_result = make_error(
-            e.msg, e.exception_type, e.trace, invoke_id, level=logging.FATAL
+            e.msg,
+            e.exception_type,
+            e.trace,
+            invoke_id,
         )
 
     except Exception:
@@ -262,7 +260,6 @@ def build_fault_result(exc_info, msg):
         msg if msg else str(value),
         etype.__name__,
         traceback.format_list(tb_tuples),
-        level=logging.FATAL,
     )
 
 
@@ -300,9 +297,7 @@ def emit(self, record):
         msg = self.format(record)
 
         self.log_sink.log(
-            msg,
-            log_level=record.levelno,
-            log_format=getattr(record, "log_format", _TEXT_FORMAT),
+            msg, frame_type=getattr(record, "_frame_type", _FRAME_TYPES.get(()))
         )
 
 
@@ -344,29 +339,15 @@ def __enter__(self):
     def __exit__(self, exc_type, exc_value, exc_tb):
         pass
 
-    def log(self, msg, log_level=None, log_format=None):
+    def log(self, msg, frame_type=None):
         sys.stdout.write(msg)
 
-    def log_error(self, message_lines, log_level=logging.ERROR):
+    def log_error(self, message_lines):
         error_message = ERROR_LOG_LINE_TERMINATE.join(message_lines) + "\n"
         sys.stdout.write(error_message)
 
 
-FRAME_TYPES = {
-    (_JSON_FORMAT, logging.NOTSET): 0xA55A0002.to_bytes(4, "big"),
-    (_JSON_FORMAT, logging.DEBUG): 0xA55A000A.to_bytes(4, "big"),
-    (_JSON_FORMAT, logging.INFO): 0xA55A000E.to_bytes(4, "big"),
-    (_JSON_FORMAT, logging.WARNING): 0xA55A0012.to_bytes(4, "big"),
-    (_JSON_FORMAT, logging.ERROR): 0xA55A0016.to_bytes(4, "big"),
-    (_JSON_FORMAT, logging.CRITICAL): 0xA55A001A.to_bytes(4, "big"),
-    (_TEXT_FORMAT, logging.NOTSET): 0xA55A0003.to_bytes(4, "big"),
-    (_TEXT_FORMAT, logging.DEBUG): 0xA55A000B.to_bytes(4, "big"),
-    (_TEXT_FORMAT, logging.INFO): 0xA55A000F.to_bytes(4, "big"),
-    (_TEXT_FORMAT, logging.WARNING): 0xA55A0013.to_bytes(4, "big"),
-    (_TEXT_FORMAT, logging.ERROR): 0xA55A0017.to_bytes(4, "big"),
-    (_TEXT_FORMAT, logging.CRITICAL): 0xA55A001B.to_bytes(4, "big"),
-}
-DEFAULT_FRAME_TYPE = 0xA55A0003.to_bytes(4, "big")
+_ERROR_FRAME_TYPE = _FRAME_TYPES[(_AWS_LAMBDA_LOG_FORMAT, logging.ERROR)]
 
 
 class FramedTelemetryLogSink(object):
@@ -395,8 +376,7 @@ def __enter__(self):
     def __exit__(self, exc_type, exc_value, exc_tb):
         self.file.close()
 
-    def log(self, msg, log_level=logging.NOTSET, log_format: int = _TEXT_FORMAT):
-        frame_type = FRAME_TYPES.get((log_format, log_level), DEFAULT_FRAME_TYPE)
+    def log(self, msg, frame_type=_DEFAULT_FRAME_TYPE):
         encoded_msg = msg.encode("utf8")
 
         timestamp = int(time.time_ns() / 1000)  # UNIX timestamp in microseconds
@@ -408,9 +388,12 @@ def log(self, msg, log_level=logging.NOTSET, log_format: int = _TEXT_FORMAT):
         )
         self.file.write(log_msg)
 
-    def log_error(self, message_lines, log_level=logging.ERROR):
+    def log_error(self, message_lines):
         error_message = "\n".join(message_lines)
-        self.log(error_message, log_level=log_level)
+        self.log(
+            error_message,
+            frame_type=_ERROR_FRAME_TYPE,
+        )
 
 
 def update_xray_env_variable(xray_trace_id):
@@ -435,7 +418,6 @@ def create_log_sink():
 
 
 def setup_logging(log_format, log_level, log_sink):
-    logging._levelToName[logging.CRITICAL] = "FATAL"
     logging.Formatter.converter = time.gmtime
     logger = logging.getLogger()
     logger_handler = LambdaLoggerHandler(log_sink)

awslambdaric/lambda_runtime_log_utils.py

@@ -4,8 +4,7 @@
 
 import logging
 import traceback
-
-from .lambda_runtime_marshaller import to_json
+import json
 
 _DATETIME_FORMAT = "%Y-%m-%dT%H:%M:%SZ"
 _RESERVED_FIELDS = {
@@ -30,10 +29,28 @@
     "processName",
     "process",
     "aws_request_id",
-    "log_format",
+    "_frame_type",
 }
 _JSON_FORMAT = 0b0
 _TEXT_FORMAT = 0b1
+_FRAME_TYPES = {
+    (_JSON_FORMAT, logging.NOTSET): 0xA55A0002.to_bytes(4, "big"),
+    (_JSON_FORMAT, logging.DEBUG): 0xA55A000A.to_bytes(4, "big"),
+    (_JSON_FORMAT, logging.INFO): 0xA55A000E.to_bytes(4, "big"),
+    (_JSON_FORMAT, logging.WARNING): 0xA55A0012.to_bytes(4, "big"),
+    (_JSON_FORMAT, logging.ERROR): 0xA55A0016.to_bytes(4, "big"),
+    (_JSON_FORMAT, logging.CRITICAL): 0xA55A001A.to_bytes(4, "big"),
+    (_TEXT_FORMAT, logging.NOTSET): 0xA55A0003.to_bytes(4, "big"),
+    (_TEXT_FORMAT, logging.DEBUG): 0xA55A000B.to_bytes(4, "big"),
+    (_TEXT_FORMAT, logging.INFO): 0xA55A000F.to_bytes(4, "big"),
+    (_TEXT_FORMAT, logging.WARNING): 0xA55A0013.to_bytes(4, "big"),
+    (_TEXT_FORMAT, logging.ERROR): 0xA55A0017.to_bytes(4, "big"),
+    (_TEXT_FORMAT, logging.CRITICAL): 0xA55A001B.to_bytes(4, "big"),
+}
+_DEFAULT_FRAME_TYPE = _FRAME_TYPES[(_TEXT_FORMAT, logging.NOTSET)]
+
+
+_encode_json = json.JSONEncoder(ensure_ascii=False).encode
 
 
 def _get_log_format_from_str(value: str):
@@ -75,12 +92,15 @@ def format_location(record: logging.LogRecord):
 
     @staticmethod
     def format_log_level(record: logging.LogRecord):
-        level = min(50, max(0, record.levelno)) // 10 * 10
-        record.levelname = logging.getLevelName(level)
+        record.levelno = min(50, max(0, record.levelno)) // 10 * 10
+        record.levelname = logging.getLevelName(record.levelno)
 
     def format(self, record: logging.LogRecord) -> str:
         self.format_log_level(record)
-        record.log_format = _JSON_FORMAT
+        record._frame_type = _FRAME_TYPES.get(
+            (_JSON_FORMAT, record.levelno),
+            _FRAME_TYPES[(_JSON_FORMAT, logging.NOTSET)],
+        )
 
         result = {
             "timestamp": self.formatTime(record, self.datefmt),
@@ -101,4 +121,4 @@ def format(self, record: logging.LogRecord) -> str:
 
         result = {k: v for k, v in result.items() if v is not None}
 
-        return to_json(result, ensure_ascii=False) + "\n"
+        return _encode_json(result) + "\n"

awslambdaric/lambda_runtime_marshaller.py

@@ -24,8 +24,8 @@ def default(self, obj):
         return super().default(obj)
 
 
-def to_json(obj, ensure_ascii=True):
-    return Encoder(ensure_ascii=ensure_ascii).encode(obj)
+def to_json(obj):
+    return Encoder().encode(obj)
 
 
 class LambdaMarshaller: