feat(kubernetes): Add log formatting options and refactor log handling in KubernetesPodOperator

- Rename log_prefix to  container_name_log_prefix_enabled parameter to control container name prefix in logs

- Extract duplicate log formatting logic into _log_message private method in PodManager

- Update docstrings with detailed parameter descriptions

Author: HsiuChuanHsu

kubernetes-tests/tests/kubernetes_tests/test_kubernetes_pod_operator.py

@@ -552,7 +552,7 @@ def test_volume_mount(self):
                 task_id=str(uuid4()),
                 in_cluster=False,
                 do_xcom_push=False,
-                log_prefix=False,
+                container_name_log_prefix_enabled=False,
             )
             context = create_context(k)
             k.execute(context=context)
@@ -1452,15 +1452,11 @@ def test_init_container_logs_filtered(self):
             ),
         ],
     )
-    def test_log_output_configurations(
-        self, mock_get_connection, log_prefix_enabled, log_formatter, expected_log_message_check
-    ):
+    def test_log_output_configurations(self, log_prefix_enabled, log_formatter, expected_log_message_check):
         """
-        Tests various log output configurations (log_prefix, log_formatter)
+        Tests various log output configurations (container_name_log_prefix_enabled, log_formatter)
         for KubernetesPodOperator.
         """
-        from airflow.providers.cncf.kubernetes.utils.pod_manager import PodLoggingStatus
-
         marker = f"test_log_{uuid4()}"
         k = KubernetesPodOperator(
             namespace="default",
@@ -1472,36 +1468,24 @@ def test_log_output_configurations(
             in_cluster=False,
             do_xcom_push=False,
             get_logs=True,
-            log_prefix=log_prefix_enabled,
+            container_name_log_prefix_enabled=log_prefix_enabled,
             log_formatter=log_formatter,
         )
-        context = create_context(k)
+
+        # Test the _log_message method directly
         logger = logging.getLogger("airflow.providers.cncf.kubernetes.utils.pod_manager.PodManager")
         with mock.patch.object(logger, "info") as mock_info:
-            with mock.patch.object(PodManager, "fetch_container_logs") as mock_fetch_logs:
-                # Mock the fetch_container_logs to simulate log output
-                def side_effect(
-                    pod,
-                    container_name,
-                    follow,
-                    since_time=None,
-                    post_termination_timeout=120,
-                    log_prefix=True,
-                    log_formatter=None,
-                ):
-                    log_message = marker
-                    if log_formatter:
-                        log_message = log_formatter(container_name, marker)
-                    else:
-                        log_message = f"[{container_name}] {marker}" if log_prefix else marker
-                    logger.info(log_message)
-                    return PodLoggingStatus(last_log_time=pendulum.now(), running=False)
-
-                mock_fetch_logs.side_effect = side_effect
-                k.execute(context)
+            k.pod_manager._log_message(
+                message=marker,
+                container_name="base",
+                container_name_log_prefix_enabled=log_prefix_enabled,
+                log_formatter=log_formatter,
+            )
 
-            captured_messages = [call_args[0][0] for call_args in mock_info.call_args_list if call_args]
-            assert any(expected_log_message_check(marker, msg) for msg in captured_messages)
+            # Check that the message was logged with the expected format
+            mock_info.assert_called_once()
+            logged_message = mock_info.call_args[0][1]  # Second argument is the message
+            assert expected_log_message_check(marker, logged_message)
 
 
 # TODO: Task SDK: https://github.com/apache/airflow/issues/45438

providers/cncf/kubernetes/src/airflow/providers/cncf/kubernetes/operators/pod.py

@@ -235,6 +235,11 @@ class KubernetesPodOperator(BaseOperator):
         resuming to fetch the latest logs. If ``None``, then the task will remain in deferred state until pod
         is done, and no logs will be visible until that time.
     :param trigger_kwargs: additional keyword parameters passed to the trigger
+    :param container_name_log_prefix_enabled: if True, will prefix container name to each log line.
+        Default to True.
+    :param log_formatter: custom log formatter function that takes two string arguments:
+        the first string is the container_name and the second string is the message_to_log.
+        The function should return a formatted string. If None, the default formatting will be used.
     """
 
     # !!! Changes in KubernetesPodOperator's arguments should be also reflected in !!!
@@ -343,7 +348,7 @@ def __init__(
         progress_callback: Callable[[str], None] | None = None,
         logging_interval: int | None = None,
         trigger_kwargs: dict | None = None,
-        log_prefix: bool = True,
+        container_name_log_prefix_enabled: bool = True,
         log_formatter: Callable[[str, str], str] | None = None,
         **kwargs,
     ) -> None:
@@ -440,7 +445,7 @@ def __init__(
         self._progress_callback = progress_callback
         self.callbacks = [] if not callbacks else callbacks if isinstance(callbacks, list) else [callbacks]
         self._killed: bool = False
-        self.log_prefix = log_prefix
+        self.container_name_log_prefix_enabled = container_name_log_prefix_enabled
         self.log_formatter = log_formatter
 
     @cached_property
@@ -754,7 +759,7 @@ def await_init_containers_completion(self, pod: k8s.V1Pod):
                     pod=pod,
                     init_containers=self.init_container_logs,
                     follow_logs=True,
-                    log_prefix=self.log_prefix,
+                    container_name_log_prefix_enabled=self.container_name_log_prefix_enabled,
                     log_formatter=self.log_formatter,
                 )
         except kubernetes.client.exceptions.ApiException as exc:
@@ -772,7 +777,7 @@ def await_pod_completion(self, pod: k8s.V1Pod):
                     pod=pod,
                     containers=self.container_logs,
                     follow_logs=True,
-                    log_prefix=self.log_prefix,
+                    container_name_log_prefix_enabled=self.container_name_log_prefix_enabled,
                     log_formatter=self.log_formatter,
                 )
             if not self.get_logs or (
@@ -922,7 +927,7 @@ def trigger_reentry(self, context: Context, event: dict[str, Any]) -> Any:
                         container_name=self.base_container_name,
                         follow=follow,
                         since_time=last_log_time,
-                        log_prefix=self.log_prefix,
+                        container_name_log_prefix_enabled=self.container_name_log_prefix_enabled,
                         log_formatter=self.log_formatter,
                     )
 

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

@@ -456,6 +456,26 @@ async def await_pod_start(
 
             await asyncio.sleep(check_interval)
 
+    def _log_message(
+        self,
+        message: str,
+        container_name: str,
+        container_name_log_prefix_enabled: bool,
+        log_formatter: Callable[[str, str], str] | None,
+    ) -> None:
+        """Log a message with appropriate formatting."""
+        if is_log_group_marker(message):
+            print(message)
+        else:
+            if log_formatter:
+                formatted_message = log_formatter(container_name, message)
+                self.log.info("%s", formatted_message)
+            else:
+                log_message = (
+                    f"[{container_name}] {message}" if container_name_log_prefix_enabled else message
+                )
+                self.log.info("%s", log_message)
+
     def fetch_container_logs(
         self,
         pod: V1Pod,
@@ -464,7 +484,7 @@ def fetch_container_logs(
         follow=False,
         since_time: DateTime | None = None,
         post_termination_timeout: int = 120,
-        log_prefix: bool = True,
+        container_name_log_prefix_enabled: bool = True,
         log_formatter: Callable[[str, str], str] | None = None,
     ) -> PodLoggingStatus:
         """
@@ -531,20 +551,12 @@ def consume_logs(*, since_time: DateTime | None = None) -> tuple[DateTime | None
                                             line=line, client=self._client, mode=ExecutionMode.SYNC
                                         )
                                 if message_to_log is not None:
-                                    if is_log_group_marker(message_to_log):
-                                        print(message_to_log)
-                                    else:
-                                        # Apply custom formatter or prefix logic
-                                        if log_formatter:
-                                            formatted_message = log_formatter(container_name, message_to_log)
-                                            self.log.info("%s", formatted_message)
-                                        else:
-                                            log_message = (
-                                                f"[{container_name}] {message_to_log}"
-                                                if log_prefix
-                                                else message_to_log
-                                            )
-                                            self.log.info("%s", log_message)
+                                    self._log_message(
+                                        message_to_log,
+                                        container_name,
+                                        container_name_log_prefix_enabled,
+                                        log_formatter,
+                                    )
                                 last_captured_timestamp = message_timestamp
                                 message_to_log = message
                                 message_timestamp = line_timestamp
@@ -560,17 +572,9 @@ def consume_logs(*, since_time: DateTime | None = None) -> tuple[DateTime | None
                                 line=line, client=self._client, mode=ExecutionMode.SYNC
                             )
                     if message_to_log is not None:
-                        if is_log_group_marker(message_to_log):
-                            print(message_to_log)
-                        else:
-                            if log_formatter:
-                                formatted_message = log_formatter(container_name, message_to_log)
-                                self.log.info("%s", formatted_message)
-                            else:
-                                log_message = (
-                                    f"[{container_name}] {message_to_log}" if log_prefix else message_to_log
-                                )
-                                self.log.info("%s", log_message)
+                        self._log_message(
+                            message_to_log, container_name, container_name_log_prefix_enabled, log_formatter
+                        )
                     last_captured_timestamp = message_timestamp
             except TimeoutError as e:
                 # in case of timeout, increment return time by 2 seconds to avoid
@@ -653,7 +657,7 @@ def fetch_requested_init_container_logs(
         pod: V1Pod,
         init_containers: Iterable[str] | str | Literal[True] | None,
         follow_logs=False,
-        log_prefix: bool = True,
+        container_name_log_prefix_enabled: bool = True,
         log_formatter: Callable[[str, str], str] | None = None,
     ) -> list[PodLoggingStatus]:
         """
@@ -678,7 +682,7 @@ def fetch_requested_init_container_logs(
                 pod=pod,
                 container_name=c,
                 follow=follow_logs,
-                log_prefix=log_prefix,
+                container_name_log_prefix_enabled=container_name_log_prefix_enabled,
                 log_formatter=log_formatter,
             )
             pod_logging_statuses.append(status)
@@ -689,7 +693,7 @@ def fetch_requested_container_logs(
         pod: V1Pod,
         containers: Iterable[str] | str | Literal[True],
         follow_logs=False,
-        log_prefix: bool = True,
+        container_name_log_prefix_enabled: bool = True,
         log_formatter: Callable[[str, str], str] | None = None,
     ) -> list[PodLoggingStatus]:
         """
@@ -711,7 +715,7 @@ def fetch_requested_container_logs(
                 pod=pod,
                 container_name=c,
                 follow=follow_logs,
-                log_prefix=log_prefix,
+                container_name_log_prefix_enabled=container_name_log_prefix_enabled,
                 log_formatter=log_formatter,
             )
             pod_logging_statuses.append(status)

providers/cncf/kubernetes/tests/unit/cncf/kubernetes/operators/test_pod.py

@@ -1734,7 +1734,11 @@ def test_get_logs_but_not_for_base_container(
 
         # check that the base container is not included in the logs
         mock_fetch_log.assert_called_once_with(
-            pod=pod, containers=["some_init_container"], follow_logs=True, log_prefix=True, log_formatter=None
+            pod=pod,
+            containers=["some_init_container"],
+            follow_logs=True,
+            container_name_log_prefix_enabled=True,
+            log_formatter=None,
         )
         # check that KPO waits for the base container to complete before proceeding to extract XCom
         mock_await_container_completion.assert_called_once_with(
@@ -2006,7 +2010,7 @@ def test_await_container_completion_refreshes_properties_on_exception(
                         pod=pod,
                         containers=k.container_logs,
                         follow_logs=True,
-                        log_prefix=True,
+                        container_name_log_prefix_enabled=True,
                         log_formatter=None,
                     )
                 ]

providers/cncf/kubernetes/tests/unit/cncf/kubernetes/operators/test_spark_kubernetes.py

@@ -729,7 +729,7 @@ def test_get_logs_from_driver(
             pod=op.pod,
             containers="spark-kubernetes-driver",
             follow_logs=True,
-            log_prefix=True,
+            container_name_log_prefix_enabled=True,
             log_formatter=None,
         )