Remove duplicated logs by reusing PodLogsConsumer

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

@@ -392,8 +392,12 @@ def fetch_container_logs(
             before=before_log(self.log, logging.INFO),
         )
         def consume_logs(
-            *, since_time: DateTime | None = None, follow: bool = True, termination_timeout: int = 120
-        ) -> DateTime | None:
+            *,
+            since_time: DateTime | None = None,
+            follow: bool = True,
+            termination_timeout: int = 120,
+            logs: PodLogsConsumer | None,
+        ) -> tuple[DateTime | None, PodLogsConsumer | None]:
             """
             Tries to follow container logs until container completes.
 
@@ -404,16 +408,17 @@ def consume_logs(
             """
             last_captured_timestamp = None
             try:
-                logs = self.read_pod_logs(
-                    pod=pod,
-                    container_name=container_name,
-                    timestamps=True,
-                    since_seconds=(
-                        math.ceil((pendulum.now() - since_time).total_seconds()) if since_time else None
-                    ),
-                    follow=follow,
-                    post_termination_timeout=termination_timeout,
-                )
+                if not logs:
+                    logs = self.read_pod_logs(
+                        pod=pod,
+                        container_name=container_name,
+                        timestamps=True,
+                        since_seconds=(
+                            math.ceil((pendulum.now() - since_time).total_seconds()) if since_time else None
+                        ),
+                        follow=follow,
+                        post_termination_timeout=post_termination_timeout,
+                    )
                 for raw_line in logs:
                     line = raw_line.decode("utf-8", errors="backslashreplace")
                     line_timestamp, message = self.parse_log_line(line)
@@ -434,15 +439,19 @@ def consume_logs(
                     pod.metadata.name,
                     exc_info=True,
                 )
-            return last_captured_timestamp or since_time
+            return last_captured_timestamp or since_time, logs
 
         # note: `read_pod_logs` follows the logs, so we shouldn't necessarily *need* to
         # loop as we do here. But in a long-running process we might temporarily lose connectivity.
         # So the looping logic is there to let us resume following the logs.
+        logs = None
         last_log_time = since_time
         while True:
-            last_log_time = consume_logs(
-                since_time=last_log_time, follow=follow, termination_timeout=post_termination_timeout
+            last_log_time, logs = consume_logs(
+                since_time=last_log_time,
+                follow=follow,
+                termination_timeout=post_termination_timeout,
+                logs=logs,
             )
             if not self.container_is_running(pod, container_name=container_name):
                 return PodLoggingStatus(running=False, last_log_time=last_log_time)

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

@@ -284,6 +284,31 @@ def test_fetch_container_logs_invoke_progress_callback(
         self.pod_manager.fetch_container_logs(mock.MagicMock(), mock.MagicMock(), follow=True)
         self.mock_progress_callback.assert_has_calls([mock.call(message), mock.call(no_ts_message)])
 
+    @mock.patch("airflow.providers.cncf.kubernetes.utils.pod_manager.PodManager.container_is_running")
+    def test_fetch_container_logs_failures(self, mock_container_is_running):
+        last_timestamp_string = "2020-10-08T14:18:17.793417674Z"
+        messages = [
+            bytes("2020-10-08T14:16:17.793417674Z message", "utf-8"),
+            bytes("2020-10-08T14:17:17.793417674Z message", "utf-8"),
+            None,
+            bytes(f"{last_timestamp_string} message", "utf-8"),
+        ]
+        expected_call_count = len([message for message in messages if message is not None])
+
+        def consumer_iter():
+            while messages:
+                message = messages.pop(0)
+                if message is None:
+                    raise BaseHTTPError("Boom")
+                yield message
+
+        with mock.patch.object(PodLogsConsumer, "__iter__") as mock_consumer_iter:
+            mock_consumer_iter.side_effect = consumer_iter
+            mock_container_is_running.side_effect = [True, True, False]
+            status = self.pod_manager.fetch_container_logs(mock.MagicMock(), mock.MagicMock(), follow=True)
+        assert status.last_log_time == cast(DateTime, pendulum.parse(last_timestamp_string))
+        assert self.mock_progress_callback.call_count == expected_call_count
+
     def test_parse_invalid_log_line(self, caplog):
         with caplog.at_level(logging.INFO):
             self.pod_manager.parse_log_line("2020-10-08T14:16:17.793417674ZInvalidmessage\n")