Merge pull request #631 from efiop/master

Add initial SSH driver

.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

dvc/cloud/base.py

@@ -205,3 +205,6 @@ class DataCloudBase(object):
 
     def connect(self):
         pass
+
+    def disconnect(self):
+        pass

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):
         """

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<user>.*)@(?P<host>.*):(?P<path>/+.*)$'
+
+    @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

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:

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

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

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(

tests/__main__.py

@@ -0,0 +1,4 @@
+import nose
+
+if __name__ == '__main__':
+    nose.main(argv=['--cover-inclusive', '--cover-erase', '--cover-package=dvc', '--with-coverage'])

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()

unittests.sh

@@ -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
-