KPO Add follow log in termination step

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

@@ -21,6 +21,7 @@
 import datetime
 import json
 import logging
+import math
 import re
 import shlex
 import string
@@ -710,10 +711,13 @@ def trigger_reentry(self, context: Context, event: dict[str, Any]) -> Any:
                     pod=self.pod, event=event, client=self.client, mode=ExecutionMode.SYNC
                 )
 
+            follow = self.logging_interval is None
+            last_log_time = event.get("last_log_time")
+
             if event["status"] in ("error", "failed", "timeout"):
                 # fetch some logs when pod is failed
                 if self.get_logs:
-                    self.write_logs(self.pod)
+                    self.write_logs(self.pod, follow=follow, since_time=last_log_time)
 
                 if self.do_xcom_push:
                     _ = self.extract_xcom(pod=self.pod)
@@ -723,13 +727,12 @@ def trigger_reentry(self, context: Context, event: dict[str, Any]) -> Any:
 
             elif event["status"] == "running":
                 if self.get_logs:
-                    last_log_time = event.get("last_log_time")
                     self.log.info("Resuming logs read from time %r", last_log_time)
 
                     pod_log_status = self.pod_manager.fetch_container_logs(
                         pod=self.pod,
                         container_name=self.BASE_CONTAINER_NAME,
-                        follow=self.logging_interval is None,
+                        follow=follow,
                         since_time=last_log_time,
                     )
 
@@ -742,7 +745,7 @@ def trigger_reentry(self, context: Context, event: dict[str, Any]) -> Any:
             elif event["status"] == "success":
                 # fetch some logs when pod is executed successfully
                 if self.get_logs:
-                    self.write_logs(self.pod)
+                    self.write_logs(self.pod, follow=follow, since_time=last_log_time)
 
                 if self.do_xcom_push:
                     xcom_sidecar_output = self.extract_xcom(pod=self.pod)
@@ -773,12 +776,18 @@ def _clean(self, event: dict[str, Any]):
     def execute_complete(self, context: Context, event: dict, **kwargs):
         self.trigger_reentry(context=context, event=event)
 
-    def write_logs(self, pod: k8s.V1Pod):
+    def write_logs(self, pod: k8s.V1Pod, follow: bool = False, since_time: DateTime | None = None):
         try:
+            since_seconds = (
+                math.ceil((datetime.datetime.now(tz=datetime.timezone.utc) - since_time).total_seconds())
+                if since_time
+                else None
+            )
             logs = self.pod_manager.read_pod_logs(
                 pod=pod,
                 container_name=self.base_container_name,
-                follow=False,
+                follow=follow,
+                since_seconds=since_seconds,
             )
             for raw_line in logs:
                 line = raw_line.decode("utf-8", errors="backslashreplace").rstrip("\n")

airflow/providers/cncf/kubernetes/triggers/pod.py

@@ -243,7 +243,12 @@ async def _wait_for_container_completion(self) -> TriggerEvent:
             container_state = self.define_container_state(pod)
             if container_state == ContainerState.TERMINATED:
                 return TriggerEvent(
-                    {"status": "success", "namespace": self.pod_namespace, "name": self.pod_name}
+                    {
+                        "status": "success",
+                        "namespace": self.pod_namespace,
+                        "name": self.pod_name,
+                        "last_log_time": self.last_log_time,
+                    }
                 )
             elif container_state == ContainerState.FAILED:
                 return TriggerEvent(
@@ -252,6 +257,7 @@ async def _wait_for_container_completion(self) -> TriggerEvent:
                         "namespace": self.pod_namespace,
                         "name": self.pod_name,
                         "message": "Container state failed",
+                        "last_log_time": self.last_log_time,
                     }
                 )
             if time_get_more_logs and datetime.datetime.now(tz=datetime.timezone.utc) > time_get_more_logs:

tests/providers/cncf/kubernetes/triggers/test_pod.py

@@ -187,6 +187,7 @@ async def test_run_loop_return_failed_event(self, mock_hook, mock_method, mock_w
                 "namespace": "default",
                 "name": "test-pod-name",
                 "message": "Container state failed",
+                "last_log_time": None,
             }
         )
         actual_event = await trigger.run().asend(None)