Use AWS_LAMBDA_LOG_LEVEL env var to set log level for text format

Author: Artem Krivonos

awslambdaric/bootstrap.py

@@ -19,13 +19,16 @@
     _JSON_FRAME_TYPES,
     JsonFormatter,
     LogFormat,
+    _get_log_level_from_env_var,
 )
 from .lambda_runtime_marshaller import to_json
 
 ERROR_LOG_LINE_TERMINATE = "\r"
 ERROR_LOG_IDENT = "\u00a0"  # NO-BREAK SPACE U+00A0
 _AWS_LAMBDA_LOG_FORMAT = LogFormat.from_str(os.environ.get("AWS_LAMBDA_LOG_FORMAT"))
-_AWS_LAMBDA_LOG_LEVEL = os.environ.get("AWS_LAMBDA_LOG_LEVEL", "").upper()
+_AWS_LAMBDA_LOG_LEVEL = _get_log_level_from_env_var(
+    os.environ.get("AWS_LAMBDA_LOG_LEVEL")
+)
 
 
 def _get_handler(handler):
@@ -419,10 +422,6 @@ def _setup_logging(log_format, log_level, log_sink):
     logger_handler = LambdaLoggerHandler(log_sink)
     if log_format == LogFormat.JSON:
         logger_handler.setFormatter(JsonFormatter())
-
-        logging.addLevelName(logging.DEBUG, "TRACE")
-        if log_level in logging._nameToLevel:
-            logger.setLevel(log_level)
     else:
         logger_handler.setFormatter(
             logging.Formatter(
@@ -431,6 +430,9 @@ def _setup_logging(log_format, log_level, log_sink):
             )
         )
 
+    if log_level in logging._nameToLevel:
+        logger.setLevel(log_level)
+
     logger_handler.addFilter(LambdaLoggerFilter())
     logger.addHandler(logger_handler)
 

awslambdaric/lambda_runtime_log_utils.py

@@ -45,6 +45,10 @@ def from_str(cls, value: str):
         return cls.TEXT.value
 
 
+def _get_log_level_from_env_var(log_level):
+    return {None: "", "TRACE": "DEBUG"}.get(log_level, log_level).upper()
+
+
 _JSON_FRAME_TYPES = {
     logging.NOTSET: 0xA55A0002.to_bytes(4, "big"),
     logging.DEBUG: 0xA55A000A.to_bytes(4, "big"),

tests/test_bootstrap.py

@@ -18,7 +18,7 @@
 
 import awslambdaric.bootstrap as bootstrap
 from awslambdaric.lambda_runtime_exception import FaultException
-from awslambdaric.lambda_runtime_log_utils import LogFormat
+from awslambdaric.lambda_runtime_log_utils import LogFormat, _get_log_level_from_env_var
 from awslambdaric.lambda_runtime_marshaller import LambdaMarshaller
 
 
@@ -1180,13 +1180,13 @@ def test_log_level(self) -> None:
             (LogFormat.JSON, "ERROR", logging.ERROR),
             (LogFormat.JSON, "FATAL", logging.CRITICAL),
             # Log level is set only for Json format
-            (LogFormat.TEXT, "TRACE", logging.NOTSET),
-            (LogFormat.TEXT, "DEBUG", logging.NOTSET),
-            (LogFormat.TEXT, "INFO", logging.NOTSET),
-            (LogFormat.TEXT, "WARN", logging.NOTSET),
-            (LogFormat.TEXT, "ERROR", logging.NOTSET),
-            (LogFormat.TEXT, "FATAL", logging.NOTSET),
-            ("Unknown format", "INFO", logging.NOTSET),
+            (LogFormat.TEXT, "TRACE", logging.DEBUG),
+            (LogFormat.TEXT, "DEBUG", logging.DEBUG),
+            (LogFormat.TEXT, "INFO", logging.INFO),
+            (LogFormat.TEXT, "WARN", logging.WARN),
+            (LogFormat.TEXT, "ERROR", logging.ERROR),
+            (LogFormat.TEXT, "FATAL", logging.CRITICAL),
+            ("Unknown format", "INFO", logging.INFO),
             # if level is unknown fall back to default
             (LogFormat.JSON, "Unknown level", logging.NOTSET),
         ]
@@ -1196,7 +1196,11 @@ def test_log_level(self) -> None:
                 logging.getLogger().handlers.clear()
                 logging.getLogger().level = logging.NOTSET
 
-                bootstrap._setup_logging(fmt, log_level, bootstrap.StandardLogSink())
+                bootstrap._setup_logging(
+                    fmt,
+                    _get_log_level_from_env_var(log_level),
+                    bootstrap.StandardLogSink(),
+                )
 
                 self.assertEqual(expected_level, logging.getLogger().level)