[AIP-94] airflowctl variables: add set command

airflow-core/src/airflow/cli/commands/variable_command.py

@@ -23,10 +23,17 @@
 import os
 from typing import TYPE_CHECKING
 
+from airflowctl.api.datamodels.generated import (
+    BulkActionOnExistence,
+    BulkBodyVariableBody,
+    BulkCreateActionVariableBody,
+    VariableBody,
+)
 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,
@@ -108,10 +115,25 @@ def variables_get(args):
 
 
 @cli_utils.action_cli
+@deprecated_for_airflowctl("airflowctl variables set")
+@suppress_logs_and_warning
 @providers_configuration_loaded
-def variables_set(args):
-    """Create new variable with a given name, value and description."""
-    Variable.set(args.key, args.value, args.description, serialize_json=args.json)
+@provide_api_client
+def variables_set(args, api_client: Client = NEW_API_CLIENT):
+    """Set a variable, creating it if it does not exist and updating it otherwise."""
+    value = args.value
+    if args.json:
+        value = json.dumps(value, indent=2)
+    bulk_body = BulkBodyVariableBody(
+        actions=[
+            BulkCreateActionVariableBody(
+                action="create",
+                entities=[VariableBody(key=args.key, value=value, description=args.description)],
+                action_on_existence=BulkActionOnExistence.OVERWRITE,
+            )
+        ]
+    )
+    api_client.variables.bulk(variables=bulk_body)
     print(f"Variable {args.key} created")
 
 

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_set, ["variables", "set", "foo", "bar"], "airflowctl variables set"),
 ]
 
 

airflow-core/tests/unit/cli/commands/test_variable_command.py

@@ -24,6 +24,7 @@
 
 import pytest
 import yaml
+from airflowctl.api.datamodels.generated import BulkActionOnExistence
 from sqlalchemy import select
 
 from airflow import models
@@ -129,28 +130,42 @@ def setup_method(self):
     def teardown_method(self):
         clear_db_variables()
 
-    def test_variables_set(self):
+    @staticmethod
+    def _set_entity(mock_cli_api_client):
+        """Return the single VariableBody sent by ``set`` through the bulk request."""
+        bulk_body = mock_cli_api_client.variables.bulk.call_args.kwargs["variables"]
+        action = bulk_body.actions[0]
+        assert action.action == "create"
+        assert action.action_on_existence == BulkActionOnExistence.OVERWRITE
+        return action.entities[0]
+
+    def test_variables_set(self, mock_cli_api_client):
         """Test variable_set command"""
         variable_command.variables_set(self.parser.parse_args(["variables", "set", "foo", "bar"]))
-        assert Variable.get("foo") is not None
-        with pytest.raises(KeyError):
-            Variable.get("foo1")
 
-    def test_variables_set_with_description(self):
+        mock_cli_api_client.variables.bulk.assert_called_once()
+        entity = self._set_entity(mock_cli_api_client)
+        assert entity.key == "foo"
+        assert entity.value.root == "bar"
+        assert entity.description is None
+
+    def test_variables_set_with_description(self, mock_cli_api_client):
         """Test variable_set command with optional description argument"""
         expected_var_desc = "foo_bar_description"
         var_key = "foo"
         variable_command.variables_set(
             self.parser.parse_args(["variables", "set", var_key, "bar", "--description", expected_var_desc])
         )
 
-        assert Variable.get(var_key) == "bar"
-        with create_session() as session:
-            actual_var_desc = session.scalar(select(Variable.description).where(Variable.key == var_key))
-            assert actual_var_desc == expected_var_desc
+        assert self._set_entity(mock_cli_api_client).description == expected_var_desc
 
-        with pytest.raises(KeyError):
-            Variable.get("foo1")
+    def test_variables_set_serialize_json(self, mock_cli_api_client):
+        """Test variable_set command with json argument"""
+        variable_command.variables_set(
+            self.parser.parse_args(["variables", "set", "foo", '{"a": 1}', "--json"])
+        )
+
+        assert self._set_entity(mock_cli_api_client).value.root == json.dumps('{"a": 1}', indent=2)
 
     def test_variables_get(self, stdout_capture):
         Variable.set("foo", {"foo": "bar"}, serialize_json=True)
@@ -171,25 +186,19 @@ def test_get_variable_missing_variable(self):
             variable_command.variables_get(self.parser.parse_args(["variables", "get", "no-existing-VAR"]))
 
     def test_variables_set_different_types(self):
-        """Test storage of various data types"""
-        # Set a dict
-        variable_command.variables_set(
-            self.parser.parse_args(["variables", "set", "dict", '{"foo": "oops"}'])
-        )
-        # Set a list
-        variable_command.variables_set(self.parser.parse_args(["variables", "set", "list", '["oops"]']))
-        # Set str
-        variable_command.variables_set(self.parser.parse_args(["variables", "set", "str", "hello string"]))
-        # Set int
-        variable_command.variables_set(self.parser.parse_args(["variables", "set", "int", "42"]))
-        # Set float
-        variable_command.variables_set(self.parser.parse_args(["variables", "set", "float", "42.0"]))
-        # Set true
-        variable_command.variables_set(self.parser.parse_args(["variables", "set", "true", "true"]))
-        # Set false
-        variable_command.variables_set(self.parser.parse_args(["variables", "set", "false", "false"]))
-        # Set none
-        variable_command.variables_set(self.parser.parse_args(["variables", "set", "null", "null"]))
+        """Test export/import round-trips storage of various data types.
+
+        ``set`` is migrated to the airflowctl client, so the variables are seeded directly
+        through the model here; ``export``/``import`` remain local DB commands.
+        """
+        Variable.set("dict", '{"foo": "oops"}')
+        Variable.set("list", '["oops"]')
+        Variable.set("str", "hello string")
+        Variable.set("int", "42")
+        Variable.set("float", "42.0")
+        Variable.set("true", "true")
+        Variable.set("false", "false")
+        Variable.set("null", "null")
 
         # Export and then import
         variable_command.variables_export(
@@ -210,8 +219,8 @@ def test_variables_set_different_types(self):
         assert Variable.get("null", deserialize_json=True) is None
 
         # test variable import skip existing
-        # set varliable list to ["airflow"] and have it skip during import
-        variable_command.variables_set(self.parser.parse_args(["variables", "set", "list", '["airflow"]']))
+        # set variable list to ["airflow"] and have it skip during import
+        Variable.set("list", '["airflow"]')
         variable_command.variables_import(
             self.parser.parse_args(
                 ["variables", "import", "variables_types.json", "--action-on-existing-key", "skip"]
@@ -325,8 +334,8 @@ def test_variables_list_edge_cases(self):
                 assert item["val"] == "***"
 
     def test_variables_delete(self):
-        """Test variable_delete command"""
-        variable_command.variables_set(self.parser.parse_args(["variables", "set", "foo", "bar"]))
+        """Test variable_delete command (``set`` is migrated, so seed via the model)"""
+        Variable.set("foo", "bar")
         variable_command.variables_delete(self.parser.parse_args(["variables", "delete", "foo"]))
         with pytest.raises(KeyError):
             Variable.get("foo")
@@ -365,13 +374,13 @@ def test_variables_isolation(self, tmp_path):
         path1 = tmp_path / "testfile1.json"
         path2 = tmp_path / "testfile2.json"
 
-        # First export
-        variable_command.variables_set(self.parser.parse_args(["variables", "set", "foo", '{"foo":"bar"}']))
-        variable_command.variables_set(self.parser.parse_args(["variables", "set", "bar", "original"]))
+        # First export (``set`` is migrated to airflowctl, so seed via the model)
+        Variable.set("foo", '{"foo":"bar"}')
+        Variable.set("bar", "original")
         variable_command.variables_export(self.parser.parse_args(["variables", "export", os.fspath(path1)]))
 
-        variable_command.variables_set(self.parser.parse_args(["variables", "set", "bar", "updated"]))
-        variable_command.variables_set(self.parser.parse_args(["variables", "set", "foo", '{"foo":"oops"}']))
+        Variable.set("bar", "updated")
+        Variable.set("foo", '{"foo":"oops"}')
         variable_command.variables_delete(self.parser.parse_args(["variables", "delete", "foo"]))
         with create_session() as session:
             variable_command.variables_import(
@@ -389,13 +398,10 @@ def test_variables_isolation(self, tmp_path):
     def test_variables_import_and_export_with_description(self, tmp_path):
         """Test variables_import with file-description parameter"""
         variables_types_file = tmp_path / "variables_types.json"
-        variable_command.variables_set(
-            self.parser.parse_args(["variables", "set", "foo", "bar", "--description", "Foo var description"])
-        )
-        variable_command.variables_set(
-            self.parser.parse_args(["variables", "set", "foo1", "bar1", "--description", "12"])
-        )
-        variable_command.variables_set(self.parser.parse_args(["variables", "set", "foo2", "bar2"]))
+        # ``set`` is migrated to airflowctl, so seed the variables via the model
+        Variable.set("foo", "bar", description="Foo var description")
+        Variable.set("foo1", "bar1", description="12")
+        Variable.set("foo2", "bar2")
         variable_command.variables_export(
             self.parser.parse_args(["variables", "export", os.fspath(variables_types_file)])
         )

airflow-ctl-tests/tests/airflowctl_tests/test_airflowctl_commands.py

@@ -120,6 +120,11 @@ def date_param():
     "variables get test_key",
     "variables get test_key -o table",
     "variables update --key=test_key --value=updated_value",
+    # set is an upsert: the first call creates the key, the second updates the existing key.
+    "variables set test_set_key set_value",
+    "variables set test_set_key updated_set_value",
+    "variables get test_set_key",
+    "variables delete test_set_key",
     "variables import tests/airflowctl_tests/fixtures/test_variables.json",
     "variables delete test_key",
     "variables delete test_import_var",

airflow-ctl/docs/images/command_hashes.txt

@@ -9,7 +9,7 @@ dagrun:c32e0011aa9a845456c778786717208e
 jobs:a5b644c5da8889443bb40ee10b599270
 pools:19efe105b9515ab1926ebcaf0e028d71
 providers:34502fe09dc0b8b0a13e7e46efdffda6
-variables:f8fc76d3d398b2780f4e97f7cd816646
+variables:68cf6c7b27960c35e5e96895053a349f
 version:31f4efdf8de0dbaaa4fac71ff7efecc3
 plugins:4864fd8f356704bd2b3cd1aec3567e35
 auth login:9fe2bb1dd5c602beea2eefb33a2b20a8