dvc machine: add a new command rename to it (#6633)

* `dvc machine`: add a new command `rename` to it

fix #6481
1. add new command `dvc machine rename`
2. add a unit test for the call.
3. add several functional tests for it.t push myself fix6481

* Move rename logic to tpi

* use tpi.state_mv instead of rename

* Rebase to a newer version

* Update tpi version

* Remove unused assert and try catch

dvc/command/machine.py

@@ -3,6 +3,7 @@ import argparse
 from dvc.command.base import CmdBase, append_doc_link, fix_subparsers
 from dvc.command.config import CmdConfig
 from dvc.config import ConfigError
+from dvc.exceptions import DvcException
 from dvc.ui import ui
 from dvc.utils import format_link
 
@@ -104,6 +105,54 @@ class CmdMachineModify(CmdMachineConfig):
         return 0
 
 
+class CmdMachineRename(CmdBase):
+    def _check_exists(self, conf):
+        if self.args.name not in conf["machine"]:
+            raise ConfigError(f"machine '{self.args.name}' doesn't exist.")
+
+    def _rename_default(self, conf):
+        if conf["core"].get("machine") == self.args.name:
+            conf["core"]["machine"] = self.args.new
+
+    def _check_before_rename(self):
+        from dvc.machine import validate_name
+
+        validate_name(self.args.new)
+
+        all_config = self.config.load_config_to_level(None)
+        if self.args.new in all_config.get("machine", {}):
+            raise ConfigError(
+                "Rename failed. Machine '{}' already exists.".format(
+                    self.args.new
+                )
+            )
+        ui.write(f"Rename machine '{self.args.name}' to '{self.args.new}'.")
+
+    def run(self):
+
+        self._check_before_rename()
+
+        with self.config.edit(self.args.level) as conf:
+            self._check_exists(conf)
+            conf["machine"][self.args.new] = conf["machine"][self.args.name]
+            try:
+                self.repo.machine.rename(self.args.name, self.args.new)
+            except DvcException as error:
+                del conf["machine"][self.args.new]
+                raise ConfigError("terraform rename failed") from error
+            del conf["machine"][self.args.name]
+            self._rename_default(conf)
+
+        up_to_level = self.args.level or "repo"
+        for level in reversed(self.config.LEVELS):
+            if level == up_to_level:
+                break
+            with self.config.edit(level) as level_conf:
+                self._rename_default(level_conf)
+
+        return 0
+
+
 class CmdMachineDefault(CmdMachineConfig):
     def run(self):
         if self.args.name is None and not self.args.unset:
@@ -296,6 +345,18 @@ def add_parser(subparsers, parent_parser):
     )
     machine_modify_parser.set_defaults(func=CmdMachineModify)
 
+    machine_RENAME_HELP = "Rename a machine "
+    machine_rename_parser = machine_subparsers.add_parser(
+        "rename",
+        parents=[parent_config_parser, parent_parser],
+        description=append_doc_link(machine_RENAME_HELP, "remote/rename"),
+        help=machine_RENAME_HELP,
+        formatter_class=argparse.RawDescriptionHelpFormatter,
+    )
+    machine_rename_parser.add_argument("name", help="Machine to be renamed")
+    machine_rename_parser.add_argument("new", help="New name of the machine")
+    machine_rename_parser.set_defaults(func=CmdMachineRename)
+
     machine_REMOVE_HELP = "Remove an machine."
     machine_remove_parser = machine_subparsers.add_parser(
         "remove",

dvc/machine/__init__.py

@@ -206,3 +206,8 @@ class MachineManager:
     def status(self, name: str) -> Iterator[Dict[Any, Any]]:
         config, backend = self.get_config_and_backend(name)
         return backend.instances(**config)
+
+    def rename(self, name: str, new: str):
+        """move configuration to a new location if the machine is running."""
+        config, backend = self.get_config_and_backend(name)
+        return backend.rename(new=new, **config)

dvc/machine/backend/base.py

@@ -48,3 +48,7 @@ class BaseMachineBackend(ABC):
     ) -> Iterator["SSHFileSystem"]:
         """Return an sshfs instance for the default directory on the
         specified machine."""
+
+    @abstractmethod
+    def rename(self, name: str, new: str, **config):
+        """Rename a machine instance."""

dvc/machine/backend/terraform.py

@@ -72,3 +72,22 @@ class TerraformBackend(BaseMachineBackend):
                 keyfile=pem,
             )
             yield fs
+
+    def rename(self, name: str, new: str, **config):
+        """rename a dvc machine instance to a new name"""
+        import shutil
+
+        mtype = "iterative_machine"
+
+        new_dir = os.path.join(self.tmp_dir, new)
+        old_dir = os.path.join(self.tmp_dir, name)
+        if os.path.exists(new_dir):
+            raise DvcException(f"rename failed: path {new_dir} already exists")
+
+        if not os.path.exists(old_dir):
+            return
+
+        with self.make_tpi(name) as tpi:
+            tpi.state_mv(f"{mtype}.{name}", f"{mtype}.{new}", **config)
+
+        shutil.move(old_dir, new_dir)

setup.cfg

@@ -111,7 +111,7 @@ ssh_gssapi = sshfs[gssapi]>=2021.8.1
 webdav = webdav4>=0.9.3
 # not to break `dvc[webhdfs]`
 webhdfs =
-terraform = tpi[ssh]>=2.0.0
+terraform = tpi[ssh]>=2.1.0
 tests =
     %(terraform)s
     wheel==0.37.0

tests/func/machine/test_machine_config.py

@@ -2,6 +2,7 @@ import os
 import textwrap
 
 import pytest
+import tpi
 
 from dvc.main import main
 
@@ -114,3 +115,62 @@ def test_machine_list(tmp_dir, dvc, capsys):
         )
         in cap.out
     )
+
+
+def test_machine_rename_success(tmp_dir, scm, dvc, capsys, mocker):
+    config_file = tmp_dir / ".dvc" / "config"
+    config_file.write_text(CONFIG_TEXT)
+
+    mocker.patch.object(
+        tpi.terraform.TerraformBackend,
+        "state_mv",
+        autospec=True,
+        return_value=True,
+    )
+
+    os.makedirs((tmp_dir / ".dvc" / "tmp" / "machine" / "terraform" / "foo"))
+
+    assert main(["machine", "rename", "foo", "bar"]) == 0
+    cap = capsys.readouterr()
+    assert "Rename machine 'foo' to 'bar'." in cap.out
+    assert config_file.read_text() == CONFIG_TEXT.replace("foo", "bar")
+    assert not (
+        tmp_dir / ".dvc" / "tmp" / "machine" / "terraform" / "foo"
+    ).exists()
+    assert (
+        tmp_dir / ".dvc" / "tmp" / "machine" / "terraform" / "bar"
+    ).exists()
+
+
+def test_machine_rename_none_exist(tmp_dir, scm, dvc, caplog):
+    config_alice = CONFIG_TEXT.replace("foo", "alice")
+    config_file = tmp_dir / ".dvc" / "config"
+    config_file.write_text(config_alice)
+    assert main(["machine", "rename", "foo", "bar"]) == 251
+    assert config_file.read_text() == config_alice
+    assert "machine 'foo' doesn't exist." in caplog.text
+
+
+def test_machine_rename_exist(tmp_dir, scm, dvc, caplog):
+    config_bar = CONFIG_TEXT + "['machine \"bar\"']\n    cloud = aws"
+    config_file = tmp_dir / ".dvc" / "config"
+    config_file.write_text(config_bar)
+    assert main(["machine", "rename", "foo", "bar"]) == 251
+    assert config_file.read_text() == config_bar
+    assert "Machine 'bar' already exists." in caplog.text
+
+
+def test_machine_rename_error(tmp_dir, scm, dvc, caplog, mocker):
+    config_file = tmp_dir / ".dvc" / "config"
+    config_file.write_text(CONFIG_TEXT)
+
+    os.makedirs((tmp_dir / ".dvc" / "tmp" / "machine" / "terraform" / "foo"))
+
+    def cmd_error(self, source, destination, **kwargs):
+        raise tpi.TPIError("test error")
+
+    mocker.patch.object(tpi.terraform.TerraformBackend, "state_mv", cmd_error)
+
+    assert main(["machine", "rename", "foo", "bar"]) == 251
+    assert config_file.read_text() == CONFIG_TEXT
+    assert "rename failed" in caplog.text

tests/unit/command/test_machine.py

@@ -8,6 +8,7 @@ from dvc.command.machine import (
     CmdMachineList,
     CmdMachineModify,
     CmdMachineRemove,
+    CmdMachineRename,
     CmdMachineSsh,
     CmdMachineStatus,
 )
@@ -116,3 +117,13 @@ def test_modified(tmp_dir):
     assert cmd.run() == 0
     config = configobj.ConfigObj(str(tmp_dir / ".dvc" / "config"))
     assert config['machine "foo"']["cloud"] == "azure"
+
+
+def test_rename(tmp_dir, scm, dvc):
+    tmp_dir.gen(DATA)
+    cli_args = parse_args(["machine", "rename", "foo", "bar"])
+    assert cli_args.func == CmdMachineRename
+    cmd = cli_args.func(cli_args)
+    assert cmd.run() == 0
+    config = configobj.ConfigObj(str(tmp_dir / ".dvc" / "config"))
+    assert config['machine "bar"']["cloud"] == "aws"