Skip to content

Use AWS_LAMBDA_LOG_LEVEL env var to set log level for text format #111

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 1 commit into from
Aug 29, 2023
Merged

Use AWS_LAMBDA_LOG_LEVEL env var to set log level for text format #111

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
39 changes: 31 additions & 8 deletions awslambdaric/bootstrap.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -17,15 +17,20 @@
_DATETIME_FORMAT,
_DEFAULT_FRAME_TYPE,
_JSON_FRAME_TYPES,
_TEXT_FRAME_TYPES,
JsonFormatter,
LogFormat,
_format_log_level,
_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):
Expand Down Expand Up @@ -122,7 +127,7 @@ def log_error(error_result, log_sink):
)

else:
_ERROR_FRAME_TYPE = _DEFAULT_FRAME_TYPE
_ERROR_FRAME_TYPE = _TEXT_FRAME_TYPES[logging.ERROR]

def log_error(error_result, log_sink):
error_description = "[ERROR]"
Expand Down Expand Up @@ -296,8 +301,22 @@ def __init__(self, log_sink):

def emit(self, record):
msg = self.format(record)
self.log_sink.log(msg)


self.log_sink.log(msg, frame_type=getattr(record, "_frame_type", None))
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),
frame_type=(
getattr(record, "_frame_type", None)
or _TEXT_FRAME_TYPES.get(_format_log_level(record))
),
)


class LambdaLoggerFilter(logging.Filter):
Expand Down Expand Up @@ -416,13 +435,14 @@ 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())

logging.addLevelName(logging.DEBUG, "TRACE")
if log_level in logging._nameToLevel:
logger.setLevel(log_level)
else:
logger_handler.setFormatter(
logging.Formatter(
Expand All @@ -431,6 +451,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)

Expand Down
26 changes: 19 additions & 7 deletions awslambdaric/lambda_runtime_log_utils.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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"),
Expand All @@ -53,12 +57,24 @@ def from_str(cls, value: str):
logging.ERROR: 0xA55A0016.to_bytes(4, "big"),
logging.CRITICAL: 0xA55A001A.to_bytes(4, "big"),
}
_DEFAULT_FRAME_TYPE = 0xA55A0003.to_bytes(4, "big")
_TEXT_FRAME_TYPES = {
logging.NOTSET: 0xA55A0003.to_bytes(4, "big"),
logging.DEBUG: 0xA55A000B.to_bytes(4, "big"),
logging.INFO: 0xA55A000F.to_bytes(4, "big"),
logging.WARNING: 0xA55A0013.to_bytes(4, "big"),
logging.ERROR: 0xA55A0017.to_bytes(4, "big"),
logging.CRITICAL: 0xA55A001B.to_bytes(4, "big"),
}
_DEFAULT_FRAME_TYPE = _TEXT_FRAME_TYPES[logging.NOTSET]

_json_encoder = json.JSONEncoder(ensure_ascii=False)
_encode_json = _json_encoder.encode


def _format_log_level(record: logging.LogRecord) -> int:
return min(50, max(0, record.levelno)) // 10 * 10


class JsonFormatter(logging.Formatter):
def __init__(self):
super().__init__(datefmt=_DATETIME_FORMAT)
Expand Down Expand Up @@ -90,13 +106,9 @@ def __format_location(record: logging.LogRecord):

return f"{record.pathname}:{record.funcName}:{record.lineno}"

@staticmethod
def __format_log_level(record: logging.LogRecord):
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.levelno = _format_log_level(record)
record.levelname = logging.getLevelName(record.levelno)
record._frame_type = _JSON_FRAME_TYPES.get(
record.levelno, _JSON_FRAME_TYPES[logging.NOTSET]
)
Expand Down
83 changes: 69 additions & 14 deletions tests/test_bootstrap.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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


Expand Down Expand Up @@ -927,7 +927,7 @@ def test_log_error_framed_log_sink(self):
content = f.read()

frame_type = int.from_bytes(content[:4], "big")
self.assertEqual(frame_type, 0xA55A0003)
self.assertEqual(frame_type, 0xA55A0017)

length = int.from_bytes(content[4:8], "big")
self.assertEqual(length, len(expected_logged_error.encode("utf8")))
Expand Down Expand Up @@ -973,7 +973,7 @@ def test_log_error_indentation_framed_log_sink(self):
content = f.read()

frame_type = int.from_bytes(content[:4], "big")
self.assertEqual(frame_type, 0xA55A0003)
self.assertEqual(frame_type, 0xA55A0017)

length = int.from_bytes(content[4:8], "big")
self.assertEqual(length, len(expected_logged_error.encode("utf8")))
Expand Down Expand Up @@ -1016,7 +1016,7 @@ def test_log_error_empty_stacktrace_line_framed_log_sink(self):
content = f.read()

frame_type = int.from_bytes(content[:4], "big")
self.assertEqual(frame_type, 0xA55A0003)
self.assertEqual(frame_type, 0xA55A0017)

length = int.from_bytes(content[4:8], "big")
self.assertEqual(length, len(expected_logged_error))
Expand Down Expand Up @@ -1053,7 +1053,7 @@ def test_log_error_invokeId_line_framed_log_sink(self):
content = f.read()

frame_type = int.from_bytes(content[:4], "big")
self.assertEqual(frame_type, 0xA55A0003)
self.assertEqual(frame_type, 0xA55A0017)

length = int.from_bytes(content[4:8], "big")
self.assertEqual(length, len(expected_logged_error))
Expand Down Expand Up @@ -1179,14 +1179,13 @@ def test_log_level(self) -> None:
(LogFormat.JSON, "WARN", logging.WARNING),
(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),
]
Expand All @@ -1196,11 +1195,67 @@ 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)


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:
Expand Down