(airflow) branch v3-0-test updated: [v3-0-test] Add exception to context for task callbacks (#52055) (#52066)

[ka...@apache.org]

This is an automated email from the ASF dual-hosted git repository.

kaxilnaik pushed a commit to branch v3-0-test
in repository https://gitbox.apache.org/repos/asf/airflow.git


The following commit(s) were added to refs/heads/v3-0-test by this push:
     new daa8b8344f2 [v3-0-test] Add exception to context for task callbacks (#52055) (#52066)
daa8b8344f2 is described below

commit daa8b8344f200e314c12d4239e33ff4b9cf75a27
Author: github-actions[bot] <41...@users.noreply.github.com>
AuthorDate: Mon Jun 23 16:59:12 2025 +0530

    [v3-0-test] Add exception to context for task callbacks (#52055) (#52066)
    
    (cherry picked from commit 8d48d53df5736c376abdcdcb8378a2f08d92905b)
    
    Co-authored-by: Tzu-ping Chung <ur...@gmail.com>
---
 task-sdk/src/airflow/sdk/execution_time/task_runner.py     |  3 ++-
 task-sdk/tests/task_sdk/execution_time/test_task_runner.py | 11 +++++++++--
 2 files changed, 11 insertions(+), 3 deletions(-)

diff --git a/task-sdk/src/airflow/sdk/execution_time/task_runner.py b/task-sdk/src/airflow/sdk/execution_time/task_runner.py
index 33583ffe9c7..a515c1a7212 100644
--- a/task-sdk/src/airflow/sdk/execution_time/task_runner.py
+++ b/task-sdk/src/airflow/sdk/execution_time/task_runner.py
@@ -1292,7 +1292,8 @@ def main():
             bundle_name=ti.bundle_instance.name,
             bundle_version=ti.bundle_instance.version,
         ):
-            state, msg, error = run(ti, context, log)
+            state, _, error = run(ti, context, log)
+            context["exception"] = error
             finalize(ti, state, context, log, error)
     except KeyboardInterrupt:
         log = structlog.get_logger(logger_name="task")
diff --git a/task-sdk/tests/task_sdk/execution_time/test_task_runner.py b/task-sdk/tests/task_sdk/execution_time/test_task_runner.py
index 09d4d9709e7..4b47fb57aa5 100644
--- a/task-sdk/tests/task_sdk/execution_time/test_task_runner.py
+++ b/task-sdk/tests/task_sdk/execution_time/test_task_runner.py
@@ -2277,6 +2277,9 @@ class TestTaskRunnerCallsListeners:
 
 @pytest.mark.usefixtures("mock_supervisor_comms")
 class TestTaskRunnerCallsCallbacks:
+    class _Failure(Exception):
+        """Exception raised in a failed execution and received by the failure callback."""
+
     def _execute_success(self, context):
         self.results.append("execute success")
 
@@ -2288,7 +2291,7 @@ class TestTaskRunnerCallsCallbacks:
 
     def _execute_failure(self, context):
         self.results.append("execute failure")
-        raise Exception("sorry!")
+        raise self._Failure("sorry!")
 
     @pytest.mark.parametrize(
         "execute_impl, should_retry, expected_state, expected_results",
@@ -2336,6 +2339,10 @@ class TestTaskRunnerCallsCallbacks:
         def custom_callback(context, *, kind):
             collected_results.append(f"on-{kind} callback")
 
+        def failure_callback(context):
+            custom_callback(context, kind="failure")
+            assert isinstance(context["exception"], self._Failure)
+
         class CustomOperator(BaseOperator):
             results = collected_results
             execute = execute_impl
@@ -2345,7 +2352,7 @@ class TestTaskRunnerCallsCallbacks:
             on_execute_callback=functools.partial(custom_callback, kind="execute"),
             on_skipped_callback=functools.partial(custom_callback, kind="skipped"),
             on_success_callback=functools.partial(custom_callback, kind="success"),
-            on_failure_callback=functools.partial(custom_callback, kind="failure"),
+            on_failure_callback=failure_callback,
             on_retry_callback=functools.partial(custom_callback, kind="retry"),
         )
         runtime_ti = create_runtime_ti(dag_id="dag", task=task, should_retry=should_retry)