From 552295d50fc29e97397f5a404cac33c6af502163 Mon Sep 17 00:00:00 2001 From: Amogh Desai Date: Wed, 21 May 2025 15:50:35 +0530 Subject: [PATCH 1/2] Deserialize should work while retrieving variables with secrects backend --- task-sdk/src/airflow/sdk/execution_time/context.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/task-sdk/src/airflow/sdk/execution_time/context.py b/task-sdk/src/airflow/sdk/execution_time/context.py index 83ad09c90fb36..44f699deb295a 100644 --- a/task-sdk/src/airflow/sdk/execution_time/context.py +++ b/task-sdk/src/airflow/sdk/execution_time/context.py @@ -176,6 +176,10 @@ def _get_variable(key: str, deserialize_json: bool) -> Any: try: var_val = secrets_backend.get_variable(key=key) # type: ignore[assignment] if var_val is not None: + if deserialize_json: + import json + + var_val = json.loads(var_val) return var_val except Exception: log.exception( From 673be197b75f29f1cd4e5a0e71316635375b9cf4 Mon Sep 17 00:00:00 2001 From: Amogh Desai Date: Wed, 21 May 2025 15:53:00 +0530 Subject: [PATCH 2/2] adding unit tests --- .../task_sdk/definitions/test_variables.py | 23 +++++++++++++++++++ 1 file changed, 23 insertions(+) diff --git a/task-sdk/tests/task_sdk/definitions/test_variables.py b/task-sdk/tests/task_sdk/definitions/test_variables.py index 242c5af407bcd..c9e4698de800e 100644 --- a/task-sdk/tests/task_sdk/definitions/test_variables.py +++ b/task-sdk/tests/task_sdk/definitions/test_variables.py @@ -109,6 +109,29 @@ def test_var_get_from_secrets_found(self, mock_supervisor_comms, tmp_path): assert retrieved_var is not None assert retrieved_var == "some_value" + def test_var_get_from_secrets_found_with_deserialize(self, mock_supervisor_comms, tmp_path): + """Tests getting a variable from secrets backend when deserialize_json is provided.""" + path = tmp_path / "var.json" + dict_data = {"num1": 23, "num2": 42} + jsonified_dict_data = json.dumps(dict_data) + data = {"VAR_A": jsonified_dict_data} + path.write_text(json.dumps(data, indent=4)) + + with conf_vars( + { + ( + "workers", + "secrets_backend", + ): "airflow.secrets.local_filesystem.LocalFilesystemBackend", + ("workers", "secrets_backend_kwargs"): f'{{"variables_file_path": "{path}"}}', + } + ): + retrieved_var = Variable.get(key="VAR_A") + assert retrieved_var == jsonified_dict_data + + retrieved_var_deser = Variable.get(key="VAR_A", deserialize_json=True) + assert retrieved_var_deser == dict_data + @mock.patch("airflow.secrets.environment_variables.EnvironmentVariablesBackend.get_variable") def test_get_variable_env_var(self, mock_env_get, mock_supervisor_comms): """Tests getting a variable from environment variable."""