From bd6688ff57388f3db7b2754cf7f730c1d8bcea7f Mon Sep 17 00:00:00 2001 From: Ruslan Kuprieiev Date: Sun, 8 Apr 2018 03:29:31 +0300 Subject: [PATCH 1/3] cloud: add ssh driver Signed-off-by: Ruslan Kuprieiev --- dvc/cloud/base.py | 3 + dvc/cloud/data_cloud.py | 8 ++- dvc/cloud/ssh.py | 120 ++++++++++++++++++++++++++++++++++++++++ dvc/utils.py | 12 +++- requirements.txt | 1 + setup.py | 1 + 6 files changed, 143 insertions(+), 2 deletions(-) create mode 100644 dvc/cloud/ssh.py diff --git a/dvc/cloud/base.py b/dvc/cloud/base.py index cce7625b9..3e41a5a2b 100644 --- a/dvc/cloud/base.py +++ b/dvc/cloud/base.py @@ -205,3 +205,6 @@ class DataCloudBase(object): def connect(self): pass + + def disconnect(self): + pass diff --git a/dvc/cloud/data_cloud.py b/dvc/cloud/data_cloud.py index 480cf8a36..319890aa4 100644 --- a/dvc/cloud/data_cloud.py +++ b/dvc/cloud/data_cloud.py @@ -9,6 +9,7 @@ from dvc.utils import map_progress from dvc.cloud.aws import DataCloudAWS from dvc.cloud.gcp import DataCloudGCP +from dvc.cloud.ssh import DataCloudSSH from dvc.cloud.local import DataCloudLOCAL from dvc.cloud.base import DataCloudBase @@ -19,6 +20,7 @@ class DataCloud(object): CLOUD_MAP = { 'aws' : DataCloudAWS, 'gcp' : DataCloudGCP, + 'ssh' : DataCloudSSH, 'local' : DataCloudLOCAL, } @@ -127,7 +129,11 @@ class DataCloud(object): if collect_cloud: collected |= self._collect(cloud, targets, jobs, False) - return map_progress(getattr(cloud, func), list(collected), jobs) + ret = map_progress(getattr(cloud, func), list(collected), jobs) + + cloud.disconnect() + + return ret def push(self, targets, jobs=1, remote=None): """ diff --git a/dvc/cloud/ssh.py b/dvc/cloud/ssh.py new file mode 100644 index 000000000..680317af9 --- /dev/null +++ b/dvc/cloud/ssh.py @@ -0,0 +1,120 @@ +import os +import errno +import filecmp +import paramiko +import posixpath + +from dvc.logger import Logger +from dvc.utils import copyfile, file_md5 +from dvc.cloud.base import DataCloudError, DataCloudBase +from dvc.cloud.aws import create_cb +from dvc.progress import progress + + +class SSHKey(object): + def __init__(self, bucket, name): + self.name = name + self.bucket = bucket + + @property + def path(self): + return os.path.join(self.bucket, self.name) + + +class DataCloudSSH(DataCloudBase): + """ + Driver for remote storage over ssh. + """ + #NOTE: ~/ paths are temporarily forbidden + REGEX = r'^(?P.*)@(?P.*):(?P/+.*)$' + + @property + def hostname(self): + return self.group('host') + + @property + def username(self): + return self.group('user') + + def connect(self): + self.ssh = paramiko.SSHClient() + + self.ssh.load_system_host_keys() + self.ssh.set_missing_host_key_policy(paramiko.AutoAddPolicy()) + + self.ssh.connect(self.hostname, username=self.username) + + def disconnect(self): + self.ssh.close() + + def get_sftp(self): + #NOTE: SFTP doesn't seem to be thread safe + return self.ssh.open_sftp() + + def cache_file_key(self, path): + relpath = os.path.relpath(os.path.abspath(path), self._cloud_settings.cache.cache_dir) + return relpath.replace('\\', '/') + + def _isfile_remote(self, path): + try: + self.get_sftp().open(path, 'r') + return True + except IOError as exc: + if exc.errno != errno.ENOENT: + raise + return False + + def _cmp_checksum(self, key, path): + #FIXME actually compare md5 + path_exists = os.path.exists(path) + key_exists = self._isfile_remote(key.path) + + if path_exists and key_exists: + return True + + return False + + def _get_key(self, path): + key_name = self.cache_file_key(path) + key = SSHKey(self.path, key_name) + if self._isfile_remote(key.path): + return key + return None + + def _makedirs_remote(self, dname): + try: + self.get_sftp().chdir(dname) + except IOError as exc: + if exc.errno != errno.ENOENT: + raise + + parent = posixpath.dirname(dname) + if len(parent): + self._makedirs_remote(parent) + self.get_sftp().mkdir(dname) + + def _new_key(self, path): + key_name = self.cache_file_key(path) + key = SSHKey(self.path, key_name) + self._makedirs_remote(posixpath.dirname(key.path)) + return key + + def _push_key(self, key, path): + self.get_sftp().put(path, key.path, callback=create_cb(key.name)) + progress.finish_target(key.name) + return path + + def _pull_key(self, key, path, no_progress_bar=False): + self._makedirs(path) + + tmp_file = self.tmp_file(path) + try: + self.get_sftp().get(key.path, tmp_file, callback=create_cb(key.name)) + except Exception as exc: + Logger.error('Failed to copy "{}": {}'.format(key.path, exc)) + return None + + os.rename(tmp_file, path) + progress.finish_target(key.name) + + return path diff --git a/dvc/utils.py b/dvc/utils.py index b9aaf62a1..56270352f 100644 --- a/dvc/utils.py +++ b/dvc/utils.py @@ -74,6 +74,14 @@ def copyfile(src, dest, no_progress_bar=False): fdest.close() +def wrap(func, t): + try: + return func(t) + except Exception as exc: + Logger.error('wrap', exc) + raise + + def map_progress(func, targets, n_threads): """ Process targets in multi-threaded mode with progress bar @@ -82,8 +90,10 @@ def map_progress(func, targets, n_threads): pool = ThreadPool(processes=n_threads) ret = [] + wrapper = lambda t: wrap(func, t) + try: - ret = pool.map(func, targets) + ret = pool.map(wrapper, targets) except Exception as exc: raise finally: diff --git a/requirements.txt b/requirements.txt index a74e6217a..7914e1987 100644 --- a/requirements.txt +++ b/requirements.txt @@ -16,3 +16,4 @@ setuptools>=34.0.0 nanotime>=0.5.2 pyasn1>=0.4.1 schema>=0.6.7 +paramiko>=2.4.1 diff --git a/setup.py b/setup.py index 51921f1af..2e9289c2f 100644 --- a/setup.py +++ b/setup.py @@ -22,6 +22,7 @@ install_requires = [ "nanotime>=0.5.2", "pyasn1>=0.4.1", "schema>=0.6.7", + "paramiko>=2.4.1", ] setup( From 36d924c4c24640d4911379c93803647d3c7e874a Mon Sep 17 00:00:00 2001 From: Ruslan Kuprieiev Date: Sun, 8 Apr 2018 21:54:29 +0300 Subject: [PATCH 2/3] test: add test for ssh cloud Signed-off-by: Ruslan Kuprieiev --- tests/test_data_cloud.py | 42 ++++++++++++++++++++++++++++++++++++++++ 1 file changed, 42 insertions(+) diff --git a/tests/test_data_cloud.py b/tests/test_data_cloud.py index 93d8a5e4e..ab87ca19a 100644 --- a/tests/test_data_cloud.py +++ b/tests/test_data_cloud.py @@ -2,6 +2,7 @@ import os import stat import uuid import shutil +import getpass import tempfile from dvc.main import main @@ -51,6 +52,27 @@ def should_test_gcp(): return False +def should_test_ssh(): + if os.getenv("DVC_TEST_SSH") == "true": + return True + + #FIXME: enable on windows + if os.name == 'nt': + return False + + #FIXME: this is ugly + if os.getenv('TRAVIS') == 'true': + os.system('ssh-keygen -t dsa -N "" -C "test key" -f mykey') + os.system('mkdir -p ~/.ssh') + os.system('cp mykey ~/.ssh/id_rsa') + os.system('cp mykey.pub ~/.ssh/id_rsa.pub') + os.system('cat mykey.pub >> ~/.ssh/authorized_keys') + os.system('ssh-keyscan 127.0.0.1 >> ~/.ssh/known_hosts') + assert os.system('ssh 127.0.0.1 ls') == 0 + + return False + + def get_local_storagepath(): return tempfile.mkdtemp() @@ -59,6 +81,10 @@ def get_local_url(): return get_local_storagepath() +def get_ssh_url(): + return '{}@127.0.0.1:{}'.format(getpass.getuser(), get_local_storagepath()) + + def get_aws_storagepath(): return TEST_AWS_REPO_BUCKET + '/' + str(uuid.uuid4()) @@ -285,6 +311,22 @@ class TestDataCloudLOCALCLI(TestDataCloudCLIBase): self._test_cloud(TEST_REMOTE) +class TestDataCloudSSHCLI(TestDataCloudCLIBase): + def should_test(self): + return should_test_ssh() + + def _test_compat(self): + # We do not support ssh in legacy mode + pass + + def _test(self): + url = get_ssh_url() + + self.main(['remote', 'add', TEST_REMOTE, url]) + + self._test_cloud(TEST_REMOTE) + + class TestDataCloudAWSCLI(TestDataCloudCLIBase): def should_test(self): return should_test_aws() From 52a88d6063e6054e887824bf4088c0a23eee960d Mon Sep 17 00:00:00 2001 From: Ruslan Kuprieiev Date: Sun, 8 Apr 2018 21:56:01 +0300 Subject: [PATCH 3/3] test: remove unittests.sh Signed-off-by: Ruslan Kuprieiev --- .appveyor.yml | 2 +- scripts/ci/script.sh | 2 +- tests/__main__.py | 4 ++++ unittests.sh | 15 --------------- 4 files changed, 6 insertions(+), 17 deletions(-) create mode 100644 tests/__main__.py delete mode 100755 unittests.sh diff --git a/.appveyor.yml b/.appveyor.yml index b57a127e9..a1d257cdd 100644 --- a/.appveyor.yml +++ b/.appveyor.yml @@ -61,7 +61,7 @@ before_test: - pip install -r test-requirements.txt test_script: - - nosetests --cover-inclusive --cover-erase --cover-package=dvc --with-coverage + - python -mtests after_test: - if "%APPVEYOR_PULL_REQUEST_NUMBER%"=="" codecov diff --git a/scripts/ci/script.sh b/scripts/ci/script.sh index 60f50763b..283f7493d 100644 --- a/scripts/ci/script.sh +++ b/scripts/ci/script.sh @@ -3,7 +3,7 @@ set -x set -e -./unittests.sh +python -mtests if [[ "$TRAVIS_PULL_REQUEST" == "false" && \ "$TRAVIS_SECURE_ENV_VARS" == "true" ]]; then diff --git a/tests/__main__.py b/tests/__main__.py new file mode 100644 index 000000000..f0cea5e3c --- /dev/null +++ b/tests/__main__.py @@ -0,0 +1,4 @@ +import nose + +if __name__ == '__main__': + nose.main(argv=['--cover-inclusive', '--cover-erase', '--cover-package=dvc', '--with-coverage']) diff --git a/unittests.sh b/unittests.sh deleted file mode 100755 index 31e89d2fc..000000000 --- a/unittests.sh +++ /dev/null @@ -1,15 +0,0 @@ -#!/bin/bash -set -e - -ulimit -n 2048 - -nosetests --cover-inclusive --cover-erase --cover-package=dvc --with-coverage - -if [ "$1" = "report" ]; then - if [ -z "$CODECLIMATE_REPO_TOKEN" ]; then - echo "Error: CODECLIMATE_REPO_TOKEN is not defined! Reporting faied." - else - codeclimate-test-reporter --token $CODECLIMATE_REPO_TOKEN - fi -fi -