Add JsonFormatter, INSTANCE_LOGGING_CONTEXT, INVOCATION_LOGGING_CONTEXT and environ variable AWS_LAMBDA_LOG_FORMAT that manages logging format

Author: Artem Krivonos

awslambdaric/bootstrap.py

@@ -14,6 +14,7 @@
 from .lambda_runtime_client import LambdaRuntimeClient
 from .lambda_runtime_exception import FaultException
 from .lambda_runtime_marshaller import to_json
+from .lambda_runtime_log_utils import JsonFormatter, INSTANCE_LOGGING_CONTEXT
 
 ERROR_LOG_LINE_TERMINATE = "\r"
 ERROR_LOG_IDENT = "\u00a0"  # NO-BREAK SPACE U+00A0
@@ -367,6 +368,7 @@ def create_log_sink():
         return StandardLogSink()
 
 
+AWS_LAMBDA_LOG_FORMAT = os.environ.get("AWS_LAMBDA_LOG_FORMAT", "TEXT")
 _GLOBAL_AWS_REQUEST_ID = None
 
 
@@ -381,12 +383,17 @@ def run(app_root, handler, lambda_runtime_api_addr):
             logging.Formatter.converter = time.gmtime
             logger = logging.getLogger()
             logger_handler = LambdaLoggerHandler(log_sink)
-            logger_handler.setFormatter(
-                logging.Formatter(
-                    "[%(levelname)s]\t%(asctime)s.%(msecs)03dZ\t%(aws_request_id)s\t%(message)s\n",
-                    "%Y-%m-%dT%H:%M:%S",
+            if AWS_LAMBDA_LOG_FORMAT == "JSON":
+                logger_handler.setFormatter(JsonFormatter())
+                logging._levelToName[logging.CRITICAL] = "FATAL"
+            else:
+                logger_handler.setFormatter(
+                    logging.Formatter(
+                        "[%(levelname)s]\t%(asctime)s.%(msecs)03dZ\t%(aws_request_id)s\t%(message)s\n",
+                        "%Y-%m-%dT%H:%M:%S",
+                    )
                 )
-            )
+
             logger_handler.addFilter(LambdaLoggerFilter())
             logger.addHandler(logger_handler)
 
@@ -403,6 +410,7 @@ def run(app_root, handler, lambda_runtime_api_addr):
 
         while True:
             event_request = lambda_runtime_client.wait_next_invocation()
+            INSTANCE_LOGGING_CONTEXT.clear()
 
             _GLOBAL_AWS_REQUEST_ID = event_request.invoke_id
 

awslambdaric/lambda_runtime_log_utils.py

@@ -0,0 +1,96 @@
+"""
+Copyright 2023 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+"""
+
+import logging
+
+from .lambda_runtime_marshaller import to_json
+
+
+RESERVED_FIELDS = {
+    "name",
+    "msg",
+    "args",
+    "levelname",
+    "levelno",
+    "pathname",
+    "filename",
+    "module",
+    "exc_info",
+    "exc_text",
+    "stack_info",
+    "lineno",
+    "funcName",
+    "created",
+    "msecs",
+    "relativeCreated",
+    "thread",
+    "threadName",
+    "processName",
+    "process",
+    "aws_request_id",
+}
+
+
+class JsonFormatter(logging.Formatter):
+    def __init__(self):
+        super().__init__(datefmt="%Y-%m-%dT%H:%M:%SZ", validate=False)
+
+    def format_stacktrace(self, exc_info):
+        if not exc_info:
+            return None
+        return self.formatException(exc_info)
+
+    @staticmethod
+    def format_exception_name(exc_info):
+        if not exc_info:
+            return None
+
+        return exc_info[0].__name__
+
+    @staticmethod
+    def format_exception(exc_info):
+        if not exc_info:
+            return None
+
+        return str(exc_info[1])
+
+    @staticmethod
+    def format_location(record: logging.LogRecord):
+        return f"{record.pathname}:{record.funcName}:{record.lineno}"
+
+    def format(self, record: logging.LogRecord) -> str:
+        result = {
+            "timestamp": self.formatTime(record, self.datefmt),
+            "level": record.levelname,
+            "message": record.getMessage(),
+            "logger": record.name,
+            "stackTrace": self.format_stacktrace(record.exc_info),
+            "errorType": self.format_exception_name(record.exc_info),
+            "errorMessage": self.format_exception(record.exc_info),
+            "requestId": getattr(record, "aws_request_id", None),
+            "location": self.format_location(record),
+        }
+        result.update(
+            (key, value)
+            for key, value in record.__dict__.items()
+            if key not in RESERVED_FIELDS and key not in result
+        )
+        result.update(INSTANCE_LOGGING_CONTEXT)
+        result.update(INVOCATION_LOGGING_CONTEXT)
+
+        result = {k: v for k, v in result.items() if v is not None}
+
+        return to_json(result)
+
+
+INSTANCE_LOGGING_CONTEXT = dict()
+INVOCATION_LOGGING_CONTEXT = dict()
+
+
+def lambda_instance_logs_update(key: str, value: str):
+    INSTANCE_LOGGING_CONTEXT[str(key)] = str(value)
+
+
+def lambda_invocation_logs_update(key: str, value: str):
+    INVOCATION_LOGGING_CONTEXT[str(key)] = str(value)