Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
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
6 changes: 6 additions & 0 deletions CHANGELOG.md
Original file line number Diff line number Diff line change
Expand Up @@ -12,6 +12,12 @@ Changes are grouped as follows
- `Fixed` for any bug fixes.
- `Security` in case of vulnerabilities.

## Next

### Added

* Added a new command-line argument, --log-level (-l) to allow users to override configured log levels for a single run.

## 7.7.0

### Added
Expand Down
24 changes: 20 additions & 4 deletions cognite/extractorutils/unstable/core/base.py
Original file line number Diff line number Diff line change
Expand Up @@ -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):
Expand Down Expand Up @@ -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}")

Expand All @@ -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
Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Is min and max level for the log should be same in case log level override? Or is it configurable too?

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I think when we use the --log-level override from the command line, the intended behavior should be to force all log handlers to a single, specific level unlike the config where there can be multiple log-handlers and hence multiple log-levels.

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)
Expand All @@ -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)

Expand All @@ -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)
Expand Down
10 changes: 10 additions & 0 deletions cognite/extractorutils/unstable/core/runtime.py
Original file line number Diff line number Diff line change
Expand Up @@ -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",
Expand Down Expand Up @@ -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()
Expand Down
14 changes: 13 additions & 1 deletion tests/test_unstable/conftest.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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))
73 changes: 71 additions & 2 deletions tests/test_unstable/test_base.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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
Loading