KubernetesPodTrigger: add exception stack trace in TriggerEvent (#35716)

apache · Dec 16, 2023 · a398d9d · a398d9d
1 parent 63e97ab
commit a398d9d
Show file tree

Hide file tree

Showing 4 changed files with 12 additions and 1 deletion.
diff --git a/airflow/providers/cncf/kubernetes/operators/pod.py b/airflow/providers/cncf/kubernetes/operators/pod.py
@@ -672,6 +672,7 @@ def invoke_defer_method(self):
         )
 
     def execute_complete(self, context: Context, event: dict, **kwargs):
+        self.log.debug("Triggered with event: %s", event)
         pod = None
         try:
             pod = self.hook.get_pod(
@@ -682,7 +683,11 @@ def execute_complete(self, context: Context, event: dict, **kwargs):
                 # fetch some logs when pod is failed
                 if self.get_logs:
                     self.write_logs(pod)
-                raise AirflowException(event["message"])
+                if "stack_trace" in event:
+                    message = f"{event['message']}\n{event['stack_trace']}"
+                else:
+                    message = event["message"]
+                raise AirflowException(message)
             elif event["status"] == "success":
                 # fetch some logs when pod is executed successfully
                 if self.get_logs:

diff --git a/airflow/providers/cncf/kubernetes/triggers/pod.py b/airflow/providers/cncf/kubernetes/triggers/pod.py
@@ -18,6 +18,7 @@
 
 import asyncio
 import datetime
+import traceback
 import warnings
 from asyncio import CancelledError
 from enum import Enum
@@ -231,6 +232,7 @@ async def run(self) -> AsyncIterator[TriggerEvent]:  # type: ignore[override]
                     "namespace": self.pod_namespace,
                     "status": "error",
                     "message": str(e),
+                    "stack_trace": traceback.format_exc(),
                 }
             )
 

diff --git a/tests/providers/cncf/kubernetes/triggers/test_pod.py b/tests/providers/cncf/kubernetes/triggers/test_pod.py
@@ -181,12 +181,14 @@ async def test_logging_in_trigger_when_exception_should_execute_successfully(
 
         generator = trigger.run()
         actual = await generator.asend(None)
+        actual_stack_trace = actual.payload.pop("stack_trace")
         assert (
             TriggerEvent(
                 {"name": POD_NAME, "namespace": NAMESPACE, "status": "error", "message": "Test exception"}
             )
             == actual
         )
+        assert actual_stack_trace.startswith("Traceback (most recent call last):")
 
     @pytest.mark.asyncio
     @mock.patch(f"{TRIGGER_PATH}.define_container_state")

diff --git a/tests/providers/google/cloud/triggers/test_kubernetes_engine.py b/tests/providers/google/cloud/triggers/test_kubernetes_engine.py
@@ -199,12 +199,14 @@ async def test_logging_in_trigger_when_exception_should_execute_successfully(
 
         generator = trigger.run()
         actual = await generator.asend(None)
+        actual_stack_trace = actual.payload.pop("stack_trace")
         assert (
             TriggerEvent(
                 {"name": POD_NAME, "namespace": NAMESPACE, "status": "error", "message": "Test exception"}
             )
             == actual
         )
+        assert actual_stack_trace.startswith("Traceback (most recent call last):")
 
     @pytest.mark.asyncio
     @mock.patch(f"{TRIGGER_KUB_PATH}.define_container_state")