mirror of
https://github.com/Stone-Red-Code/dvc.git
synced 2026-11-02 07:18:25 +01:00
Merge pull request #1610 from mroutis/fix-830
ssh: Use SFTP when possible
This commit is contained in:
Ruslan Kuprieiev
+2
-2
@@ -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(
|
||||
|
||||
@@ -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
|
||||
|
||||
|
||||
+2
-2
@@ -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"])
|
||||
|
||||
+2
-2
@@ -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"])
|
||||
|
||||
|
||||
+3
-3
@@ -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(
|
||||
|
||||
+2
-2
@@ -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:
|
||||
|
||||
@@ -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
|
||||
@@ -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))
|
||||
@@ -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))
|
||||
@@ -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
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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}
|
||||
|
||||
|
||||
Reference in New Issue
Block a user