diff --git a/dvc/command/machine.py b/dvc/command/machine.py index 24f0e22e9..1be9ad3e4 100644 --- a/dvc/command/machine.py +++ b/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", diff --git a/dvc/machine/__init__.py b/dvc/machine/__init__.py index 4bb8d4cce..211d72373 100644 --- a/dvc/machine/__init__.py +++ b/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) diff --git a/dvc/machine/backend/base.py b/dvc/machine/backend/base.py index 833217b49..9a0c8aa28 100644 --- a/dvc/machine/backend/base.py +++ b/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.""" diff --git a/dvc/machine/backend/terraform.py b/dvc/machine/backend/terraform.py index 51fb173f0..b4a9189f9 100644 --- a/dvc/machine/backend/terraform.py +++ b/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) diff --git a/setup.cfg b/setup.cfg index 3b196441d..a08d51bfd 100644 --- a/setup.cfg +++ b/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 diff --git a/tests/func/machine/test_machine_config.py b/tests/func/machine/test_machine_config.py index af90b924f..701c69d33 100644 --- a/tests/func/machine/test_machine_config.py +++ b/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 diff --git a/tests/unit/command/test_machine.py b/tests/unit/command/test_machine.py index 33159197f..4177a89ab 100644 --- a/tests/unit/command/test_machine.py +++ b/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"