From 9fc4c607499e3d6458debdeebe5fd3e27845edd1 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Peter=20Rowlands=20=28=EB=B3=80=EA=B8=B0=ED=98=B8=29?= Date: Wed, 1 Sep 2021 07:20:13 +0000 Subject: [PATCH] machine: add `dvc machine ssh` (#6501) * machine: add naive asyncssh shell client * machine: add `dvc machine ssh` command * add ssh requirements to terraform install extras * update asyncssh comments * fix py <3.7 * fix async mocking in py<3.8 * use mock instead of unittest.mock --- dvc/command/machine.py | 22 +++++++++++++++++ dvc/machine/__init__.py | 4 ++++ dvc/machine/backend/base.py | 36 ++++++++++++++++++++++++++++ dvc/machine/backend/terraform.py | 20 +++++++++++++--- setup.py | 2 +- tests/unit/command/test_machine.py | 14 +++++++++++ tests/unit/machine/test_terraform.py | 14 +++++++++++ 7 files changed, 108 insertions(+), 4 deletions(-) diff --git a/dvc/command/machine.py b/dvc/command/machine.py index c5a51f3da..9d7c72b65 100644 --- a/dvc/command/machine.py +++ b/dvc/command/machine.py @@ -136,6 +136,15 @@ class CmdMachineDestroy(CmdBase): return 0 +class CmdMachineSsh(CmdBase): + def run(self): + if self.repo.machine is None: + raise MachineDisabledError + + self.repo.machine.run_shell(self.args.name) + return 0 + + def add_parser(subparsers, parent_parser): from dvc.command.config import parent_config_parser @@ -272,3 +281,16 @@ def add_parser(subparsers, parent_parser): "name", help="Name of the machine instance to destroy." ) machine_destroy_parser.set_defaults(func=CmdMachineDestroy) + + machine_SSH_HELP = "Connect to a machine via SSH." + machine_ssh_parser = machine_subparsers.add_parser( + "ssh", + parents=[parent_config_parser, parent_parser], + description=append_doc_link(machine_SSH_HELP, "machine/ssh"), + help=machine_SSH_HELP, + formatter_class=argparse.RawDescriptionHelpFormatter, + ) + machine_ssh_parser.add_argument( + "name", help="Name of the machine instance to connect to." + ) + machine_ssh_parser.set_defaults(func=CmdMachineSsh) diff --git a/dvc/machine/__init__.py b/dvc/machine/__init__.py index 164864257..20148d5e3 100644 --- a/dvc/machine/__init__.py +++ b/dvc/machine/__init__.py @@ -164,3 +164,7 @@ class MachineManager: def get_sshfs(self, name: Optional[str]): config, backend = self.get_config_and_backend(name) return backend.get_sshfs(**config) + + def run_shell(self, name: Optional[str]): + config, backend = self.get_config_and_backend(name) + return backend.run_shell(**config) diff --git a/dvc/machine/backend/base.py b/dvc/machine/backend/base.py index 4946823bd..5a90353c6 100644 --- a/dvc/machine/backend/base.py +++ b/dvc/machine/backend/base.py @@ -1,8 +1,12 @@ +import asyncio import logging +import os +import sys from abc import ABC, abstractmethod from contextlib import contextmanager from typing import TYPE_CHECKING, Iterator, Optional +from dvc.exceptions import DvcException from dvc.types import StrPath from dvc.utils.fs import makedirs @@ -50,3 +54,35 @@ class BaseMachineBackend(ABC): ) -> Iterator["SSHFileSystem"]: """Return an sshfs instance for the default directory on the specified machine.""" + + @abstractmethod + def run_shell(self, name: Optional[str] = None, **config): + """Spawn an interactive SSH shell for the specified machine.""" + + def _shell(self, *args, **kwargs): + """Sync wrapper for an asyncssh shell session. + + Args will be passed into asyncssh.connect(). + """ + import asyncssh + + loop = asyncio.new_event_loop() + try: + asyncio.set_event_loop(loop) + loop.run_until_complete(self._shell_async(*args, **kwargs)) + except (OSError, asyncssh.Error) as exc: + raise DvcException("SSH connection failed") from exc + finally: + asyncio.set_event_loop(None) + loop.close() + + async def _shell_async(self, *args, **kwargs): + import asyncssh + + async with asyncssh.connect(*args, **kwargs) as conn: + await conn.run( + term_type=os.environ.get("TERM", "xterm"), + stdin=sys.stdin, + stdout=sys.stdout, + stderr=sys.stderr, + ) diff --git a/dvc/machine/backend/terraform.py b/dvc/machine/backend/terraform.py index 848a97e74..7accbb705 100644 --- a/dvc/machine/backend/terraform.py +++ b/dvc/machine/backend/terraform.py @@ -109,13 +109,17 @@ class TerraformBackend(BaseMachineBackend): ) -> "BaseExecutor": raise NotImplementedError + def _default_resource(self, name): + resource = first(self.instances(name)) + if not resource: + raise TerraformError(f"No active '{name}' instances") + return resource + @contextmanager def get_sshfs( self, name: Optional[str] = None, **config ) -> Iterator["SSHFileSystem"]: - resource = first(self.instances(name)) - if not resource: - raise TerraformError(f"No active '{name}' instances") + resource = self._default_resource(name) with DvcTerraform.pemfile(resource) as pem: fs = SSHFileSystem( host=resource["instance_ip"], @@ -123,3 +127,13 @@ class TerraformBackend(BaseMachineBackend): keyfile=pem, ) yield fs + + def run_shell(self, name: Optional[str] = None, **config): + resource = self._default_resource(name) + with DvcTerraform.pemfile(resource) as pem: + self._shell( + host=resource["instance_ip"], + username="ubuntu", + client_keys=pem, + known_hosts=None, + ) diff --git a/setup.py b/setup.py index 926565c0f..ed13f905b 100644 --- a/setup.py +++ b/setup.py @@ -114,7 +114,7 @@ webdav = ["webdav4>=0.9.0"] ssh_gssapi = ["sshfs[gssapi]>=2021.7.1"] all_remotes = gs + s3 + azure + ssh + oss + gdrive + hdfs + webhdfs + webdav -terraform = ["python-terraform>=0.10.1", "jinja2>=2.0.0"] +terraform = ssh + ["python-terraform>=0.10.1", "jinja2>=2.0.0"] tests_requirements = ( Path("test_requirements.txt").read_text().strip().splitlines() diff --git a/tests/unit/command/test_machine.py b/tests/unit/command/test_machine.py index 2e4458a0a..02af8feeb 100644 --- a/tests/unit/command/test_machine.py +++ b/tests/unit/command/test_machine.py @@ -6,6 +6,7 @@ from dvc.command.machine import ( CmdMachineCreate, CmdMachineDestroy, CmdMachineRemove, + CmdMachineSsh, ) @@ -64,3 +65,16 @@ def test_destroy(tmp_dir, dvc, mocker): assert cmd.run() == 0 m.assert_called_once_with("foo") + + +def test_ssh(tmp_dir, dvc, mocker): + cli_args = parse_args(["machine", "ssh", "foo"]) + assert cli_args.func == CmdMachineSsh + + cmd = cli_args.func(cli_args) + m = mocker.patch.object( + cmd.repo.machine, "run_shell", autospec=True, return_value=0 + ) + + assert cmd.run() == 0 + m.assert_called_once_with("foo") diff --git a/tests/unit/machine/test_terraform.py b/tests/unit/machine/test_terraform.py index cfa81d4e3..d45e1344f 100644 --- a/tests/unit/machine/test_terraform.py +++ b/tests/unit/machine/test_terraform.py @@ -2,6 +2,7 @@ import json import os import pytest +from mock import ANY, MagicMock from python_terraform import IsFlagged from dvc.machine.backend.terraform import DvcTerraform, TerraformError @@ -127,3 +128,16 @@ def test_instances(tmp_dir, terraform, resource): data = json.loads(TEST_RESOURCE_STATE) expected = data["resources"][0]["instances"][0]["attributes"] assert list(terraform.instances(resource)) == [expected] + + +def test_shell(tmp_dir, terraform, resource, mocker): + data = json.loads(TEST_RESOURCE_STATE) + expected = data["resources"][0]["instances"][0]["attributes"] + mock_connect = mocker.patch("asyncssh.connect", return_value=MagicMock()) + terraform.run_shell(name=resource) + mock_connect.assert_called_once_with( + host=expected["instance_ip"], + username="ubuntu", + client_keys=ANY, + known_hosts=None, + )