diff --git a/dvc/config.py b/dvc/config.py index a1c27d9f6..eb0d5dbeb 100644 --- a/dvc/config.py +++ b/dvc/config.py @@ -42,6 +42,7 @@ class Config(object): schema.Optional(SECTION_CACHE_S3): str, schema.Optional(SECTION_CACHE_GS): str, schema.Optional(SECTION_CACHE_HDFS): str, + schema.Optional(SECTION_CACHE_SSH): str, # backward compatibility schema.Optional(SECTION_CACHE_DIR, default='cache'): str, diff --git a/dvc/dependency/__init__.py b/dvc/dependency/__init__.py index 50222f029..73490fecd 100644 --- a/dvc/dependency/__init__.py +++ b/dvc/dependency/__init__.py @@ -13,6 +13,7 @@ from dvc.dependency.s3 import DependencyS3 from dvc.dependency.gs import DependencyGS from dvc.dependency.local import DependencyLOCAL from dvc.dependency.hdfs import DependencyHDFS +from dvc.dependency.ssh import DependencySSH from dvc.remote import Remote from dvc.remote.local import RemoteLOCAL @@ -21,9 +22,10 @@ from dvc.remote.gs import RemoteGS from dvc.remote.ssh import RemoteSSH from dvc.remote.hdfs import RemoteHDFS -DEPS = [DependencyHDFS, DependencyS3, DependencyGS, DependencyLOCAL] +DEPS = [DependencyHDFS, DependencyS3, DependencyGS, DependencySSH, DependencyLOCAL] DEP_MAP = {'': DependencyLOCAL, + 'ssh': DependencySSH, 's3': DependencyS3, 'gs': DependencyGS, 'hdfs': DependencyHDFS,} diff --git a/dvc/dependency/ssh.py b/dvc/dependency/ssh.py new file mode 100644 index 000000000..2be6b1d46 --- /dev/null +++ b/dvc/dependency/ssh.py @@ -0,0 +1,47 @@ +import posixpath + +try: + from urlparse import urlparse +except ImportError: + from urllib.parse import urlparse + +from dvc.dependency.base import DependencyBase +from dvc.remote.ssh import RemoteSSH + + +class DependencySSH(DependencyBase): + REGEX = RemoteSSH.REGEX + + def __init__(self, stage, path, info=None, remote=None): + super(DependencySSH, self).__init__(stage, path) + self.info = info + self.remote = remote if remote else RemoteSSH(stage.project, {}) + + host = remote.host if remote else self.group('host') + user = remote.user if remote else self.group('user') + if remote: + path = posixpath.join(remote.prefix, urlparse(path).path.lstrip('/')) + else: + path = self.match(self.path).group('path') + + self.path_info = {'scheme': 'ssh', + 'host': host, + 'user': user, + 'path': path} + + def group(self, name): + match = self.match(self.path) + if not match: + return None + return match.group(name) + + def changed(self): + return self.info != self.remote.save_info(self.path_info) + + def save(self): + self.info = self.remote.save_info(self.path_info) + + def dumpd(self): + ret = self.info + ret[self.PARAM_PATH] = self.path + return ret diff --git a/dvc/output/__init__.py b/dvc/output/__init__.py index 057293409..9c478d6c2 100644 --- a/dvc/output/__init__.py +++ b/dvc/output/__init__.py @@ -9,15 +9,17 @@ from dvc.output.s3 import OutputS3 from dvc.output.gs import OutputGS from dvc.output.local import OutputLOCAL from dvc.output.hdfs import OutputHDFS +from dvc.output.ssh import OutputSSH from dvc.remote import Remote -OUTS = [OutputHDFS, OutputS3, OutputGS, OutputLOCAL] +OUTS = [OutputHDFS, OutputS3, OutputGS, OutputSSH, OutputLOCAL] OUTS_MAP = {'hdfs': OutputHDFS, 's3': OutputS3, 'gs': OutputGS, + 'ssh': OutputSSH, '': OutputLOCAL} SCHEMA[schema.Optional(OutputLOCAL.PARAM_CACHE)] = bool diff --git a/dvc/output/ssh.py b/dvc/output/ssh.py new file mode 100644 index 000000000..c7318881d --- /dev/null +++ b/dvc/output/ssh.py @@ -0,0 +1,46 @@ +from dvc.dependency.ssh import DependencySSH +from dvc.exceptions import DvcException + + +class OutputSSH(DependencySSH): + PARAM_CACHE = 'cache' + PARAM_METRIC = 'metric' + + def __init__(self, stage, path, info=None, remote=None, cache=True, metric=False): + super(OutputSSH, self).__init__(stage, path, info, remote=remote) + self.use_cache = cache + self.metric = metric + if cache and self.project.cache.ssh == None: + raise DvcException("No cache location setup for \'ssh\' outputs.") + + def dumpd(self): + ret = super(OutputSSH, self).dumpd() + ret[self.PARAM_CACHE] = self.use_cache + ret[self.PARAM_METRIC] = self.metric + return ret + + def changed(self): + if super(OutputSSH, self).changed(): + return True + + if self.use_cache and self.info != self.project.cache.ssh.save_info(self.path_info): + return True + + return False + + def checkout(self): + if not self.use_cache: + return + + self.project.cache.ssh.checkout(self.path_info, self.info) + + def save(self): + super(OutputSSH, self).save() + + if not self.use_cache: + return + + self.info = self.project.cache.ssh.save(self.path_info) + + def remove(self, ignore_remove=False): + self.remote.remove(self.path_info) diff --git a/dvc/remote/ssh.py b/dvc/remote/ssh.py index 05c618247..77686c86a 100644 --- a/dvc/remote/ssh.py +++ b/dvc/remote/ssh.py @@ -1,5 +1,7 @@ import os +import re import errno +import getpass import filecmp import paramiko import posixpath @@ -36,19 +38,25 @@ def create_cb(name): class RemoteSSH(RemoteBase): scheme = 'ssh' - #NOTE: ~/ paths are temporarily forbidden - REGEX = r'^(?P.*)@(?P.*):(?P/+.*)$' + #NOTE: temporarily only absolute paths are allowed + REGEX = r'^ssh://((?P.*)@)?(?P[^/]*):(?P/.*)$' + + PARAM_MD5 = 'md5' def __init__(self, project, config): self.project = project - storagepath = config.get(Config.SECTION_AWS_STORAGEPATH, '/') - self.url = config.get(Config.SECTION_REMOTE_URL, storagepath) + self.url = config.get(Config.SECTION_REMOTE_URL, '/') self.host = self.group('host') self.user = self.group('user') + if not self.user: + self.user = config.get(Config.SECTION_REMOTE_USER, getpass.getuser()) self.prefix = self.group('path') def group(self, group): - return self.match(self.url).group(group) + m = self.match(self.url) + if not m: + return None + return m.group(group) def cache_file_key(self, path): relpath = os.path.relpath(os.path.abspath(path), self.project.cache.local.cache_dir) @@ -116,6 +124,67 @@ class RemoteSSH(RemoteBase): self._makedirs_remote(ret) return ret + def md5(self, path_info): + if path_info['scheme'] != 'ssh': + raise NotImplementedError + + ssh = self.ssh(host=path_info['host'], user=path_info['user']) + stdin, stdout, stderr = ssh.exec_command('md5sum {}'.format(path_info['path'])) + md5 = re.match(r'^(?P.*) .*$', stdout.read().decode('utf-8')).group('md5') + ssh.close() + + return md5 + + def cp(self, from_info, to_info): + if from_info['scheme'] != 'ssh' or to_info['scheme'] != 'ssh': + raise NotImplementedError + + assert from_info['host'] == to_info['host'] + assert from_info['user'] == to_info['user'] + + ssh = self.ssh(host=from_info['host'], user=from_info['user']) + ssh.exec_command('cp {} {}'.format(from_info['path'], to_info['path'])) + ssh.close() + + def save_info(self, path_info): + if path_info['scheme'] != 'ssh': + raise NotImplementedError + + return {self.PARAM_MD5: self.md5(path_info)} + + def save(self, path_info): + if path_info['scheme'] != 'ssh': + raise NotImplementedError + + md5 = self.md5(path_info) + dest = path_info.copy() + dest['path'] = posixpath.join(self.prefix, md5[0:2], md5[2:]) + + self.cp(path_info, dest) + + return {self.PARAM_MD5: md5} + + def checkout(self, path_info, checksum_info): + if path_info['scheme'] != 'ssh': + raise NotImplementedError + + md5 = checksum_info.get(self.PARAM_MD5, None) + if not md5: + return + + src = path_info.copy() + src['path'] = posixpath.join(self.prefix, md5[0:2], md5[2:]) + + self.cp(src, path_info) + + def remove(self, path_info): + if path_info['scheme'] != 'ssh': + raise NotImplementedError + + ssh = self.ssh(host=path_info['host'], user=path_info['user']) + ssh.open_sftp().remove(path_info['path']) + ssh.close() + def download(self, path_info, path, no_progress_bar=False, name=None): if path_info['scheme'] != 'ssh': raise NotImplementedError diff --git a/tests/test_data_cloud.py b/tests/test_data_cloud.py index e5bfb5dad..9f111a414 100644 --- a/tests/test_data_cloud.py +++ b/tests/test_data_cloud.py @@ -5,10 +5,11 @@ import uuid import shutil import getpass import tempfile +import platform from dvc.main import main from dvc.config import Config, ConfigError -from dvc.data_cloud import DataCloud, RemoteS3, RemoteGS, RemoteLOCAL +from dvc.data_cloud import DataCloud, RemoteS3, RemoteGS, RemoteLOCAL, RemoteSSH from dvc.remote.base import STATUS_UNKNOWN, STATUS_OK, STATUS_MODIFIED, STATUS_NEW, STATUS_DELETED from tests.basic_env import TestDvc @@ -54,8 +55,8 @@ def _should_test_ssh(): if os.getenv("DVC_TEST_SSH") == "true": return True - #FIXME: enable on windows - if os.name == 'nt': + #FIXME: enable on windows and osx + if platform.system() != 'Linux': return False assert os.system('ssh 127.0.0.1 ls &> /dev/null') == 0 @@ -72,7 +73,7 @@ def get_local_url(): def get_ssh_url(): - return '{}@127.0.0.1:{}'.format(getpass.getuser(), get_local_storagepath()) + return 'ssh://{}@127.0.0.1:{}'.format(getpass.getuser(), get_local_storagepath()) def get_aws_storagepath(): @@ -108,7 +109,10 @@ class TestDataCloud(TestDvc): def test(self): config = TEST_CONFIG - for scheme, cl in [('s3://', RemoteS3), ('gs://', RemoteGS), (tempfile.mkdtemp(), RemoteLOCAL)]: + for scheme, cl in [('s3://', RemoteS3), + ('gs://', RemoteGS), + ('ssh://user@localhost:/', RemoteSSH), + (tempfile.mkdtemp(), RemoteLOCAL)]: config[TEST_SECTION][Config.SECTION_REMOTE_URL] = scheme + str(uuid.uuid4()) self._test_cloud(config, cl) @@ -252,6 +256,20 @@ class TestRemoteLOCAL(TestDataCloudBase): self.assertTrue(os.path.isdir(self.dname)) +class TestRemoteSSH(TestDataCloudBase): + def _should_test(self): + return _should_test_ssh() + + def _get_url(self): + return get_ssh_url() + + def _get_cloud_class(self): + return RemoteSSH + + def test(self): + super(TestRemoteSSH, self).test() + + class TestDataCloudCLIBase(TestDvc): def main(self, args): ret = main(args) diff --git a/tests/test_repro.py b/tests/test_repro.py index 7bc5b4b14..f4e15d8a7 100644 --- a/tests/test_repro.py +++ b/tests/test_repro.py @@ -3,7 +3,9 @@ import yaml import stat import shutil import filecmp +import getpass import tempfile +import posixpath from subprocess import Popen, PIPE import boto3 @@ -20,6 +22,7 @@ from dvc.stage import Stage from tests.basic_env import TestDvc from tests.test_data_cloud import _should_test_aws, TEST_AWS_REPO_BUCKET from tests.test_data_cloud import _should_test_gcp, TEST_GCP_REPO_BUCKET +from tests.test_data_cloud import _should_test_ssh class TestRepro(TestDvc): @@ -298,7 +301,7 @@ class TestReproExternalBase(TestDvc): cmd_stage = self.dvc.run(outs=[out_bar_path], deps=[out_foo_path], - cmd='{} {} {}'.format(self.cmd, foo_path, bar_path)) + cmd=self.cmd(foo_path, bar_path)) self.write(self.bucket, foo_key, 'bar') @@ -323,9 +326,8 @@ class TestReproExternalS3(TestReproExternalBase): def bucket(self): return TEST_AWS_REPO_BUCKET - @property - def cmd(self): - return 'aws s3 cp' + def cmd(self, i, o): + return 'aws s3 cp {} {}'.format(i, o) def write(self, bucket, key, body): s3 = boto3.resource('s3') @@ -347,9 +349,8 @@ class TestReproExternalGS(TestReproExternalBase): def bucket(self): return TEST_GCP_REPO_BUCKET - @property - def cmd(self): - return 'gsutil cp' + def cmd(self, i, o): + return 'gsutil cp {} {}'.format(i, o) def write(self, bucket, key, body): client = gc.Client() @@ -373,9 +374,8 @@ class TestReproExternalHDFS(TestReproExternalBase): def bucket(self): return os.getenv('HADOOP_CONTAINER_IP') - @property - def cmd(self): - return 'HADOOP_USER_NAME=root hadoop fs -cp' + def cmd(self, i, o): + return 'HADOOP_USER_NAME=root hadoop fs -cp {} {}'.format(i, o) def write(self, bucket, key, body): url = self.scheme + '://' + bucket + '/' + key @@ -396,6 +396,57 @@ class TestReproExternalHDFS(TestReproExternalBase): p.communicate() self.assertEqual(p.returncode, 0) + +class TestReproExternalSSH(TestReproExternalBase): + _dir = None + + def should_test(self): + return _should_test_ssh() + + @property + def scheme(self): + return 'ssh' + + @property + def bucket(self): + if not self._dir: + self._dir = tempfile.mkdtemp() + return '{}@127.0.0.1:{}'.format(getpass.getuser(), self._dir) + + def cmd(self, i, o): + i = i.strip('ssh://') + o = o.strip('ssh://') + return 'scp {} {}'.format(i, o) + + def write(self, bucket, key, body): + dest = '{}@127.0.0.1'.format(getpass.getuser()) + path = posixpath.join(self._dir, key) + p = Popen('ssh {} rm {}'.format(dest, path), + shell=True, + executable=os.getenv('SHELL'), + stdin=PIPE, + stdout=PIPE, + stderr=PIPE) + p.communicate() + + p = Popen('ssh {} "mkdir -p $(dirname {})"'.format(dest, path), + shell=True, + executable=os.getenv('SHELL'), + stdin=PIPE, + stdout=PIPE, + stderr=PIPE) + p.communicate() + self.assertEqual(p.returncode, 0) + + p = Popen('echo "{}" | ssh {} "tr -d \'\\n\' > {}"'.format(body, dest, path), + shell=True, + executable=os.getenv('SHELL'), + stdin=PIPE, + stdout=PIPE, + stderr=PIPE) + p.communicate() + self.assertEqual(p.returncode, 0) + class TestReproExternalLOCAL(TestReproExternalBase): def setUp(self): @@ -425,11 +476,10 @@ class TestReproExternalLOCAL(TestReproExternalBase): def bucket(self): return self.tmpdir - @property - def cmd(self): + def cmd(self, i, o): if os.name == 'nt': - return 'copy' - return 'cp' + return 'copy {} {}'.format(i, o) + return 'cp {} {}'.format(i, o) def write(self, bucket, key, body): path = os.path.join(bucket, key)