diff --git a/datadog_lambda/__init__.py b/datadog_lambda/__init__.py index 20b42443..cfb5f8e3 100644 --- a/datadog_lambda/__init__.py +++ b/datadog_lambda/__init__.py @@ -1,6 +1,5 @@ -import os -import logging from datadog_lambda.cold_start import initialize_cold_start_tracing +from datadog_lambda.logger import initialize_logging initialize_cold_start_tracing() @@ -13,5 +12,4 @@ __version__ = importlib_metadata.version(__name__) -logger = logging.getLogger(__name__) -logger.setLevel(logging.getLevelName(os.environ.get("DD_LOG_LEVEL", "INFO").upper())) +initialize_logging(__name__) diff --git a/datadog_lambda/logger.py b/datadog_lambda/logger.py new file mode 100644 index 00000000..ae2b5d6d --- /dev/null +++ b/datadog_lambda/logger.py @@ -0,0 +1,27 @@ +import logging +import os + +try: + _level_mappping = logging.getLevelNamesMapping() +except AttributeError: + # python 3.8 + _level_mappping = {name: num for num, name in logging._levelToName.items()} +# https://docs.datadoghq.com/agent/troubleshooting/debug_mode/?tab=agentv6v7#agent-log-level +_level_mappping.update( + { + "TRACE": 5, + "WARN": logging.WARNING, + "OFF": 100, + } +) + + +def initialize_logging(name): + logger = logging.getLogger(name) + str_level = (os.environ.get("DD_LOG_LEVEL") or "INFO").upper() + level = _level_mappping.get(str_level) + if level is None: + logger.setLevel(logging.INFO) + logger.warning("Invalid log level: %s Defaulting to INFO", str_level) + else: + logger.setLevel(level) diff --git a/tests/test_logger.py b/tests/test_logger.py new file mode 100644 index 00000000..eb2822f2 --- /dev/null +++ b/tests/test_logger.py @@ -0,0 +1,43 @@ +import io +import logging +import pytest + +from datadog_lambda.logger import initialize_logging + +_test_initialize_logging = ( + ("TRACE", (10, 20, 30, 40, 50)), + ("DEBUG", (10, 20, 30, 40, 50)), + ("debug", (10, 20, 30, 40, 50)), + ("INFO", (20, 30, 40, 50)), + ("WARNING", (30, 40, 50)), + ("WARN", (30, 40, 50)), + ("ERROR", (40, 50)), + ("CRITICAL", (50,)), + ("OFF", ()), + ("", (20, 30, 40, 50)), + (None, (20, 30, 40, 50)), + ("PURPLE", (30, 20, 30, 40, 50)), # log warning then default to INFO +) + + +@pytest.mark.parametrize("level,logged_levels", _test_initialize_logging) +def test_initialize_logging(level, logged_levels, monkeypatch): + if level is not None: + monkeypatch.setenv("DD_LOG_LEVEL", level) + + stream = io.StringIO() + handler = logging.StreamHandler(stream) + handler.setFormatter(logging.Formatter("%(levelno)s")) + logger = logging.getLogger(__name__) + logger.addHandler(handler) + + initialize_logging(__name__) + + logger.debug("debug") + logger.info("info") + logger.warning("warning") + logger.error("error") + logger.critical("critical") + + logged = tuple(map(int, stream.getvalue().strip().split())) + assert logged == logged_levels