diff --git a/airflow-core/src/airflow/cli/commands/variable_command.py b/airflow-core/src/airflow/cli/commands/variable_command.py index 5216aabb446c6..2c4dff95fcd1f 100644 --- a/airflow-core/src/airflow/cli/commands/variable_command.py +++ b/airflow-core/src/airflow/cli/commands/variable_command.py @@ -23,10 +23,12 @@ import os from typing import TYPE_CHECKING +from airflowctl.api.operations import ServerResponseError from sqlalchemy import select +from airflow.cli.api_client import NEW_API_CLIENT, Client, provide_api_client from airflow.cli.simple_table import AirflowConsole -from airflow.cli.utils import SENSITIVE_PLACEHOLDER, print_export_output +from airflow.cli.utils import SENSITIVE_PLACEHOLDER, deprecated_for_airflowctl, print_export_output from airflow.exceptions import ( AirflowFileParseException, AirflowUnsupportedFileTypeException, @@ -92,19 +94,26 @@ def _mapper(var): AirflowConsole().print_as(data=variables, output=args.output, mapper=None) +@deprecated_for_airflowctl("airflowctl variables get") @suppress_logs_and_warning @providers_configuration_loaded -def variables_get(args): +@provide_api_client +def variables_get(args, api_client: Client = NEW_API_CLIENT): """Display variable by a given name.""" try: - if args.default is None: - var = Variable.get(args.key, deserialize_json=args.json) - print(var) - else: - var = Variable.get(args.key, deserialize_json=args.json, default_var=args.default) - print(var) - except (ValueError, KeyError) as e: - raise SystemExit(str(e).strip("'\"")) + variable = api_client.variables.get(variable_key=args.key) + except ServerResponseError as e: + if e.response.status_code == 404: + if args.default is not None: + print(args.default) + return + raise SystemExit(f"Variable {args.key} does not exist") + raise + value = variable.value + if args.json and value is not None: + print(json.loads(value)) + else: + print(value) @cli_utils.action_cli diff --git a/airflow-core/tests/unit/cli/commands/test_command_deprecations.py b/airflow-core/tests/unit/cli/commands/test_command_deprecations.py index b4eb6840c9069..291183b61794f 100644 --- a/airflow-core/tests/unit/cli/commands/test_command_deprecations.py +++ b/airflow-core/tests/unit/cli/commands/test_command_deprecations.py @@ -30,7 +30,7 @@ import pytest -from airflow.cli.commands import asset_command, dag_command, pool_command +from airflow.cli.commands import asset_command, dag_command, pool_command, variable_command from airflow.exceptions import RemovedInAirflow4Warning # (command callable, argv to parse, expected airflowctl replacement named in the warning) @@ -52,6 +52,7 @@ ["assets", "materialize", "--name=foo"], "airflowctl assets materialize", ), + (variable_command.variables_get, ["variables", "get", "foo"], "airflowctl variables get"), ] diff --git a/airflow-core/tests/unit/cli/commands/test_variable_command.py b/airflow-core/tests/unit/cli/commands/test_variable_command.py index a02c95aa31722..27dea7bae0970 100644 --- a/airflow-core/tests/unit/cli/commands/test_variable_command.py +++ b/airflow-core/tests/unit/cli/commands/test_variable_command.py @@ -22,8 +22,11 @@ from contextlib import redirect_stdout from io import StringIO +import httpx import pytest import yaml +from airflowctl.api.datamodels.generated import VariableResponse +from airflowctl.api.operations import ServerResponseError from sqlalchemy import select from airflow import models @@ -37,6 +40,12 @@ pytestmark = pytest.mark.db_test +def _server_error(status_code: int) -> ServerResponseError: + request = httpx.Request("GET", "http://testserver/api/v2/variables/foo") + response = httpx.Response(status_code, request=request, json={"detail": "boom"}) + return ServerResponseError(message="boom", request=request, response=response) + + # Test data fixtures @pytest.fixture def simple_variable_data(): @@ -152,24 +161,44 @@ def test_variables_set_with_description(self): with pytest.raises(KeyError): Variable.get("foo1") - def test_variables_get(self, stdout_capture): - Variable.set("foo", {"foo": "bar"}, serialize_json=True) + def test_variables_get(self, mock_cli_api_client, stdout_capture): + mock_cli_api_client.variables.get.return_value = VariableResponse( + key="foo", value='{\n "foo": "bar"\n}', is_encrypted=False + ) with stdout_capture as stdout: variable_command.variables_get(self.parser.parse_args(["variables", "get", "foo"])) assert stdout.getvalue() == '{\n "foo": "bar"\n}\n' + mock_cli_api_client.variables.get.assert_called_once_with(variable_key="foo") + + def test_variables_get_json_deserializes_value(self, mock_cli_api_client, stdout_capture): + mock_cli_api_client.variables.get.return_value = VariableResponse( + key="foo", value='{"foo": "bar"}', is_encrypted=False + ) + + with stdout_capture as stdout: + variable_command.variables_get(self.parser.parse_args(["variables", "get", "foo", "--json"])) + assert stdout.getvalue() == "{'foo': 'bar'}\n" + + def test_get_variable_default_value(self, mock_cli_api_client, stdout_capture): + mock_cli_api_client.variables.get.side_effect = _server_error(404) - def test_get_variable_default_value(self, stdout_capture): with stdout_capture as stdout: variable_command.variables_get( self.parser.parse_args(["variables", "get", "baz", "--default", "bar"]) ) assert stdout.getvalue() == "bar\n" - def test_get_variable_missing_variable(self): + def test_get_variable_missing_variable(self, mock_cli_api_client): + mock_cli_api_client.variables.get.side_effect = _server_error(404) with pytest.raises(SystemExit): variable_command.variables_get(self.parser.parse_args(["variables", "get", "no-existing-VAR"])) + def test_get_variable_other_error_reraised(self, mock_cli_api_client): + mock_cli_api_client.variables.get.side_effect = _server_error(500) + with pytest.raises(ServerResponseError): + variable_command.variables_get(self.parser.parse_args(["variables", "get", "foo"])) + def test_variables_set_different_types(self): """Test storage of various data types""" # Set a dict