[unstable] Add log level override

cognite/extractorutils/unstable/core/base.py

@@ -97,10 +97,12 @@ def __init__(
         connection_config: ConnectionConfig,
         application_config: _T,
         current_config_revision: ConfigRevision,
+        log_level_override: str | None = None,
     ) -> None:
         self.connection_config = connection_config
         self.application_config = application_config
         self.current_config_revision = current_config_revision
+        self.log_level_override = log_level_override
 
 
 class Extractor(Generic[ConfigType], CogniteLogger):
@@ -132,6 +134,7 @@ def __init__(self, config: FullConfig[ConfigType]) -> None:
         self.connection_config = config.connection_config
         self.application_config = config.application_config
         self.current_config_revision = config.current_config_revision
+        self.log_level_override = config.log_level_override
 
         self.cognite_client = self.connection_config.get_cognite_client(f"{self.EXTERNAL_ID}-{self.VERSION}")
 
@@ -147,8 +150,15 @@ def __init__(self, config: FullConfig[ConfigType]) -> None:
         self.__init_tasks__()
 
     def _setup_logging(self) -> None:
-        min_level = min([_resolve_log_level(h.level.value) for h in self.application_config.log_handlers])
-        max_level = max([_resolve_log_level(h.level.value) for h in self.application_config.log_handlers])
+        if self.log_level_override:
+            # Use the override level if provided
+            level_to_set = _resolve_log_level(self.log_level_override)
+            min_level = level_to_set
+            max_level = level_to_set
+        else:
+            # Otherwise, use the levels from the config file
+            min_level = min([_resolve_log_level(h.level.value) for h in self.application_config.log_handlers])
+            max_level = max([_resolve_log_level(h.level.value) for h in self.application_config.log_handlers])
 
         root = logging.getLogger()
         root.setLevel(min_level)
@@ -173,7 +183,10 @@ def _setup_logging(self) -> None:
                 case LogConsoleHandlerConfig() as console_handler:
                     sh = logging.StreamHandler()
                     sh.setFormatter(fmt)
-                    sh.setLevel(_resolve_log_level(console_handler.level.value))
+                    level_for_handler = (
+                        level_to_set if self.log_level_override else _resolve_log_level(console_handler.level.value)
+                    )
+                    sh.setLevel(level_for_handler)
 
                     root.addHandler(sh)
 
@@ -184,7 +197,10 @@ def _setup_logging(self) -> None:
                         utc=True,
                         backupCount=file_handler.retention,
                     )
-                    fh.setLevel(_resolve_log_level(file_handler.level.value))
+                    level_for_handler = (
+                        level_to_set if self.log_level_override else _resolve_log_level(file_handler.level.value)
+                    )
+                    fh.setLevel(level_for_handler)
                     fh.setFormatter(fmt)
 
                     root.addHandler(fh)

cognite/extractorutils/unstable/core/runtime.py

@@ -113,6 +113,15 @@ def _create_argparser(self) -> ArgumentParser:
             default=None,
             help="Include to use a local application configuration instead of fetching it from CDF",
         )
+        argparser.add_argument(
+            "-l",
+            "--log-level",
+            choices=["debug", "info", "warning", "error", "critical"],
+            type=str,
+            required=False,
+            default="info",
+            help="Set the logging level for the runtime. Default is 'info'.",
+        )
         argparser.add_argument(
             "--skip-init-checks",
             action="store_true",
@@ -334,6 +343,7 @@ def run(self) -> None:
                     connection_config=connection_config,
                     application_config=application_config,
                     current_config_revision=current_config_revision,
+                    log_level_override=args.log_level,
                 )
             )
             process.join()

tests/test_unstable/conftest.py

@@ -16,7 +16,7 @@
     Scopes,
     _ClientCredentialsConfig,
 )
-from cognite.extractorutils.unstable.core.base import Extractor
+from cognite.extractorutils.unstable.core.base import Extractor, StartupTask, TaskContext
 
 
 @pytest.fixture
@@ -106,3 +106,15 @@ class TestExtractor(Extractor[TestConfig]):
     DESCRIPTION = "Test of the new runtime"
     VERSION = "1.0.0"
     CONFIG_TYPE = TestConfig
+
+    def __init_tasks__(self) -> None:
+        """
+        A simple task that runs on startup and logs messages at different levels.
+        """
+
+        def log_messages_task(ctx: TaskContext) -> None:
+            ctx.debug("This is a debug message.")
+            ctx.info("This is an info message.")
+            ctx.warning("This is a warning message.")
+
+        self.add_task(StartupTask(name="log_task", target=log_messages_task))

tests/test_unstable/test_base.py

@@ -2,9 +2,15 @@
 
 import pytest
 
-from cognite.extractorutils.unstable.configuration.models import ConnectionConfig, IntervalConfig, TimeIntervalConfig
+from cognite.extractorutils.unstable.configuration.models import (
+    ConnectionConfig,
+    IntervalConfig,
+    LogConsoleHandlerConfig,
+    LogLevel,
+    TimeIntervalConfig,
+)
 from cognite.extractorutils.unstable.core.base import FullConfig
-from cognite.extractorutils.unstable.core.tasks import ScheduledTask
+from cognite.extractorutils.unstable.core.tasks import ScheduledTask, TaskContext
 from cognite.extractorutils.util import now
 
 from .conftest import MockFunction, TestConfig, TestExtractor
@@ -88,3 +94,66 @@ def test_simple_task_report(
     assert res["items"][0]["errorCount"] == 0
     assert start_time <= res["items"][0]["startTime"] < mid_way
     assert mid_way < res["items"][0]["endTime"] < end_time
+
+
+@pytest.mark.parametrize(
+    "config_level, override_level, expected_logs, unexpected_logs",
+    [
+        (
+            "INFO",
+            None,
+            ["This is an info message.", "This is a warning message."],
+            ["This is a debug message."],
+        ),
+        (
+            "INFO",
+            "DEBUG",
+            ["This is a debug message.", "This is an info message.", "This is a warning message."],
+            [],
+        ),
+        (
+            "INFO",
+            "WARNING",
+            ["This is a warning message."],
+            ["This is a debug message.", "This is an info message."],
+        ),
+    ],
+)
+def test_log_level_override(
+    capsys: pytest.CaptureFixture[str],
+    connection_config: ConnectionConfig,
+    config_level: str,
+    override_level: str | None,
+    expected_logs: list[str],
+    unexpected_logs: list[str],
+) -> None:
+    """
+    Tests that the log level override parameter correctly overrides the log level
+    set in the application configuration.
+    """
+    app_config = TestConfig(
+        parameter_one=1,
+        parameter_two="a",
+        log_handlers=[LogConsoleHandlerConfig(type="console", level=LogLevel(config_level))],
+    )
+
+    full_config = FullConfig(
+        connection_config=connection_config,
+        application_config=app_config,
+        current_config_revision=1,
+        log_level_override=override_level,
+    )
+    extractor = TestExtractor(full_config)
+
+    with extractor:
+        startup_task = next(t for t in extractor._tasks if t.name == "log_task")
+        task_context = TaskContext(task=startup_task, extractor=extractor)
+        startup_task.target(task_context)
+
+    captured = capsys.readouterr()
+    console_output = captured.err
+
+    for log in expected_logs:
+        assert log in console_output
+    for log in unexpected_logs:
+        assert log not in console_output