diff --git a/airflow-ctl/src/airflowctl/api/operations.py b/airflow-ctl/src/airflowctl/api/operations.py index 00f760134827c..7350b1aa2d657 100644 --- a/airflow-ctl/src/airflowctl/api/operations.py +++ b/airflow-ctl/src/airflowctl/api/operations.py @@ -37,6 +37,7 @@ BulkBodyConnectionBody, BulkBodyPoolBody, BulkBodyVariableBody, + BulkResponse, Config, ConnectionBody, ConnectionCollectionResponse, @@ -672,11 +673,11 @@ def create(self, variable: VariableBody) -> VariableResponse | ServerResponseErr except ServerResponseError as e: raise e - def bulk(self, variables: BulkBodyVariableBody) -> BulkActionResponse | ServerResponseError: + def bulk(self, variables: BulkBodyVariableBody) -> BulkResponse | ServerResponseError: """CRUD multiple variables.""" try: self.response = self.client.patch("variables", json=variables.model_dump()) - return BulkActionResponse.model_validate_json(self.response.content) + return BulkResponse.model_validate_json(self.response.content) except ServerResponseError as e: raise e diff --git a/airflow-ctl/src/airflowctl/ctl/commands/variable_command.py b/airflow-ctl/src/airflowctl/ctl/commands/variable_command.py index 1c31828a372fd..d5ffed31d8c32 100644 --- a/airflow-ctl/src/airflowctl/ctl/commands/variable_command.py +++ b/airflow-ctl/src/airflowctl/ctl/commands/variable_command.py @@ -23,9 +23,9 @@ import rich -from airflow.api_fastapi.core_api.datamodels.common import BulkActionOnExistence from airflowctl.api.client import NEW_API_CLIENT, ClientKind, provide_api_client from airflowctl.api.datamodels.generated import ( + BulkActionOnExistence, BulkBodyVariableBody, BulkCreateActionVariableBody, VariableBody, @@ -33,9 +33,11 @@ @provide_api_client(kind=ClientKind.CLI) -def import_(args, api_client=NEW_API_CLIENT): +def import_(args, api_client=NEW_API_CLIENT) -> list[str]: """Import variables from a given file.""" - success_message = "[green]Import successful! success: {success}, errors: {errors}[/green]" + success_message = "[green]Import successful! success: {success}[/green]" + errors_message = "[red]Import failed! errors: {errors}[/red]" + if not os.path.exists(args.file): rich.print(f"[red]Missing variable file: {args.file}") sys.exit(1) @@ -71,12 +73,16 @@ def import_(args, api_client=NEW_API_CLIENT): ] ) result = api_client.variables.bulk(variables=bulk_body) - rich.print(success_message.format(success=result.success, errors=result.errors)) - return result.success, result.errors + if result.create.errors: + rich.print(errors_message.format(errors=result.create.errors)) + sys.exit(1) + + rich.print(success_message.format(success=result.create.success)) + return result.create.success @provide_api_client(kind=ClientKind.CLI) -def export(args, api_client=NEW_API_CLIENT): +def export(args, api_client=NEW_API_CLIENT) -> None: """Export all the variables to the file.""" success_message = "[green]Export successful! {total_entries} variable(s) to {file}[/green]" var_dict = {} diff --git a/airflow-ctl/tests/airflow_ctl/api/test_operations.py b/airflow-ctl/tests/airflow_ctl/api/test_operations.py index 7cc058c1952c3..a727783bbbaa7 100644 --- a/airflow-ctl/tests/airflow_ctl/api/test_operations.py +++ b/airflow-ctl/tests/airflow_ctl/api/test_operations.py @@ -44,6 +44,7 @@ BulkCreateActionConnectionBody, BulkCreateActionPoolBody, BulkCreateActionVariableBody, + BulkResponse, Config, ConfigOption, ConfigSection, @@ -1054,9 +1055,10 @@ class TestVariablesOperations: ) ] ) - variable_bulk_response = BulkActionResponse( - success=[key], - errors=[], + variable_bulk_response = BulkResponse( + create=BulkActionResponse(success=[key], errors=[]), + update=None, + delete=None, ) def test_get(self): diff --git a/airflow-ctl/tests/airflow_ctl/ctl/commands/test_variable_command.py b/airflow-ctl/tests/airflow_ctl/ctl/commands/test_variable_command.py index 72116b73f58a0..e0c9ce16d8d4f 100644 --- a/airflow-ctl/tests/airflow_ctl/ctl/commands/test_variable_command.py +++ b/airflow-ctl/tests/airflow_ctl/ctl/commands/test_variable_command.py @@ -19,9 +19,12 @@ import json import os +import pytest + from airflowctl.api.client import ClientKind from airflowctl.api.datamodels.generated import ( BulkActionResponse, + BulkResponse, VariableCollectionResponse, VariableResponse, ) @@ -46,12 +49,24 @@ class TestCliVariableCommands: ], total_entries=1, ) - bulk_action_response = BulkActionResponse(success=[key], errors=[]) + bulk_response_success = BulkResponse( + create=BulkActionResponse(success=[key], errors=[]), update=None, delete=None + ) + bulk_response_error = BulkResponse( + create=BulkActionResponse( + success=[], + errors=[ + {"error": f"The variables with these keys: {{'{key}'}} already exist.", "status_code": 409} + ], + ), + update=None, + delete=None, + ) - def test_import(self, api_client_maker, tmp_path, monkeypatch): + def test_import_success(self, api_client_maker, tmp_path, monkeypatch): api_client = api_client_maker( path="/api/v2/variables", - response_json=self.bulk_action_response.model_dump(), + response_json=self.bulk_response_success.model_dump(), expected_http_status_code=200, kind=ClientKind.CLI, ) @@ -67,7 +82,28 @@ def test_import(self, api_client_maker, tmp_path, monkeypatch): self.parser.parse_args(["variables", "import", expected_json_path.as_posix()]), api_client=api_client, ) - assert response == ([self.key], []) + assert response == [self.key] + + def test_import_error(self, api_client_maker, tmp_path, monkeypatch): + api_client = api_client_maker( + path="/api/v2/variables", + response_json=self.bulk_response_error.model_dump(), + expected_http_status_code=200, + kind=ClientKind.CLI, + ) + + monkeypatch.chdir(tmp_path) + expected_json_path = tmp_path / self.export_file_name + variable_file = { + self.key: self.value, + } + + expected_json_path.write_text(json.dumps(variable_file)) + with pytest.raises(SystemExit): + variable_command.import_( + self.parser.parse_args(["variables", "import", expected_json_path.as_posix()]), + api_client=api_client, + ) def test_export(self, api_client_maker, tmp_path, monkeypatch): api_client = api_client_maker(