From e5390fdd69eb78033be2229046d1b6616d51f644 Mon Sep 17 00:00:00 2001 From: Jayachandra Kasarla Date: Sun, 14 Jun 2026 22:44:25 +0530 Subject: [PATCH 1/3] modified variable set and delete methods to raise exceptions, and added unit tests --- .../src/airflow/sdk/definitions/variable.py | 2 ++ .../task_sdk/definitions/test_variables.py | 35 ++++++++++++++++++- 2 files changed, 36 insertions(+), 1 deletion(-) diff --git a/task-sdk/src/airflow/sdk/definitions/variable.py b/task-sdk/src/airflow/sdk/definitions/variable.py index a379022a90918..4d2d3b17e8440 100644 --- a/task-sdk/src/airflow/sdk/definitions/variable.py +++ b/task-sdk/src/airflow/sdk/definitions/variable.py @@ -67,6 +67,7 @@ def set(cls, key: str, value: Any, description: str | None = None, serialize_jso return _set_variable(key, value, description, serialize_json=serialize_json) except AirflowRuntimeError as e: log.exception(e) + raise @classmethod def keys(cls, prefix: str | None = None) -> Sequence[str]: @@ -101,3 +102,4 @@ def delete(cls, key: str) -> None: _delete_variable(key=key) except AirflowRuntimeError as e: log.exception(e) + raise diff --git a/task-sdk/tests/task_sdk/definitions/test_variables.py b/task-sdk/tests/task_sdk/definitions/test_variables.py index 6e94ccf503f8c..2f9d3d9e4d6e6 100644 --- a/task-sdk/tests/task_sdk/definitions/test_variables.py +++ b/task-sdk/tests/task_sdk/definitions/test_variables.py @@ -25,7 +25,15 @@ from airflow.sdk import Variable from airflow.sdk.configuration import initialize_secrets_backends -from airflow.sdk.execution_time.comms import GetVariableKeys, PutVariable, VariableKeysResult, VariableResult +from airflow.sdk.exceptions import AirflowRuntimeError, ErrorType +from airflow.sdk.execution_time.comms import ( + DeleteVariable, + ErrorResponse, + GetVariableKeys, + PutVariable, + VariableKeysResult, + VariableResult, +) from airflow.sdk.execution_time.secrets import DEFAULT_SECRETS_SEARCH_PATH_WORKERS from tests_common.test_utils.config import conf_vars @@ -89,6 +97,31 @@ def test_var_set(self, key, value, description, serialize_json, mock_supervisor_ ), ) + def test_var_set_raises_on_runtime_error(self, mock_supervisor_comms): + with mock.patch( + "airflow.sdk.execution_time.context._set_variable", + side_effect=AirflowRuntimeError( + ErrorResponse(error=ErrorType.GENERIC_ERROR, detail={"message": "API_SERVER_ERROR"}) + ), + ): + with pytest.raises(AirflowRuntimeError): + Variable.set(key="my_key", value="my_value") + + def test_var_delete(self, mock_supervisor_comms): + Variable.delete(key="my_key") + + mock_supervisor_comms.send.assert_called_once_with(msg=DeleteVariable(key="my_key")) + + def test_var_delete_raises_on_runtime_error(self, mock_supervisor_comms): + with mock.patch( + "airflow.sdk.execution_time.context._delete_variable", + side_effect=AirflowRuntimeError( + ErrorResponse(error=ErrorType.GENERIC_ERROR, detail={"message": "API_SERVER_ERROR"}) + ), + ): + with pytest.raises(AirflowRuntimeError): + Variable.delete(key="my_key") + class TestVariableKeys: @pytest.mark.parametrize( From 19631f5f2ab5e3996f2b754d1de7c30ad0b8c2fe Mon Sep 17 00:00:00 2001 From: Jayachandra Kasarla Date: Wed, 17 Jun 2026 13:08:22 +0530 Subject: [PATCH 2/3] removed try-except blocks to skip logging the error messages twice --- task-sdk/src/airflow/sdk/definitions/variable.py | 12 ++---------- 1 file changed, 2 insertions(+), 10 deletions(-) diff --git a/task-sdk/src/airflow/sdk/definitions/variable.py b/task-sdk/src/airflow/sdk/definitions/variable.py index 4d2d3b17e8440..9e769c350c3ee 100644 --- a/task-sdk/src/airflow/sdk/definitions/variable.py +++ b/task-sdk/src/airflow/sdk/definitions/variable.py @@ -63,11 +63,7 @@ def set(cls, key: str, value: Any, description: str | None = None, serialize_jso from airflow.sdk.exceptions import AirflowRuntimeError from airflow.sdk.execution_time.context import _set_variable - try: - return _set_variable(key, value, description, serialize_json=serialize_json) - except AirflowRuntimeError as e: - log.exception(e) - raise + _set_variable(key, value, description, serialize_json=serialize_json) @classmethod def keys(cls, prefix: str | None = None) -> Sequence[str]: @@ -98,8 +94,4 @@ def delete(cls, key: str) -> None: from airflow.sdk.exceptions import AirflowRuntimeError from airflow.sdk.execution_time.context import _delete_variable - try: - _delete_variable(key=key) - except AirflowRuntimeError as e: - log.exception(e) - raise + _delete_variable(key=key) From dc85480653cf071dd3bf766c49fde6c4b390d112 Mon Sep 17 00:00:00 2001 From: Jayachandra Kasarla Date: Wed, 17 Jun 2026 13:14:50 +0530 Subject: [PATCH 3/3] removed redundant imports and tests --- .../src/airflow/sdk/definitions/variable.py | 2 -- .../task_sdk/definitions/test_variables.py | 22 ------------------- 2 files changed, 24 deletions(-) diff --git a/task-sdk/src/airflow/sdk/definitions/variable.py b/task-sdk/src/airflow/sdk/definitions/variable.py index 9e769c350c3ee..f1db76c3fae60 100644 --- a/task-sdk/src/airflow/sdk/definitions/variable.py +++ b/task-sdk/src/airflow/sdk/definitions/variable.py @@ -60,7 +60,6 @@ def get(cls, key: str, default: Any = NOTSET, deserialize_json: bool = False): @classmethod def set(cls, key: str, value: Any, description: str | None = None, serialize_json: bool = False) -> None: - from airflow.sdk.exceptions import AirflowRuntimeError from airflow.sdk.execution_time.context import _set_variable _set_variable(key, value, description, serialize_json=serialize_json) @@ -91,7 +90,6 @@ def keys(cls, prefix: str | None = None) -> Sequence[str]: @classmethod def delete(cls, key: str) -> None: - from airflow.sdk.exceptions import AirflowRuntimeError from airflow.sdk.execution_time.context import _delete_variable _delete_variable(key=key) diff --git a/task-sdk/tests/task_sdk/definitions/test_variables.py b/task-sdk/tests/task_sdk/definitions/test_variables.py index 2f9d3d9e4d6e6..29b6ac0cb97ca 100644 --- a/task-sdk/tests/task_sdk/definitions/test_variables.py +++ b/task-sdk/tests/task_sdk/definitions/test_variables.py @@ -97,31 +97,11 @@ def test_var_set(self, key, value, description, serialize_json, mock_supervisor_ ), ) - def test_var_set_raises_on_runtime_error(self, mock_supervisor_comms): - with mock.patch( - "airflow.sdk.execution_time.context._set_variable", - side_effect=AirflowRuntimeError( - ErrorResponse(error=ErrorType.GENERIC_ERROR, detail={"message": "API_SERVER_ERROR"}) - ), - ): - with pytest.raises(AirflowRuntimeError): - Variable.set(key="my_key", value="my_value") - def test_var_delete(self, mock_supervisor_comms): Variable.delete(key="my_key") mock_supervisor_comms.send.assert_called_once_with(msg=DeleteVariable(key="my_key")) - def test_var_delete_raises_on_runtime_error(self, mock_supervisor_comms): - with mock.patch( - "airflow.sdk.execution_time.context._delete_variable", - side_effect=AirflowRuntimeError( - ErrorResponse(error=ErrorType.GENERIC_ERROR, detail={"message": "API_SERVER_ERROR"}) - ), - ): - with pytest.raises(AirflowRuntimeError): - Variable.delete(key="my_key") - class TestVariableKeys: @pytest.mark.parametrize( @@ -204,8 +184,6 @@ def test_keys_paginates_when_results_exceed_page_size(self, mock_supervisor_comm ) def test_keys_raises_on_error_response(self, mock_supervisor_comms): - from airflow.sdk.exceptions import AirflowRuntimeError, ErrorType - from airflow.sdk.execution_time.comms import ErrorResponse mock_supervisor_comms.send.return_value = ErrorResponse( error=ErrorType.GENERIC_ERROR, detail={"message": "boom"}