Skip to content

Commit 277f01c

Browse files
Vishwaspatel2401Vishwaspatel2401
authored
Fix KubernetesPodOperator to forward pod log levels to Airflow task logs (#64829)
* Fix KubernetesPodOperator to forward pod log levels to Airflow task logs Previously all pod container output was logged at INFO level regardless of the original log level. Now ERROR, WARNING, DEBUG, CRITICAL prefixes are detected and forwarded at the correct Python logging level. closes: #64820 * Fix _log_message to use log.info() for INFO level messages Use self.log.info() for INFO-level messages and self.log.log(level, ...) only for non-INFO levels, restoring backward compatibility with existing tests that mock log.info directly. * Address review feedback: simplify log dispatch and add level detection to async path - Consolidate if/else in _log_message to a single self.log.log(level, msg) call - Apply _parse_log_level() in AsyncPodManager.fetch_container_logs_before_current_sec() so deferrable/async task logs are also forwarded at the correct level instead of always INFO - Update test to assert self.log.log() instead of self.log.info() for async path * Fix test assertions for log level detection - Add missing assert_called_once() in test_log_message_defaults_to_info_for_plain_lines for consistency with test_log_message_uses_detected_log_level - Use _parse_log_level() to derive expected level in async test assertions instead of hardcoding logging.INFO, matching what production code does * Fix system test assertions broken by log.log() change test_volume_mount and test_log_output_configurations were asserting log.info() calls, but _log_message now always uses self.log.log(level, ...) Update both tests to assert log.log() calls instead.
1 parent 6c598e2 commit 277f01c

3 files changed

Lines changed: 90 additions & 14 deletions

File tree

‎kubernetes-tests/tests/kubernetes_tests/test_kubernetes_pod_operator.py‎

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -554,7 +554,7 @@ def test_volume_mount(self):
554554
)
555555
context = create_context(k)
556556
k.execute(context=context)
557-
mock_logger.info.assert_any_call("%s", StringContainingId("retrieved from mount"))
557+
mock_logger.log.assert_any_call(logging.INFO, StringContainingId("retrieved from mount"))
558558
actual_pod = self.api_client.sanitize_for_serialization(k.pod)
559559
self.expected_pod["spec"]["containers"][0]["args"] = args
560560
self.expected_pod["spec"]["containers"][0]["volumeMounts"] = [
@@ -1480,7 +1480,7 @@ def test_log_output_configurations(self, log_prefix_enabled, log_formatter, expe
14801480
)
14811481

14821482
# Test the _log_message method directly
1483-
with mock.patch.object(k.pod_manager.log, "info") as mock_info:
1483+
with mock.patch.object(k.pod_manager.log, "log") as mock_log:
14841484
k.pod_manager._log_message(
14851485
message=marker,
14861486
container_name="base",
@@ -1489,8 +1489,8 @@ def test_log_output_configurations(self, log_prefix_enabled, log_formatter, expe
14891489
)
14901490

14911491
# Check that the message was logged with the expected format
1492-
mock_info.assert_called_once()
1493-
logged_message = mock_info.call_args[0][1] # Second argument is the message
1492+
mock_log.assert_called_once()
1493+
logged_message = mock_log.call_args[0][1] # Second argument is the message
14941494
assert expected_log_message_check(marker, logged_message)
14951495

14961496

‎providers/cncf/kubernetes/src/airflow/providers/cncf/kubernetes/utils/pod_manager.py‎

Lines changed: 37 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -21,7 +21,9 @@
2121
import asyncio
2222
import enum
2323
import json
24+
import logging
2425
import math
26+
import re
2527
import time
2628
from collections.abc import Callable, Generator, Iterable
2729
from contextlib import closing
@@ -74,6 +76,33 @@
7476
:meta private:
7577
"""
7678

79+
_POD_LOG_LEVEL_PATTERN = re.compile(
80+
r"^\s*(?:\[)?(DEBUG|INFO|WARNING|WARN|ERROR|CRITICAL|FATAL)(?:\])?\s*[:\-]?\s*",
81+
re.IGNORECASE,
82+
)
83+
_POD_LOG_LEVEL_MAP: dict[str, int] = {
84+
"DEBUG": logging.DEBUG,
85+
"INFO": logging.INFO,
86+
"WARNING": logging.WARNING,
87+
"WARN": logging.WARNING,
88+
"ERROR": logging.ERROR,
89+
"CRITICAL": logging.CRITICAL,
90+
"FATAL": logging.CRITICAL,
91+
}
92+
93+
94+
def _parse_log_level(message: str) -> int:
95+
"""
96+
Detect the Python logging level from a pod log line's prefix.
97+
98+
Recognises common formats: ``ERROR:``, ``[ERROR]``, ``WARNING -``, etc.
99+
Returns ``logging.INFO`` when no known prefix is found (backwards-compatible).
100+
"""
101+
match = _POD_LOG_LEVEL_PATTERN.match(message)
102+
if match:
103+
return _POD_LOG_LEVEL_MAP.get(match.group(1).upper(), logging.INFO)
104+
return logging.INFO
105+
77106

78107
class XComRetrievalError(AirflowException):
79108
"""When not possible to get xcom."""
@@ -436,18 +465,18 @@ def _log_message(
436465
container_name_log_prefix_enabled: bool,
437466
log_formatter: Callable[[str, str], str] | None,
438467
) -> None:
439-
"""Log a message with appropriate formatting."""
468+
"""Log a message at the level detected from its prefix, with appropriate formatting."""
440469
if is_log_group_marker(message):
441470
print(message)
442471
else:
472+
level = _parse_log_level(message)
443473
if log_formatter:
444474
formatted_message = log_formatter(container_name, message)
445-
self.log.info("%s", formatted_message)
446475
else:
447-
log_message = (
476+
formatted_message = (
448477
f"[{container_name}] {message}" if container_name_log_prefix_enabled else message
449478
)
450-
self.log.info("%s", log_message)
479+
self.log.log(level, formatted_message)
451480

452481
def fetch_container_logs(
453482
self,
@@ -1180,7 +1209,8 @@ async def fetch_container_logs_before_current_sec(
11801209
if is_log_group_marker(message_to_log):
11811210
print(message_to_log)
11821211
else:
1183-
self.log.info("[%s] %s", container_name, message_to_log)
1212+
level = _parse_log_level(message_to_log)
1213+
self.log.log(level, "[%s] %s", container_name, message_to_log)
11841214
message_to_log = message
11851215
elif message_to_log: # continuation of the previous log line
11861216
message_to_log = f"{message_to_log}\n{message}"
@@ -1190,5 +1220,6 @@ async def fetch_container_logs_before_current_sec(
11901220
if is_log_group_marker(message_to_log):
11911221
print(message_to_log)
11921222
else:
1193-
self.log.info("[%s] %s", container_name, message_to_log)
1223+
level = _parse_log_level(message_to_log)
1224+
self.log.log(level, "[%s] %s", container_name, message_to_log)
11941225
return now # Return the current time as the last log time to ensure logs from the current second are read in the next fetch.

‎providers/cncf/kubernetes/tests/unit/cncf/kubernetes/utils/test_pod_manager.py‎

Lines changed: 49 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -36,6 +36,7 @@
3636
PodManager,
3737
PodPhase,
3838
XComRetrievalError,
39+
_parse_log_level,
3940
log_pod_event,
4041
parse_log_line,
4142
)
@@ -153,6 +154,48 @@ def setup_method(self):
153154
callbacks=[MockKubernetesPodOperatorCallback],
154155
)
155156

157+
@pytest.mark.parametrize(
158+
("message", "expected_level"),
159+
[
160+
("ERROR: something went wrong", logging.ERROR),
161+
("WARNING: low disk space", logging.WARNING),
162+
("WARN: deprecated usage", logging.WARNING),
163+
("DEBUG: entering function", logging.DEBUG),
164+
("CRITICAL: system failure", logging.CRITICAL),
165+
("FATAL: unrecoverable error", logging.CRITICAL),
166+
("INFO: starting up", logging.INFO),
167+
("[ERROR] bracketed prefix", logging.ERROR),
168+
("plain log line with no level", logging.INFO),
169+
("", logging.INFO),
170+
],
171+
)
172+
def test_parse_log_level(self, message, expected_level):
173+
assert _parse_log_level(message) == expected_level
174+
175+
def test_log_message_uses_detected_log_level(self):
176+
"""_log_message should forward ERROR lines at ERROR level, not INFO."""
177+
with mock.patch.object(self.pod_manager.log, "log") as mock_log:
178+
self.pod_manager._log_message(
179+
message="ERROR: something failed",
180+
container_name="base",
181+
container_name_log_prefix_enabled=True,
182+
log_formatter=None,
183+
)
184+
mock_log.assert_called_once()
185+
assert mock_log.call_args[0][0] == logging.ERROR
186+
187+
def test_log_message_defaults_to_info_for_plain_lines(self):
188+
"""_log_message should use INFO for lines without a known level prefix."""
189+
with mock.patch.object(self.pod_manager.log, "log") as mock_log:
190+
self.pod_manager._log_message(
191+
message="Starting application",
192+
container_name="base",
193+
container_name_log_prefix_enabled=True,
194+
log_formatter=None,
195+
)
196+
mock_log.assert_called_once()
197+
assert mock_log.call_args[0][0] == logging.INFO
198+
156199
def test_read_pod_logs_successfully_returns_logs(self):
157200
mock.sentinel.metadata = mock.MagicMock()
158201
self.mock_kube_client.read_namespaced_pod_log.return_value = mock.sentinel.logs
@@ -1503,17 +1546,19 @@ async def test_fetch_container_logs_before_current_sec_various_logs(
15031546
async_hook=mock_async_hook,
15041547
callbacks=[],
15051548
)
1506-
with mock.patch.object(async_pod_manager.log, "info") as mock_log_info:
1549+
with mock.patch.object(async_pod_manager.log, "log") as mock_log:
15071550
result = await async_pod_manager.fetch_container_logs_before_current_sec(
15081551
pod=pod, container_name=container_name, since_time=since_time
15091552
)
15101553
assert result == now
15111554

15121555
for expected in expected_log_messages:
1513-
mock_log_info.assert_any_call("[%s] %s", container_name, expected)
1556+
mock_log.assert_any_call(_parse_log_level(expected), "[%s] %s", container_name, expected)
15141557
for not_expected in not_expected_log_messages:
1515-
unexpected_call = mock.call("[%s] %s", container_name, not_expected)
1516-
assert unexpected_call not in mock_log_info.mock_calls
1558+
unexpected_call = mock.call(
1559+
_parse_log_level(not_expected), "[%s] %s", container_name, not_expected
1560+
)
1561+
assert unexpected_call not in mock_log.mock_calls
15171562

15181563
@pytest.mark.asyncio
15191564
async def test_fetch_container_logs_before_current_sec_error_handling(self):

0 commit comments

Comments
 (0)