From 5e7a152b530e9bfaa9641a4d0f1d50b5a01efabb Mon Sep 17 00:00:00 2001 From: "Mr. Outis" Date: Wed, 13 Feb 2019 17:06:16 -0600 Subject: [PATCH] ssh: use sftp when possible Related #830 [1/2] --- dvc/remote/azure.py | 4 +- dvc/remote/base.py | 9 - dvc/remote/gs.py | 4 +- dvc/remote/http.py | 4 +- dvc/remote/local.py | 6 +- dvc/remote/s3.py | 4 +- dvc/remote/ssh.py | 388 ----------------------------------- dvc/remote/ssh/__init__.py | 208 +++++++++++++++++++ dvc/remote/ssh/connection.py | 257 +++++++++++++++++++++++ dvc/utils/compat.py | 25 +++ tests/test_data_cloud.py | 2 - tests/unit/remote/ssh.py | 2 +- 12 files changed, 502 insertions(+), 411 deletions(-) delete mode 100644 dvc/remote/ssh.py create mode 100644 dvc/remote/ssh/__init__.py create mode 100644 dvc/remote/ssh/connection.py diff --git a/dvc/remote/azure.py b/dvc/remote/azure.py index d73d431be..43a978d02 100644 --- a/dvc/remote/azure.py +++ b/dvc/remote/azure.py @@ -10,7 +10,7 @@ except ImportError: BlockBlobService = None import dvc.logger as logger -from dvc.utils.compat import urlparse +from dvc.utils.compat import urlparse, makedirs from dvc.progress import progress from dvc.config import Config from dvc.remote.base import RemoteBase @@ -176,7 +176,7 @@ class RemoteAzure(RemoteBase): cb = None if no_progress_bar else Callback(name) - self._makedirs(to_info["path"]) + makedirs(os.path.dirname(to_info["path"]), exist_ok=True) try: self.blob_service.get_blob_to_path( diff --git a/dvc/remote/base.py b/dvc/remote/base.py index 5793751e1..6dda03602 100644 --- a/dvc/remote/base.py +++ b/dvc/remote/base.py @@ -197,15 +197,6 @@ class RemoteBase(object): def copy(self, from_info, to_info): raise RemoteActionNotImplemented("copy", self.scheme) - def _makedirs(self, fname): - dname = os.path.dirname(fname) - - try: - os.makedirs(dname) - except OSError as e: - if e.errno != errno.EEXIST: - raise - def exists(self, path_infos): raise NotImplementedError diff --git a/dvc/remote/gs.py b/dvc/remote/gs.py index c2d4c3dbe..df9e5af7c 100644 --- a/dvc/remote/gs.py +++ b/dvc/remote/gs.py @@ -9,7 +9,7 @@ except ImportError: storage = None import dvc.logger as logger -from dvc.utils.compat import urlparse +from dvc.utils.compat import urlparse, makedirs from dvc.remote.base import RemoteBase from dvc.config import Config from dvc.progress import progress @@ -204,7 +204,7 @@ class RemoteGS(RemoteBase): # lets at least update progress at pathpoints(start, finish) progress.update_target(name, 0, None) - self._makedirs(to_info["path"]) + makedirs(os.path.dirname(to_info["path"]), exist_ok=True) try: bucket = gs.bucket(from_info["bucket"]) diff --git a/dvc/remote/http.py b/dvc/remote/http.py index 7e2b3855b..c6232e29e 100644 --- a/dvc/remote/http.py +++ b/dvc/remote/http.py @@ -1,5 +1,5 @@ from __future__ import unicode_literals -from dvc.utils.compat import open +from dvc.utils.compat import open, makedirs import os import threading @@ -67,7 +67,7 @@ class RemoteHTTP(RemoteBase): if not name: name = os.path.basename(to_info["path"]) - self._makedirs(to_info["path"]) + makedirs(os.path.dirname(to_info["path"]), exist_ok=True) total = self._content_length(from_info["path"]) diff --git a/dvc/remote/local.py b/dvc/remote/local.py index b5cefc611..bc698041b 100644 --- a/dvc/remote/local.py +++ b/dvc/remote/local.py @@ -1,6 +1,6 @@ from __future__ import unicode_literals -from dvc.utils.compat import str +from dvc.utils.compat import str, makedirs import os import stat @@ -484,7 +484,7 @@ class RemoteLOCAL(RemoteBase): if not name: name = os.path.basename(from_info["path"]) - self._makedirs(to_info["path"]) + makedirs(os.path.dirname(to_info["path"]), exist_ok=True) try: copyfile(from_info["path"], to_info["path"], name=name) @@ -521,7 +521,7 @@ class RemoteLOCAL(RemoteBase): if not name: name = os.path.basename(to_info["path"]) - self._makedirs(to_info["path"]) + makedirs(os.path.dirname(to_info["path"]), exist_ok=True) tmp_file = self.tmp_file(to_info["path"]) try: copyfile( diff --git a/dvc/remote/s3.py b/dvc/remote/s3.py index ad8349fa6..44186b52c 100644 --- a/dvc/remote/s3.py +++ b/dvc/remote/s3.py @@ -9,7 +9,7 @@ except ImportError: boto3 = None import dvc.logger as logger -from dvc.utils.compat import urlparse +from dvc.utils.compat import urlparse, makedirs from dvc.progress import progress from dvc.config import Config from dvc.remote.base import RemoteBase @@ -236,7 +236,7 @@ class RemoteS3(RemoteBase): if not name: name = os.path.basename(to_info["path"]) - self._makedirs(to_info["path"]) + makedirs(os.path.dirname(to_info["path"]), exist_ok=True) try: if no_progress_bar: diff --git a/dvc/remote/ssh.py b/dvc/remote/ssh.py deleted file mode 100644 index 2b175e45c..000000000 --- a/dvc/remote/ssh.py +++ /dev/null @@ -1,388 +0,0 @@ -from __future__ import unicode_literals - -import os -import getpass -import posixpath - -try: - import paramiko -except ImportError: - paramiko = None - -import dvc.prompt as prompt -import dvc.logger as logger - -from dvc.config import Config -from dvc.progress import progress -from dvc.utils.compat import urlparse -from dvc.exceptions import DvcException -from dvc.remote.base import RemoteBase, RemoteCmdError - - -def sizeof_fmt(num, suffix="B"): - """ Convert number of bytes to human-readable string """ - for unit in ["", "K", "M", "G", "T", "P", "E", "Z"]: - if abs(num) < 1024.0: - return "%3.1f%s%s" % (num, unit, suffix) - num /= 1024.0 - return "%.1f%s%s" % (num, "Y", suffix) - - -def percent_cb(name, complete, total): - """ Callback for updating target progress """ - logger.debug( - "{}: {} transferred out of {}".format( - name, sizeof_fmt(complete), sizeof_fmt(total) - ) - ) - progress.update_target(name, complete, total) - - -def create_cb(name): - """ Create callback function for multipart object """ - return lambda cur, tot: percent_cb(name, cur, tot) - - -class RemoteSSH(RemoteBase): - scheme = "ssh" - - # NOTE: we support both URL-like (ssh://[user@]host.xz[:port]/path) and - # SCP-like (ssh://[user@]host.xz:path) urls. - REGEX = r"^ssh://.*$" - - REQUIRES = {"paramiko": paramiko} - - JOBS = 4 - - PARAM_CHECKSUM = "md5" - - DEFAULT_PORT = 22 - TIMEOUT = 1800 - - def __init__(self, repo, config): - super(RemoteSSH, self).__init__(repo, config) - self.url = config.get(Config.SECTION_REMOTE_URL, "ssh://") - - parsed = urlparse(self.url) - self.host = parsed.hostname - self.user = ( - config.get(Config.SECTION_REMOTE_USER) - or parsed.username - or getpass.getuser() - ) - self.prefix = parsed.path or "/" - self.port = ( - config.get(Config.SECTION_REMOTE_PORT) - or parsed.port - or self.DEFAULT_PORT - ) - self.keyfile = config.get(Config.SECTION_REMOTE_KEY_FILE, None) - self.timeout = config.get(Config.SECTION_REMOTE_TIMEOUT, self.TIMEOUT) - self.password = config.get(Config.SECTION_REMOTE_PASSWORD, None) - self.ask_password = config.get( - Config.SECTION_REMOTE_ASK_PASSWORD, False - ) - - self.path_info = { - "scheme": "ssh", - "host": self.host, - "user": self.user, - "port": self.port, - } - - def ssh(self, host=None, user=None, port=None): - msg = ( - "Establishing ssh connection with '{}' " - "through port '{}' as user '{}'" - ) - logger.debug(msg.format(host, port, user)) - - ssh = paramiko.SSHClient() - - ssh.load_system_host_keys() - ssh.set_missing_host_key_policy(paramiko.AutoAddPolicy()) - - if self.ask_password and self.password is None: - msg = ( - "Enter a private key passphrase or a password for " - "host '{}' port '{}' user '{}'" - ).format(host, port, user) - self.password = prompt.password(msg) - - ssh.connect( - host, - username=user, - port=port, - key_filename=self.keyfile, - timeout=self.timeout, - password=self.password, - ) - - return ssh - - def exists(self, path_info): - assert not isinstance(path_info, list) - assert path_info["scheme"] == "ssh" - - with self.ssh( - path_info["host"], path_info["user"], path_info["port"] - ) as ssh: - try: - self._exec(ssh, "test -e {}".format(path_info["path"])) - exists = True - except RemoteCmdError: - exists = False - - return exists - - def _exec(self, ssh, cmd): - stdin, stdout, stderr = ssh.exec_command(cmd) - channel = stdout.channel - - stdin.close() - channel.shutdown_write() - - stdout_chunks = [] - stderr_chunks = [] - while ( - not channel.closed - or channel.recv_ready() - or channel.recv_stderr_ready() - ): - import select - - got_chunk = False - readq, _, _ = select.select([stdout.channel], [], [], self.timeout) - for c in readq: - if c.recv_ready(): - stdout_chunks.append(stdout.channel.recv(len(c.in_buffer))) - got_chunk = True - - if c.recv_stderr_ready(): - stderr_len = len(c.in_stderr_buffer) - s = stderr.channel.recv_stderr(stderr_len) - stderr_chunks.append(s) - got_chunk = True - - if ( - not got_chunk - and stdout.channel.exit_status_ready() - and not stderr.channel.recv_stderr_ready() - and not stdout.channel.recv_ready() - ): - stdout.channel.shutdown_read() - stdout.channel.close() - break - - stdout.close() - stderr.close() - - ret = stdout.channel.recv_exit_status() - if ret != 0: - err = b"".join(stderr_chunks).decode("utf-8") - raise RemoteCmdError(self.scheme, cmd, ret, err) - - return b"".join(stdout_chunks).decode("utf-8") - - def md5(self, path_info): - if path_info["scheme"] != "ssh": - raise NotImplementedError - - ssh = self.ssh( - host=path_info["host"], - user=path_info["user"], - port=path_info["port"], - ) - - # Use different md5 commands depending on os - stdout = self._exec(ssh, "uname").strip() - if stdout == "Darwin": - md5cmd = "md5" - index = -1 - elif stdout == "Linux": - md5cmd = "md5sum" - index = 0 - else: - msg = "'{}' is not supported as a remote".format(stdout) - raise DvcException(msg) - - stdout = self._exec(ssh, "{} {}".format(md5cmd, path_info["path"])) - md5 = stdout.split()[index] - ssh.close() - - assert len(md5) == 32 - - return md5 - - def copy(self, from_info, to_info, ssh=None): - 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"] - - s = ( - ssh - if ssh - else self.ssh( - host=from_info["host"], - user=from_info["user"], - port=from_info["port"], - ) - ) - - dname = posixpath.dirname(to_info["path"]) - self._exec(s, "mkdir -p {}".format(dname)) - self._exec(s, "cp {} {}".format(from_info["path"], to_info["path"])) - - if not ssh: - s.close() - - def save_info(self, path_info): - if path_info["scheme"] != "ssh": - raise NotImplementedError - - return {self.PARAM_CHECKSUM: self.md5(path_info)} - - def save(self, path_info, checksum_info): - if path_info["scheme"] != "ssh": - raise NotImplementedError - - md5 = checksum_info[self.PARAM_CHECKSUM] - dest = path_info.copy() - dest["path"] = self.checksum_to_path(md5) - - self.copy(path_info, dest) - - def remove(self, path_info): - if path_info["scheme"] != "ssh": - raise NotImplementedError - - logger.debug( - "Removing ssh://{}@{}/{}".format( - path_info["user"], path_info["host"], path_info["path"] - ) - ) - - ssh = self.ssh( - host=path_info["host"], - user=path_info["user"], - port=path_info["port"], - ) - ssh.open_sftp().remove(path_info["path"]) - ssh.close() - - def download( - self, - from_infos, - to_infos, - no_progress_bar=False, - names=None, - resume=False, - ): - names = self._verify_path_args(from_infos, to_infos, names) - - ssh = self.ssh( - host=from_infos[0]["host"], - user=from_infos[0]["user"], - port=from_infos[0]["port"], - ) - - for to_info, from_info, name in zip(to_infos, from_infos, names): - if from_info["scheme"] != "ssh": - raise NotImplementedError - - if to_info["scheme"] == "ssh": - assert from_info["host"] == to_info["host"] - assert from_info["port"] == to_info["port"] - assert from_info["user"] == to_info["user"] - self.copy(from_info, to_info, ssh=ssh) - continue - - if to_info["scheme"] != "local": - raise NotImplementedError - - msg = "Downloading '{}/{}' to '{}'".format( - from_info["host"], from_info["path"], to_info["path"] - ) - logger.debug(msg) - - if not name: - name = os.path.basename(to_info["path"]) - - self._makedirs(to_info["path"]) - tmp_file = self.tmp_file(to_info["path"]) - try: - ssh.open_sftp().get( - from_info["path"], tmp_file, callback=create_cb(name) - ) - except Exception: - msg = "failed to download '{}/{}' to '{}'" - logger.error( - msg.format( - from_info["host"], from_info["path"], to_info["path"] - ) - ) - continue - - os.rename(tmp_file, to_info["path"]) - progress.finish_target(name) - - ssh.close() - - def upload(self, from_infos, to_infos, names=None): - names = self._verify_path_args(to_infos, from_infos, names) - - ssh = self.ssh( - host=to_infos[0]["host"], - user=to_infos[0]["user"], - port=to_infos[0]["port"], - ) - sftp = ssh.open_sftp() - - for from_info, to_info, name in zip(from_infos, to_infos, names): - if to_info["scheme"] != "ssh": - raise NotImplementedError - - if from_info["scheme"] != "local": - raise NotImplementedError - - logger.debug( - "Uploading '{}' to '{}/{}'".format( - from_info["path"], to_info["host"], to_info["path"] - ) - ) - - if not name: - name = os.path.basename(from_info["path"]) - - dname = posixpath.dirname(to_info["path"]) - self._exec(ssh, "mkdir -p {}".format(dname)) - - try: - sftp.put( - from_info["path"], - to_info["path"], - callback=create_cb(name), - ) - except Exception: - msg = "failed to upload '{}' to '{}/{}'" - logger.error( - msg.format( - from_info["path"], to_info["host"], to_info["path"] - ) - ) - continue - - progress.finish_target(name) - - sftp.close() - ssh.close() - - def list_cache_paths(self): - ssh = self.ssh(host=self.host, user=self.user, port=self.port) - cmd = "find {} -type f -follow -print".format(self.prefix) - stdout = self._exec(ssh, cmd) - flist = stdout.split() - ssh.close() - return flist diff --git a/dvc/remote/ssh/__init__.py b/dvc/remote/ssh/__init__.py new file mode 100644 index 000000000..5340d7e6b --- /dev/null +++ b/dvc/remote/ssh/__init__.py @@ -0,0 +1,208 @@ +from __future__ import unicode_literals + +import getpass + +try: + import paramiko +except ImportError: + paramiko = None + +import dvc.prompt as prompt +import dvc.logger as logger +from dvc.remote.ssh.connection import SSHConnection +from dvc.config import Config +from dvc.utils.compat import urlparse +from dvc.remote.base import RemoteBase + + +class RemoteSSH(RemoteBase): + scheme = "ssh" + + # NOTE: we support both URL-like (ssh://[user@]host.xz[:port]/path) and + # SCP-like (ssh://[user@]host.xz:/absolute/path) urls. + REGEX = r"^ssh://.*$" + + REQUIRES = {"paramiko": paramiko} + + JOBS = 4 + PARAM_CHECKSUM = "md5" + DEFAULT_PORT = 22 + TIMEOUT = 1800 + + def __init__(self, repo, config): + super(RemoteSSH, self).__init__(repo, config) + self.url = config.get(Config.SECTION_REMOTE_URL, "ssh://") + + parsed = urlparse(self.url) + self.host = parsed.hostname + self.user = ( + config.get(Config.SECTION_REMOTE_USER) + or parsed.username + or getpass.getuser() + ) + self.prefix = parsed.path or "/" + self.port = ( + config.get(Config.SECTION_REMOTE_PORT) + or parsed.port + or self.DEFAULT_PORT + ) + self.keyfile = config.get(Config.SECTION_REMOTE_KEY_FILE, None) + self.timeout = config.get(Config.SECTION_REMOTE_TIMEOUT, self.TIMEOUT) + self.password = config.get(Config.SECTION_REMOTE_PASSWORD, None) + self.ask_password = config.get( + Config.SECTION_REMOTE_ASK_PASSWORD, False + ) + + self.path_info = { + "scheme": "ssh", + "host": self.host, + "user": self.user, + "port": self.port, + } + + def ssh(self, host=None, user=None, port=None, **kwargs): + logger.debug( + "Establishing ssh connection with '{host}' " + "through port '{port}' as user '{user}'".format( + host=host, user=user, port=port + ) + ) + + if self.ask_password and not self.password: + self.password = prompt.password( + "Enter a private key passphrase or a password for " + "host '{host}' port '{port}' user '{user}'".format( + host=host, port=port, user=user + ) + ) + + return SSHConnection( + host, + username=user, + port=port, + key_filename=self.keyfile, + timeout=self.timeout, + password=self.password, + ) + + def exists(self, path_info): + assert not isinstance(path_info, list) + assert path_info["scheme"] == self.scheme + + with self.ssh(**path_info) as ssh: + return ssh.file_exists(path_info["path"]) + + def md5(self, path_info): + if path_info["scheme"] != self.scheme: + raise NotImplementedError + + with self.ssh(**path_info) as ssh: + return ssh.md5(path_info["path"]) + + def copy(self, from_info, to_info): + if ( + from_info["scheme"] != self.scheme + or to_info["scheme"] != self.scheme + ): + raise NotImplementedError + + with self.ssh(**from_info) as ssh: + ssh.cp(from_info["path"], to_info["path"]) + + def save_info(self, path_info): + if path_info["scheme"] != self.scheme: + raise NotImplementedError + + return {self.PARAM_CHECKSUM: self.md5(path_info)} + + def save(self, path_info, checksum_info): + if path_info["scheme"] != self.scheme: + raise NotImplementedError + + md5 = checksum_info[self.PARAM_CHECKSUM] + dest = path_info.copy() + dest["path"] = self.checksum_to_path(md5) + + self.copy(path_info, dest) + + def remove(self, path_info): + if path_info["scheme"] != self.scheme: + raise NotImplementedError + + with self.ssh(**path_info) as ssh: + ssh.remove(path_info["path"]) + + def download( + self, + from_infos, + to_infos, + no_progress_bar=False, + names=None, + resume=False, + ): + names = self._verify_path_args(from_infos, to_infos, names) + ssh = self.ssh(**from_infos[0]) + + for to_info, from_info, name in zip(to_infos, from_infos, names): + if from_info["scheme"] != self.scheme: + raise NotImplementedError + + if to_info["scheme"] == self.scheme: + ssh.cp(from_info["path"], to_info["path"]) + continue + + if to_info["scheme"] != "local": + raise NotImplementedError + + logger.debug( + "Downloading '{host}/{path}' to '{dest}'".format( + host=from_info["host"], + path=from_info["path"], + dest=to_info["path"], + ) + ) + + try: + ssh.download( + from_info["path"], to_info["path"], progress_title=name + ) + except Exception: + logger.error( + "failed to download '{host}/{path}' to '{dest}'".format( + host=from_info["host"], + path=from_info["path"], + dest=to_info["path"], + ) + ) + continue + + ssh.close() + + def upload(self, from_infos, to_infos, names=None): + names = self._verify_path_args(to_infos, from_infos, names) + + with self.ssh(**to_infos[0]) as ssh: + for from_info, to_info, name in zip(from_infos, to_infos, names): + if to_info["scheme"] != self.scheme: + raise NotImplementedError + + if from_info["scheme"] != "local": + raise NotImplementedError + + try: + ssh.upload( + from_info["path"], to_info["path"], progress_title=name + ) + except Exception: + logger.error( + "failed to upload '{host}/{path}' to '{dest}'".format( + host=from_info["host"], + path=from_info["path"], + dest=to_info["path"], + ) + ) + pass + + def list_cache_paths(self): + with self.ssh(**self.path_info) as ssh: + return list(ssh.walk_files(self.prefix)) diff --git a/dvc/remote/ssh/connection.py b/dvc/remote/ssh/connection.py new file mode 100644 index 000000000..10f714f8c --- /dev/null +++ b/dvc/remote/ssh/connection.py @@ -0,0 +1,257 @@ +import os +import posixpath + +try: + import paramiko +except ImportError: + paramiko = None + +import dvc.logger as logger +from dvc.utils.compat import makedirs +from dvc.progress import progress +from dvc.exceptions import DvcException +from dvc.remote.base import RemoteCmdError + + +def sizeof_fmt(num, suffix="B"): + """ Convert number of bytes to human-readable string """ + for unit in ["", "K", "M", "G", "T", "P", "E", "Z"]: + if abs(num) < 1024.0: + return "%3.1f%s%s" % (num, unit, suffix) + num /= 1024.0 + return "%.1f%s%s" % (num, "Y", suffix) + + +def percent_cb(name, complete, total): + """ Callback for updating target progress """ + logger.debug( + "{}: {} transferred out of {}".format( + name, sizeof_fmt(complete), sizeof_fmt(total) + ) + ) + progress.update_target(name, complete, total) + + +def create_cb(name): + """ Create callback function for multipart object """ + return lambda cur, tot: percent_cb(name, cur, tot) + + +class SSHConnection: + def __init__(self, *args, **kwargs): + self.timeout = kwargs.get("timeout", 1800) + + self._ssh = paramiko.SSHClient() + self._ssh.load_system_host_keys() + self._ssh.set_missing_host_key_policy(paramiko.AutoAddPolicy()) + + self._ssh.connect(*args, **kwargs) + self._sftp = None + self._sftp_alive = False + + def __enter__(self): + return self + + def __exit__(self, exception_type, exception_value, traceback): + self.close() + + def _sftp_connect(self): + if not self._sftp or not self._sftp_alive: + self._sftp = self._ssh.open_sftp() + self._sftp_alive = True + + def close(self): + if self._sftp: + self._sftp.close() + self._sftp_alive = False + + self._ssh.close() + + def file_exists(self, path): + self._sftp_connect() + try: + return self._sftp.stat(path) + except IOError: + return False + pass + + def isdir(self, path): + from stat import S_ISDIR + + self._sftp_connect() + try: + return S_ISDIR(self._sftp.stat(path).st_mode) + except IOError: + return False + + def isfile(self, path): + from stat import S_ISREG + + self._sftp_connect() + try: + return S_ISREG(self._sftp.stat(path).st_mode) + except IOError: + return False + + def islink(self, path): + from stat import S_ISLNK + + self._sftp_connect() + try: + return S_ISLNK(self._sftp.stat(path).st_mode) + except IOError: + return False + + def makedirs(self, path): + self._sftp_connect() + + if self.isdir(path): + return + + if self.isfile(path) or self.islink(path): + raise DvcException( + "a file with the same name '{}' already exists".format(path) + ) + + head, tail = os.path.split(path) + + if head: + self.makedirs(head) + + if tail: + self._sftp.mkdir(path) + + def walk_files(self, directory): + from stat import S_ISLNK, S_ISDIR, S_ISREG + + self._sftp_connect() + + for entry in self._sftp.listdir_attr(directory): + path = os.path.join(directory, entry.filename) + + if S_ISLNK(entry.st_mode): + path = self._sftp.readlink(directory) + entry = self._sftp.stat(path) + + if S_ISDIR(entry.st_mode): + for inner_path in self.walk_files(path): + yield inner_path + + elif S_ISREG(entry.st_mode): + yield path + + def remove(self, path): + self._sftp_connect() + self._sftp.remove(path) + + def download(self, src, dest, no_progress_bar=False, progress_title=None): + from uuid import uuid4 + + self._sftp_connect() + + makedirs(posixpath.dirname(dest), exist_ok=True) + tmp_file = dest + "." + str(uuid4()) + + if no_progress_bar: + self._sftp.get(src, tmp_file) + else: + if not progress_title: + progress_title = os.path.basename(dest) + + self._sftp.get(src, tmp_file, callback=create_cb(progress_title)) + progress.finish_target(progress_title) + + os.rename(tmp_file, dest) + + def upload(self, src, dest, no_progress_bar=False, progress_title=None): + self._sftp_connect() + + self.makedirs(posixpath.dirname(dest)) + + if no_progress_bar: + self._sftp.put(src, dest) + else: + if not progress_title: + progress_title = os.path.basename(dest) + + self._sftp.put(src, dest, callback=create_cb(progress_title)) + progress.finish_target(progress_title) + + def execute(self, cmd): + stdin, stdout, stderr = self._ssh.exec_command(cmd) + channel = stdout.channel + + stdin.close() + channel.shutdown_write() + + stdout_chunks = [] + stderr_chunks = [] + while ( + not channel.closed + or channel.recv_ready() + or channel.recv_stderr_ready() + ): + import select + + got_chunk = False + readq, _, _ = select.select([stdout.channel], [], [], self.timeout) + for c in readq: + if c.recv_ready(): + stdout_chunks.append(stdout.channel.recv(len(c.in_buffer))) + got_chunk = True + + if c.recv_stderr_ready(): + stderr_len = len(c.in_stderr_buffer) + s = stderr.channel.recv_stderr(stderr_len) + stderr_chunks.append(s) + got_chunk = True + + if ( + not got_chunk + and stdout.channel.exit_status_ready() + and not stderr.channel.recv_stderr_ready() + and not stdout.channel.recv_ready() + ): + stdout.channel.shutdown_read() + stdout.channel.close() + break + + stdout.close() + stderr.close() + + ret = stdout.channel.recv_exit_status() + if ret != 0: + err = b"".join(stderr_chunks).decode("utf-8") + raise RemoteCmdError("ssh", cmd, ret, err) + + return b"".join(stdout_chunks).decode("utf-8") + + def md5(self, path): + """ + Use different md5 commands depending on the OS: + + - Darwin's `md5` returns BSD-style checksums by default + - Linux's `md5sum` needs the `--tag` flag for a similar output + + Example: + MD5 (foo.txt) = f3d220a856b52aabbf294351e8a24300 + """ + uname = self.execute("uname").strip() + + command = { + "Darwin": "md5 {}".format(path), + "Linux": "md5sum --tag {}".format(path), + }.get(uname) + + if not command: + raise DvcException( + "'{uname}' is not supported as a remote".format(uname=uname) + ) + + md5 = self.execute(command).split()[-1] + assert len(md5) == 32 + return md5 + + def cp(self, src, dest): + self.makedirs(posixpath.dirname(dest)) + self.execute("cp {} {}".format(src, dest)) diff --git a/dvc/utils/compat.py b/dvc/utils/compat.py index 011928a90..da16b3afc 100644 --- a/dvc/utils/compat.py +++ b/dvc/utils/compat.py @@ -1,6 +1,7 @@ """Handle import compatibility between Python 2 and Python 3""" import sys +import os # Syntax sugar. @@ -13,6 +14,7 @@ is_py2 = _ver[0] == 2 is_py3 = _ver[0] == 3 # NOTE: cast_bytes_py2 is taken from https://github.com/ipython/ipython_genutils +# NOTE _makedirs is taken from https://github.com/python/cpython/blob/3ce3dea60646d8a5a1c952469a2eb65f937875b3/Lib/os.py#L196-L226 # simplified version of ipython_genutils/encoding.py DEFAULT_ENCODING = sys.getdefaultencoding() @@ -33,6 +35,27 @@ def cast_bytes(s, encoding=None): return s +def _makedirs(name, mode=0o777, exist_ok=False): + head, tail = os.path.split(name) + if not tail: + head, tail = os.path.split(head) + if head and tail and not os.path.exists(head): + try: + _makedirs(head, exist_ok=exist_ok) + except FileExistsError: + pass + cdir = os.curdir + if isinstance(tail, bytes): + cdir = bytes(os.curdir, "ASCII") + if tail == cdir: + return + try: + os.mkdir(name, mode) + except OSError: + if not exist_ok or not os.path.isdir(name): + raise + + if is_py2: from urlparse import urlparse, urljoin # noqa: F401 from StringIO import StringIO # noqa: F401 @@ -49,8 +72,10 @@ if is_py2: integer_types = (int, long) # noqa: F821 input = raw_input # noqa: F821 cast_bytes_py2 = cast_bytes + makedirs = _makedirs elif is_py3: + from os import makedirs # noqa: F401 from urllib.parse import urlparse, urljoin # noqa: F401 from io import StringIO # noqa: F401 from http.server import ( # noqa: F401 diff --git a/tests/test_data_cloud.py b/tests/test_data_cloud.py index d7208ba96..66fdc5205 100644 --- a/tests/test_data_cloud.py +++ b/tests/test_data_cloud.py @@ -672,8 +672,6 @@ class TestRecursiveSyncOperations(TestDataCloudBase): self.main( ["remote", "add", "-d", TEST_REMOTE, self.cloud._cloud.cache_dir] ) - self.create(self.DATA, self.DATA_CONTENTS) - self.create(self.DATA_SUB, self.DATA_SUB_CONTENTS) self.dvc.add(self.DATA) self.dvc.add(self.DATA_SUB_CONTENTS) diff --git a/tests/unit/remote/ssh.py b/tests/unit/remote/ssh.py index 30a646a3e..9afd9fe7d 100644 --- a/tests/unit/remote/ssh.py +++ b/tests/unit/remote/ssh.py @@ -40,7 +40,7 @@ class TestRemoteSSH(TestCase): self.assertEqual(remote.port, port) self.assertEqual(remote.prefix, path) - # SCP-like URL ssh://[user@]host.xz:path + # SCP-like URL ssh://[user@]host.xz:/absolute/path url = "ssh://{}@{}:{}".format(user, host, path) config = {"url": url}