From c8761d01e167e14cbd4d49b6b1e3dad6f1631a88 Mon Sep 17 00:00:00 2001 From: alwaysgaurav1 Date: Mon, 6 Jul 2026 12:08:52 +0530 Subject: [PATCH] fix(task-sdk): raise exception on Variable set/delete failures --- .../src/airflow/sdk/execution_time/context.py | 12 +++++++++--- .../tests/task_sdk/definitions/test_variables.py | 16 ++++++++++++++++ 2 files changed, 25 insertions(+), 3 deletions(-) diff --git a/task-sdk/src/airflow/sdk/execution_time/context.py b/task-sdk/src/airflow/sdk/execution_time/context.py index 7b89164f89560..aaea0fcc52d3b 100644 --- a/task-sdk/src/airflow/sdk/execution_time/context.py +++ b/task-sdk/src/airflow/sdk/execution_time/context.py @@ -377,8 +377,9 @@ def _set_variable(key: str, value: Any, description: str | None = None, serializ # keep Task SDK as a separate package than execution time mods. import json + from airflow.sdk.exceptions import AirflowRuntimeError from airflow.sdk.execution_time.cache import SecretCache - from airflow.sdk.execution_time.comms import PutVariable + from airflow.sdk.execution_time.comms import ErrorResponse, PutVariable from airflow.sdk.execution_time.secrets.execution_api import ExecutionAPISecretsBackend from airflow.sdk.execution_time.supervisor import ensure_secrets_backend_loaded from airflow.sdk.execution_time.task_runner import SUPERVISOR_COMMS @@ -412,7 +413,9 @@ def _set_variable(key: str, value: Any, description: str | None = None, serializ except Exception as e: log.exception(e) - SUPERVISOR_COMMS.send(PutVariable(key=key, value=value, description=description)) + res = SUPERVISOR_COMMS.send(PutVariable(key=key, value=value, description=description)) + if isinstance(res, ErrorResponse): + raise AirflowRuntimeError(res) # Invalidate cache after setting the variable SecretCache.invalidate_variable(key) @@ -424,11 +427,14 @@ def _delete_variable(key: str) -> None: # A reason to not move it to `airflow.sdk.execution_time.comms` is that it # will make that module depend on Task SDK, which is not ideal because we intend to # keep Task SDK as a separate package than execution time mods. + from airflow.sdk.exceptions import AirflowRuntimeError from airflow.sdk.execution_time.cache import SecretCache - from airflow.sdk.execution_time.comms import DeleteVariable + from airflow.sdk.execution_time.comms import DeleteVariable, ErrorResponse from airflow.sdk.execution_time.task_runner import SUPERVISOR_COMMS msg = SUPERVISOR_COMMS.send(DeleteVariable(key=key)) + if isinstance(msg, ErrorResponse): + raise AirflowRuntimeError(msg) if TYPE_CHECKING: assert isinstance(msg, OKResponse) diff --git a/task-sdk/tests/task_sdk/definitions/test_variables.py b/task-sdk/tests/task_sdk/definitions/test_variables.py index 29b6ac0cb97ca..2f9a6c9eb8a1d 100644 --- a/task-sdk/tests/task_sdk/definitions/test_variables.py +++ b/task-sdk/tests/task_sdk/definitions/test_variables.py @@ -102,6 +102,22 @@ def test_var_delete(self, mock_supervisor_comms): mock_supervisor_comms.send.assert_called_once_with(msg=DeleteVariable(key="my_key")) + def test_var_set_raises_on_error_response(self, mock_supervisor_comms): + mock_supervisor_comms.send.return_value = ErrorResponse( + error=ErrorType.GENERIC_ERROR, detail={"message": "boom"} + ) + + with pytest.raises(AirflowRuntimeError): + Variable.set(key="key", value="value") + + def test_var_delete_raises_on_error_response(self, mock_supervisor_comms): + mock_supervisor_comms.send.return_value = ErrorResponse( + error=ErrorType.GENERIC_ERROR, detail={"message": "boom"} + ) + + with pytest.raises(AirflowRuntimeError): + Variable.delete(key="key") + class TestVariableKeys: @pytest.mark.parametrize(