Stone_Red/dvc
1
0
Fork 0
mirror of https://github.com/Stone-Red-Code/dvc.git synced 2026-11-02 07:18:25 +01:00
Code Issues Packages Projects Releases Wiki Activity

ssh: migrate to sshfs (#6295)

Browse Source
This commit is contained in:
Batuhan Taskaya
2021-07-28 12:49:53 +03:00
committed by GitHub
parent b294bca383
commit 9b4b6e5b31
15 changed files with 366 additions and 1197 deletions
Download Patch File Download Diff File Expand all files Collapse all files
dvc/fs/fsspec_wrapper.py
+11 -2
View File
@@ -29,6 +29,15 @@ class FSSpecWrapper(BaseFileSystem):
def _strip_bucket(self, entry):
return entry
@staticmethod
def _get_kwargs_from_urls(urlpath):
from fsspec.utils import infer_storage_options
options = infer_storage_options(urlpath)
options.pop("path", None)
options.pop("protocol", None)
return options
def _strip_buckets(self, entries, detail=False):
for entry in entries:
if detail:
@@ -236,7 +245,7 @@ class CallbackMixin:
**pbar_args,
) as pbar:
self.fs.put_file(
self._with_bucket(from_file),
os.fspath(from_file),
self._with_bucket(to_info),
callback=pbar.as_callback(_LOCAL_FS, from_file),
)
@@ -254,6 +263,6 @@ class CallbackMixin:
) as pbar:
self.fs.get_file(
self._with_bucket(from_info),
self._with_bucket(to_file),
os.fspath(to_file),
callback=pbar.as_callback(self, from_info),
)
dvc/fs/ssh.py
+125
View File
@@ -0,0 +1,125 @@
import getpass
import os.path
import threading
from funcy import cached_property, first, memoize, silent, wrap_prop, wrap_with
import dvc.prompt as prompt
from dvc.scheme import Schemes
from dvc.utils.fs import as_atomic
from .fsspec_wrapper import CallbackMixin, FSSpecWrapper
_SSH_TIMEOUT = 60 * 30
_SSH_CONFIG_FILE = os.path.expanduser(os.path.join("~", ".ssh", "config"))
@wrap_with(threading.Lock())
@memoize
def ask_password(host, user, port):
return prompt.password(
"Enter a private key passphrase or a password for "
"host '{host}' port '{port}' user '{user}'".format(
host=host, port=port, user=user
)
)
# pylint:disable=abstract-method
class SSHFileSystem(CallbackMixin, FSSpecWrapper):
scheme = Schemes.SSH
REQUIRES = {"sshfs": "sshfs"}
DEFAULT_PORT = 22
PARAM_CHECKSUM = "md5"
def _with_bucket(self, path):
if isinstance(path, self.PATH_CLS):
path_str = path.path
else:
path_str = super()._with_bucket(path)
return os.path.normpath(path_str)
def _prepare_credentials(self, **config):
self.CAN_TRAVERSE = True
from sshfs.config import parse_config
login_info = {}
try:
user_ssh_config = parse_config(host=config["host"])
except FileNotFoundError:
user_ssh_config = {}
login_info["host"] = user_ssh_config.get("Hostname", config["host"])
login_info["username"] = (
config.get("user")
or config.get("username")
or user_ssh_config.get("User")
or getpass.getuser()
)
login_info["port"] = (
config.get("port")
or silent(int)(user_ssh_config.get("Port"))
or self.DEFAULT_PORT
)
login_info["password"] = config.get("password")
if user_ssh_config.get("IdentityFile"):
config.setdefault(
"keyfile", first(user_ssh_config.get("IdentityFile"))
)
keyfile = config.get("keyfile")
if keyfile:
login_info["client_keys"] = [os.path.expanduser(keyfile)]
login_info["timeout"] = config.get("timeout", _SSH_TIMEOUT)
# These two settings fine tune the asyncssh to use the
# fastest encryption algorithm and disable compression
# altogether (since it blocking, it is slowing down
# the transfers in a considerable rate, and even for
# compressible data it is making it extremely slow).
# See: https://github.com/ronf/asyncssh/issues/374
login_info["encryption_algs"] = [
"[email protected]",
"aes256-ctr",
"aes192-ctr",
"aes128-ctr",
]
login_info["compression_algs"] = None
login_info["gss_auth"] = config.get("gss_auth", False)
login_info["agent_forwarding"] = config.get("agent_forwarding", True)
login_info["proxy_command"] = user_ssh_config.get("ProxyCommand")
if config.get("ask_password") and login_info["password"] is None:
login_info["password"] = ask_password(
login_info["host"], login_info["username"], login_info["port"]
)
# We are going to automatically add stuff to known_hosts
# something like paramiko's AutoAddPolicy()
login_info["known_hosts"] = None
return login_info
@wrap_prop(threading.Lock())
@cached_property
def fs(self):
from sshfs import SSHFileSystem as _SSHFileSystem
return _SSHFileSystem(**self.fs_args)
# Ensure that if an interrupt happens during the transfer, we don't
# pollute the cache.
def _upload_fobj(self, fobj, to_info, *args, **kwargs):
with as_atomic(self, to_info) as tmp_file:
super()._upload_fobj(fobj, tmp_file, *args, **kwargs)
def _upload(self, from_file, to_info, *args, **kwargs):
with as_atomic(self, to_info) as tmp_file:
super()._upload(from_file, tmp_file, *args, **kwargs)
dvc/fs/ssh/__init__.py
-263
View File
@@ -1,263 +0,0 @@
import getpass
import io
import logging
import os
import shutil
import threading
from contextlib import closing, contextmanager
from funcy import first, memoize, silent, wrap_with
from dvc import prompt
from dvc.hash_info import HashInfo
from dvc.scheme import Schemes
from ..base import BaseFileSystem
from ..pool import get_connection
logger = logging.getLogger(__name__)
@wrap_with(threading.Lock())
@memoize
def ask_password(host, user, port):
return prompt.password(
"Enter a private key passphrase or a password for "
"host '{host}' port '{port}' user '{user}'".format(
host=host, port=port, user=user
)
)
class SSHFileSystem(BaseFileSystem): # pylint:disable=abstract-method
scheme = Schemes.SSH
REQUIRES = {"paramiko": "paramiko"}
_JOBS = 4
PARAM_CHECKSUM = "md5"
# At any given time some of the connections will go over network and
# paramiko stuff, so we would ideally have it double of server processors.
# We use conservative setting of 4 instead to not exhaust max sessions.
CHECKSUM_JOBS = 4
DEFAULT_CACHE_TYPES = ["copy"]
TRAVERSE_PREFIX_LEN = 2
DEFAULT_PORT = 22
TIMEOUT = 1800
def __init__(self, **config):
super().__init__(**config)
user_ssh_config = self._load_user_ssh_config(config["host"])
self.host = user_ssh_config.get("hostname", config["host"])
self.user = (
config.get("user")
or user_ssh_config.get("user")
or getpass.getuser()
)
self.port = (
config.get("port")
or self._try_get_ssh_config_port(user_ssh_config)
or self.DEFAULT_PORT
)
self.keyfile = config.get(
"keyfile"
) or self._try_get_ssh_config_keyfile(user_ssh_config)
self.timeout = config.get("timeout", self.TIMEOUT)
self.password = config.get("password", None)
self.ask_password = config.get("ask_password", False)
self.gss_auth = config.get("gss_auth", False)
proxy_command = user_ssh_config.get("proxycommand", False)
if proxy_command:
import paramiko
self.sock = paramiko.ProxyCommand(proxy_command)
else:
self.sock = None
self.allow_agent = config.get("allow_agent", True)
@staticmethod
def _get_kwargs_from_urls(urlpath):
from fsspec.implementations.sftp import SFTPFileSystem
# pylint:disable=protected-access
kwargs = SFTPFileSystem._get_kwargs_from_urls(urlpath)
if "username" in kwargs:
kwargs["user"] = kwargs.pop("username")
return kwargs
@staticmethod
def ssh_config_filename():
return os.path.expanduser(os.path.join("~", ".ssh", "config"))
@staticmethod
def _load_user_ssh_config(hostname):
import paramiko
user_config_file = SSHFileSystem.ssh_config_filename()
user_ssh_config = {}
if hostname and os.path.exists(user_config_file):
ssh_config = paramiko.SSHConfig()
with open(user_config_file) as f:
# For whatever reason parsing directly from f is unreliable
f_copy = io.StringIO(f.read())
ssh_config.parse(f_copy)
user_ssh_config = ssh_config.lookup(hostname)
return user_ssh_config
@staticmethod
def _try_get_ssh_config_port(user_ssh_config):
return silent(int)(user_ssh_config.get("port"))
@staticmethod
def _try_get_ssh_config_keyfile(user_ssh_config):
return first(user_ssh_config.get("identityfile") or ())
def ensure_credentials(self):
# NOTE: we use the same password regardless of the server :(
if self.ask_password and self.password is None:
self.password = ask_password(self.host, self.user, self.port)
def ssh(self):
self.ensure_credentials()
from .connection import SSHConnection
return get_connection(
SSHConnection,
self.host,
username=self.user,
port=self.port,
key_filename=self.keyfile,
timeout=self.timeout,
password=self.password,
gss_auth=self.gss_auth,
sock=self.sock,
allow_agent=self.allow_agent,
)
@contextmanager
def open(self, path_info, mode="r", encoding=None, **kwargs):
assert mode in {"r", "rt", "rb", "wb"}
with self.ssh() as ssh, closing(
ssh.sftp.open(path_info.path, mode)
) as fd:
if "b" in mode:
yield fd
else:
yield io.TextIOWrapper(fd, encoding=encoding)
def exists(self, path_info) -> bool:
with self.ssh() as ssh:
return ssh.exists(path_info.path)
def isdir(self, path_info):
with self.ssh() as ssh:
return ssh.isdir(path_info.path)
def isfile(self, path_info):
with self.ssh() as ssh:
return ssh.isfile(path_info.path)
def walk_files(self, path_info, **kwargs):
with self.ssh() as ssh:
for fname in ssh.walk_files(path_info.path):
yield path_info.replace(path=fname)
def remove(self, path_info):
if path_info.scheme != self.scheme:
raise NotImplementedError
with self.ssh() as ssh:
ssh.remove(path_info.path)
def makedirs(self, path_info, **kwargs):
with self.ssh() as ssh:
ssh.makedirs(path_info.path)
def move(self, from_info, to_info):
if from_info.scheme != self.scheme or to_info.scheme != self.scheme:
raise NotImplementedError
with self.ssh() as ssh:
ssh.move(from_info.path, to_info.path)
def copy(self, from_info, to_info):
if not from_info.scheme == to_info.scheme == self.scheme:
raise NotImplementedError
with self.ssh() as ssh:
ssh.atomic_copy(from_info.path, to_info.path)
def symlink(self, from_info, to_info):
if not from_info.scheme == to_info.scheme == self.scheme:
raise NotImplementedError
with self.ssh() as ssh:
ssh.symlink(from_info.path, to_info.path)
def hardlink(self, from_info, to_info):
if not from_info.scheme == to_info.scheme == self.scheme:
raise NotImplementedError
# See dvc/remote/local/__init__.py - hardlink()
if self.getsize(from_info) == 0:
with self.ssh() as ssh:
ssh.sftp.open(to_info.path, "w").close()
logger.debug(
"Created empty file: {src} -> {dest}".format(
src=str(from_info), dest=str(to_info)
)
)
return
with self.ssh() as ssh:
ssh.hardlink(from_info.path, to_info.path)
def reflink(self, from_info, to_info):
if from_info.scheme != self.scheme or to_info.scheme != self.scheme:
raise NotImplementedError
with self.ssh() as ssh:
ssh.reflink(from_info.path, to_info.path)
def md5(self, path_info):
with self.ssh() as ssh:
return HashInfo(
"md5",
ssh.md5(path_info.path),
size=ssh.getsize(path_info.path),
)
def info(self, path_info):
with self.ssh() as ssh:
return ssh.info(path_info.path)
def _upload_fobj(self, fobj, to_info, **kwargs):
self.makedirs(to_info.parent)
with self.open(to_info, mode="wb") as fdest:
shutil.copyfileobj(fobj, fdest)
def _download(self, from_info, to_file, name=None, no_progress_bar=False):
with self.ssh() as ssh:
ssh.download(
from_info.path,
to_file,
progress_title=name,
no_progress_bar=no_progress_bar,
)
def _upload(
self, from_file, to_info, name=None, no_progress_bar=False, **_kwargs
):
with self.ssh() as ssh:
ssh.upload(
from_file,
to_info.path,
progress_title=name,
no_progress_bar=no_progress_bar,
)
dvc/fs/ssh/connection.py
-356
View File
@@ -1,356 +0,0 @@
import errno
import logging
import os
import posixpath
import shlex
import stat
from contextlib import suppress
from funcy import cached_property
from dvc.exceptions import DvcException
from dvc.progress import Tqdm
from dvc.utils import tmp_fname
from ..base import RemoteCmdError
try:
import paramiko
except ImportError:
paramiko = None # type: ignore
logger = logging.getLogger(__name__)
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 f"{num:3.1f}{unit}{suffix}"
num /= 1024.0
return "{:.1f}{}{}".format(num, "Y", suffix)
class SSHConnection:
def __init__(self, host, *args, **kwargs):
logger.debug(
"Establishing ssh connection with '{host}' "
"through port '{port}' as user '{username}'".format(
host=host, **kwargs
)
)
kwargs.setdefault("timeout", 1800)
self.timeout = kwargs["timeout"]
self._ssh = paramiko.SSHClient()
# Explicitly disable paramiko logger. Due to how paramiko dynamically
# loads loggers, it is not disabled by DVC disable_other_loggers().
# See https://github.com/iterative/dvc/issues/3482
self._ssh.set_log_channel("dvc.paramiko")
logging.getLogger("dvc.paramiko").disabled = True
logging.getLogger("dvc.paramiko.sftp").disabled = True
self._ssh.load_system_host_keys()
self._ssh.set_missing_host_key_policy(paramiko.AutoAddPolicy())
self._ssh.connect(host, *args, **kwargs)
self._ssh.get_transport().set_keepalive(10)
self._sftp_channels = []
@property
def sftp(self):
if not self._sftp_channels:
self._sftp_channels = [self._ssh.open_sftp()]
return self._sftp_channels[0]
def close(self):
for sftp in self._sftp_channels:
sftp.close()
self._ssh.close()
def st_mode(self, path):
lstat = None
with suppress(FileNotFoundError):
lstat = self.sftp.lstat(path)
return lstat.st_mode if lstat else 0
def info(self, path):
st = self.sftp.stat(path)
return {
"size": st.st_size,
"type": "directory" if stat.S_ISDIR(st.st_mode) else "file",
}
def getsize(self, path):
return self.info(path)["size"]
def exists(self, path):
return bool(self.st_mode(path))
def isdir(self, path):
return stat.S_ISDIR(self.st_mode(path))
def isfile(self, path):
return stat.S_ISREG(self.st_mode(path))
def islink(self, path):
return stat.S_ISLNK(self.st_mode(path))
def makedirs(self, path):
# Single stat call will say whether this is a dir, a file or a link
st_mode = self.st_mode(path)
if stat.S_ISDIR(st_mode):
return
if stat.S_ISREG(st_mode) or stat.S_ISLNK(st_mode):
raise DvcException(
f"a file with the same name '{path}' already exists"
)
head, tail = posixpath.split(path)
if head:
self.makedirs(head)
if tail:
try:
self.sftp.mkdir(path)
except OSError as exc:
# Since paramiko errors are very vague we need to recheck
# whether it's because path already exists or something else
if exc.errno == errno.EACCES or not self.exists(path):
raise DvcException(
f"unable to create remote directory '{path}'"
) from exc
def walk(self, directory, topdown=True):
# NOTE: original os.walk() implementation [1] with default options was
# used as a template.
#
# [1] https://github.com/python/cpython/blob/master/Lib/os.py
try:
dir_entries = self.sftp.listdir_attr(directory)
except OSError as exc:
raise DvcException(
"couldn't get the '{}' remote directory files list".format(
directory
)
) from exc
dirs = []
nondirs = []
for entry in dir_entries:
name = entry.filename
if stat.S_ISDIR(entry.st_mode):
dirs.append(name)
else:
nondirs.append(name)
if topdown:
yield directory, dirs, nondirs
for dname in dirs:
newpath = posixpath.join(directory, dname)
yield from self.walk(newpath, topdown=topdown)
if not topdown:
yield directory, dirs, nondirs
def walk_files(self, directory):
for root, _, files in self.walk(directory):
for fname in files:
yield posixpath.join(root, fname)
def _remove_file(self, path):
with suppress(FileNotFoundError):
self.sftp.remove(path)
def _remove_dir(self, path):
for root, dirs, files in self.walk(path, topdown=False):
for fname in files:
with suppress(FileNotFoundError):
self._remove_file(posixpath.join(root, fname))
for dname in dirs:
with suppress(FileNotFoundError):
self.sftp.rmdir(posixpath.join(root, dname))
with suppress(FileNotFoundError):
self.sftp.rmdir(path)
def remove(self, path):
if self.isdir(path):
self._remove_dir(path)
else:
self._remove_file(path)
def download(self, src, dest, no_progress_bar=False, progress_title=None):
with Tqdm(
desc=progress_title or os.path.basename(src),
disable=no_progress_bar,
bytes=True,
total=self.getsize(src),
) as pbar:
self.sftp.get(src, dest, callback=pbar.update_to)
def move(self, src, dst):
"""Rename src to dst, if it is not possible (in case src and dst are
on different filesystems) and actual physical copying of data is
happening.
"""
self.makedirs(posixpath.dirname(dst))
try:
self.sftp.rename(src, dst)
except OSError:
self.atomic_copy(src, dst)
self.remove(src)
def atomic_copy(self, src, dst):
tmp = tmp_fname(dst)
try:
self.copy(src, tmp)
self.sftp.rename(tmp, dst)
finally:
self.remove(tmp)
def upload(self, src, dest, no_progress_bar=False, progress_title=None):
self.makedirs(posixpath.dirname(dest))
tmp_file = tmp_fname(dest)
if not progress_title:
progress_title = posixpath.basename(dest)
with Tqdm(
desc=progress_title,
disable=no_progress_bar,
bytes=True,
total=os.path.getsize(src),
) as pbar:
self.sftp.put(src, tmp_file, callback=pbar.update_to)
self.sftp.rename(tmp_file, dest)
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 selectors
selector = selectors.DefaultSelector()
selector.register(stdout.channel, selectors.EVENT_READ)
got_chunk = False
events = selector.select(self.timeout)
for key, _ in events:
c = key.fileobj
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")
@cached_property
def uname(self):
return self.execute("uname").strip()
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
"""
path = shlex.quote(path)
if self.uname == "Linux":
md5 = self.execute("md5sum " + path).split()[0]
elif self.uname == "Darwin":
md5 = self.execute("md5 " + path).split()[-1]
else:
raise DvcException(
f"'{self.uname}' is not supported as a SSH remote"
)
assert len(md5) == 32
return md5
def copy(self, src, dest):
dest = shlex.quote(dest)
src = shlex.quote(src)
self.execute(f"cp {src} {dest}")
def open_max_sftp_channels(self):
# If there are more than 1 it means we've already opened max amount
if len(self._sftp_channels) <= 1:
while True:
try:
self._sftp_channels.append(self._ssh.open_sftp())
except paramiko.ssh_exception.ChannelException:
if not self._sftp_channels:
raise
break
return self._sftp_channels
def open(self, *args, **kwargs):
return self.sftp.open(*args, **kwargs)
def symlink(self, src, dest):
self.sftp.symlink(src, dest)
def reflink(self, src, dest):
dest = shlex.quote(dest)
src = shlex.quote(src)
if self.uname == "Linux":
return self.execute(f"cp --reflink {src} {dest}")
if self.uname == "Darwin":
return self.execute(f"cp -c {src} {dest}")
raise DvcException(f"'{self.uname}' is not supported as a SSH remote")
def hardlink(self, src, dest):
dest = shlex.quote(dest)
src = shlex.quote(src)
self.execute(f"ln {src} {dest}")
dvc/objects/db/__init__.py
-4
View File
@@ -11,14 +11,10 @@ def get_odb(fs, path_info, **config):
from .gdrive import GDriveObjectDB
from .local import LocalObjectDB
from .oss import OSSObjectDB
from .ssh import SSHObjectDB
if fs.scheme == Schemes.LOCAL:
return LocalObjectDB(fs, path_info, **config)
if fs.scheme == Schemes.SSH:
return SSHObjectDB(fs, path_info, **config)
if fs.scheme == Schemes.GDRIVE:
return GDriveObjectDB(fs, path_info, **config)
dvc/objects/db/ssh.py
-90
View File
@@ -1,90 +0,0 @@
import errno
import itertools
import logging
import posixpath
from concurrent.futures import ThreadPoolExecutor
from dvc.progress import Tqdm
from dvc.utils import to_chunks
from .base import ObjectDB
logger = logging.getLogger(__name__)
class SSHObjectDB(ObjectDB):
def batch_exists(self, path_infos, callback):
def _exists(chunk_and_channel):
chunk, channel = chunk_and_channel
ret = []
for path in chunk:
try:
channel.stat(path)
ret.append(True)
except OSError as exc:
if exc.errno != errno.ENOENT:
raise
ret.append(False)
callback(path)
return ret
with self.fs.ssh() as ssh:
channels = ssh.open_max_sftp_channels()
max_workers = len(channels)
with ThreadPoolExecutor(max_workers=max_workers) as executor:
paths = [path_info.path for path_info in path_infos]
chunks = to_chunks(paths, num_chunks=max_workers)
chunks_and_channels = zip(chunks, channels)
outcome = executor.map(_exists, chunks_and_channels)
results = list(itertools.chain.from_iterable(outcome))
return results
def hashes_exist(self, hashes, jobs=None, name=None):
"""This is older implementation used in remote/base.py
We are reusing it in RemoteSSH, because SSH's batch_exists proved to be
faster than current approach (relying on exists(path_info)) applied in
remote/base.
"""
if not self.fs.CAN_TRAVERSE:
return list(set(hashes) & set(self.all()))
# possibly prompt for credentials before "Querying" progress output
self.fs.ensure_credentials()
with Tqdm(
desc="Querying "
+ ("cache in " + name if name else "remote cache"),
total=len(hashes),
unit="file",
) as pbar:
def exists_with_progress(chunks):
return self.batch_exists(chunks, callback=pbar.update_msg)
with ThreadPoolExecutor(
max_workers=jobs or self.fs.jobs
) as executor:
path_infos = [self.hash_to_path_info(x) for x in hashes]
chunks = to_chunks(path_infos, num_chunks=self.fs.jobs)
results = executor.map(exists_with_progress, chunks)
in_remote = itertools.chain.from_iterable(results)
ret = list(itertools.compress(hashes, in_remote))
return ret
def _list_paths(self, prefix=None, progress_callback=None):
if prefix:
root = posixpath.join(self.path_info.path, prefix[:2])
else:
root = self.path_info.path
with self.fs.ssh() as ssh:
if not ssh.exists(root):
return
# If we simply return an iterator then with above closes instantly
if progress_callback:
for path in ssh.walk_files(root):
progress_callback()
yield path
else:
yield from ssh.walk_files(root)
dvc/utils/fs.py
+19
View File
@@ -4,6 +4,7 @@ import os
import shutil
import stat
import sys
from contextlib import contextmanager, suppress
from typing import TYPE_CHECKING
from dvc.exceptions import DvcException
@@ -233,3 +234,21 @@ def walk_files(directory):
for root, _, files in os.walk(directory):
for f in files:
yield os.path.join(root, f)
@contextmanager
def as_atomic(fs, to_info):
from dvc.utils import tmp_fname
tmp_info = to_info.parent / tmp_fname()
try:
yield tmp_info
except BaseException:
# Handle stuff like KeyboardInterrupt
# as well as other errors that might
# arise during file transfer.
with suppress(FileNotFoundError):
fs.remove(tmp_info)
raise
else:
fs.move(tmp_info, to_info)
setup.py
+2 -2
View File
@@ -101,7 +101,7 @@ s3 = ["s3fs==2021.7.0", "aiobotocore[boto3]==1.3.0"]
azure = ["adlfs==2021.7.1", "azure-identity>=1.4.0", "knack"]
# https://github.com/Legrandin/pycryptodome/issues/465
oss = ["ossfs==2021.7.5"]
ssh = ["paramiko[invoke]>=2.7.0"]
ssh = ["sshfs>=2021.7.0"]
# Remove the env marker if/when pyarrow is available for Python3.9
hdfs = ["pyarrow>=2.0.0"]
@@ -111,7 +111,7 @@ webdav = ["webdav4>=0.9.0"]
# for linux and mac, so it will fail to compile if user doesn't have all the
# requirements, including kerberos itself. Once all the wheels are available,
# we can start shipping it by default.
ssh_gssapi = ["paramiko[invoke,gssapi]>=2.7.0"]
ssh_gssapi = ["sshfs[gssapi]>=2021.7.0"]
all_remotes = gs + s3 + azure + ssh + oss + gdrive + hdfs + webhdfs + webdav
tests_requirements = (
tests/remotes/ssh.py
+53 -31
View File
@@ -1,9 +1,11 @@
import locale
import os
from contextlib import contextmanager
import subprocess
import mockssh
import pytest
from funcy import cached_property
from mockssh.server import Handler
from dvc.path_info import URLInfo
@@ -16,6 +18,40 @@ TEST_SSH_KEY_PATH = os.path.join(
)
# See: https://github.com/carletes/mock-ssh-server/issues/22
class SSHMockHandler(Handler):
def handle_client(self, channel):
try:
command = self.command_queues[channel.chanid].get(block=True)
self.log.debug("Executing %s", command)
if isinstance(command, bytes):
command = command.decode()
p = subprocess.Popen(
command,
shell=True,
stdin=subprocess.PIPE,
stdout=subprocess.PIPE,
stderr=subprocess.PIPE,
)
stdout, stderr = p.communicate()
channel.sendall(stdout)
channel.sendall_stderr(stderr)
channel.send_exit_status(p.returncode)
except Exception: # pylint: disable=broad-except
self.log.error(
"Error handling client (channel: %s)", channel, exc_info=True
)
finally:
try:
channel.close()
except EOFError:
self.log.debug("Tried to close already closed channel")
class SSHMockServer(mockssh.Server):
handler_cls = SSHMockHandler
class SSHMocked(Base, URLInfo):
@staticmethod
def get_url(user, port): # pylint: disable=arguments-differ
@@ -48,52 +84,40 @@ class SSHMocked(Base, URLInfo):
def config(self):
return {"url": self.url, "keyfile": TEST_SSH_KEY_PATH}
@contextmanager
@cached_property
def _ssh(self):
from dvc.fs.ssh.connection import SSHConnection
from sshfs import SSHFileSystem
conn = SSHConnection(
return SSHFileSystem(
host=self.host,
port=self.port,
username=TEST_SSH_USER,
key_filename=TEST_SSH_KEY_PATH,
client_keys=[TEST_SSH_KEY_PATH],
)
try:
yield conn
finally:
conn.close()
def is_file(self):
with self._ssh() as _ssh:
return _ssh.isfile(self.path)
return self._ssh.isfile(self.path)
def is_dir(self):
with self._ssh() as _ssh:
return _ssh.isdir(self.path)
return self._ssh.isdir(self.path)
def exists(self):
with self._ssh() as _ssh:
return _ssh.exists(self.path)
return self._ssh.exists(self.path)
def mkdir(self, mode=0o777, parents=False, exist_ok=False):
assert mode == 0o777
assert parents
with self._ssh() as _ssh:
_ssh.makedirs(self.path)
self._ssh.makedirs(self.path, exist_ok=exist_ok)
def write_bytes(self, contents):
assert isinstance(contents, bytes)
with self._ssh() as _ssh:
with _ssh.open(self.path, "w+") as fobj:
# NOTE: accepts both str and bytes
fobj.write(contents)
with self._ssh.open(self.path, "wb") as fobj:
fobj.write(contents)
def read_bytes(self):
with self._ssh() as _ssh:
# NOTE: sftp always reads in binary format
with _ssh.open(self.path, "r") as fobj:
return fobj.read()
with self._ssh.open(self.path, "rb") as fobj:
return fobj.read()
def read_text(self, encoding=None, errors=None):
if not encoding:
@@ -104,23 +128,21 @@ class SSHMocked(Base, URLInfo):
@pytest.fixture
def ssh_server(test_config):
import mockssh
test_config.requires("ssh")
users = {TEST_SSH_USER: TEST_SSH_KEY_PATH}
with mockssh.Server(users) as s:
with SSHMockServer(users) as s:
yield s
@pytest.fixture
def ssh_connection(ssh_server):
from dvc.fs.ssh.connection import SSHConnection
from sshfs import SSHFileSystem
yield SSHConnection(
yield SSHFileSystem(
host=ssh_server.host,
port=ssh_server.port,
username=TEST_SSH_USER,
key_filename=TEST_SSH_KEY_PATH,
client_files=[TEST_SSH_KEY_PATH],
)
tests/unit/fs/test_ssh.py
+153
View File
@@ -0,0 +1,153 @@
# pylint: disable=W0613
import getpass
import os
import pytest
from mock import mock_open, patch
from dvc.fs.ssh import SSHFileSystem
def test_get_kwargs_from_urls():
user = "test"
host = "123.45.67.89"
port = 1234
path = "/path/to/dir"
# URL ssh://[user@]host.xz[:port]/path
url = f"ssh://{user}@{host}:{port}{path}"
assert SSHFileSystem._get_kwargs_from_urls(url) == {
"username": user,
"host": host,
"port": port,
}
# SCP-like URL ssh://[user@]host.xz:/absolute/path
url = f"ssh://{user}@{host}:{path}"
assert SSHFileSystem._get_kwargs_from_urls(url) == {
"username": user,
"host": host,
}
def test_init():
fs = SSHFileSystem(user="test", host="12.34.56.78", port="1234")
assert fs.fs_args["username"] == "test"
assert fs.fs_args["host"] == "12.34.56.78"
assert fs.fs_args["port"] == "1234"
mock_ssh_config = """
Host example.com
User ubuntu
HostName 1.2.3.4
Port 1234
IdentityFile ~/.ssh/not_default.key
"""
@pytest.mark.parametrize(
"config,expected_host",
[
({"host": "example.com"}, "1.2.3.4"),
({"host": "not_in_ssh_config.com"}, "not_in_ssh_config.com"),
],
)
@patch(
"builtins.open",
new_callable=mock_open,
read_data=mock_ssh_config,
)
def test_ssh_host_override_from_config(mock_file, config, expected_host):
fs = SSHFileSystem(**config)
assert fs.fs_args["host"] == expected_host
@pytest.mark.parametrize(
"config,expected_user",
[
({"host": "example.com", "user": "test1"}, "test1"),
({"host": "example.com"}, "ubuntu"),
({"host": "not_in_ssh_config.com", "user": "test1"}, "test1"),
({"host": "not_in_ssh_config.com"}, getpass.getuser()),
],
)
@patch(
"builtins.open",
new_callable=mock_open,
read_data=mock_ssh_config,
)
def test_ssh_user(mock_file, config, expected_user):
fs = SSHFileSystem(**config)
assert fs.fs_args["username"] == expected_user
@pytest.mark.parametrize(
"config,expected_port",
[
({"host": "example.com"}, 1234),
({"host": "example.com", "port": 4321}, 4321),
({"host": "not_in_ssh_config.com"}, SSHFileSystem.DEFAULT_PORT),
({"host": "not_in_ssh_config.com", "port": 2222}, 2222),
],
)
@patch(
"builtins.open",
new_callable=mock_open,
read_data=mock_ssh_config,
)
def test_ssh_port(mock_file, config, expected_port):
fs = SSHFileSystem(**config)
assert fs.fs_args["port"] == expected_port
@pytest.mark.parametrize(
"config,expected_keyfile",
[
(
{"host": "example.com", "keyfile": "dvc_config.key"},
["dvc_config.key"],
),
(
{"host": "example.com"},
[os.path.expanduser("~/.ssh/not_default.key")],
),
(
{
"host": "not_in_ssh_config.com",
"keyfile": "dvc_config.key",
},
["dvc_config.key"],
),
({"host": "not_in_ssh_config.com"}, None),
],
)
@patch(
"builtins.open",
new_callable=mock_open,
read_data=mock_ssh_config,
)
def test_ssh_keyfile(mock_file, config, expected_keyfile):
fs = SSHFileSystem(**config)
assert fs.fs_args.get("client_keys") == expected_keyfile
@pytest.mark.parametrize(
"config,expected_gss_auth",
[
({"host": "example.com", "gss_auth": True}, True),
({"host": "example.com", "gss_auth": False}, False),
({"host": "not_in_ssh_config.com"}, False),
],
)
@patch(
"builtins.open",
new_callable=mock_open,
read_data=mock_ssh_config,
)
def test_ssh_gss_auth(mock_file, config, expected_gss_auth):
fs = SSHFileSystem(**config)
assert fs.fs_args["gss_auth"] == expected_gss_auth
tests/unit/fs/test_tree.py
+3 -5
View File
@@ -39,11 +39,9 @@ def test_get_cloud_fs_validate(tmp_dir, dvc):
default=False,
)
assert get_cloud_fs(dvc, name="base")[1] == {"host": "example.com"}
assert get_cloud_fs(dvc, name="first")[1] == {
"host": "example.com",
"type": ["symlink"],
}
assert get_cloud_fs(dvc, name="base")[1]["host"] == "example.com"
assert get_cloud_fs(dvc, name="first")[1]["host"] == "example.com"
assert get_cloud_fs(dvc, name="first")[1]["type"] == ["symlink"]
with pytest.raises(ConfigError):
get_cloud_fs(dvc, name="second")
tests/unit/remote/ssh/__init__.py
View File
tests/unit/remote/ssh/test_connection.py
-199
View File
@@ -1,199 +0,0 @@
import filecmp
import os
import platform
import posixpath
import tempfile
import pytest
from dvc.fs.ssh.connection import SSHConnection
from dvc.info import get_fs_type
from dvc.system import System
here = os.path.abspath(os.path.dirname(__file__))
SRC_PATH_WITH_SPECIAL_CHARACTERS = "Escape me [' , ']"
ESCAPED_SRC_PATH_WITH_SPECIAL_CHARACTERS = "'Escape me ['\"'\"' , '\"'\"']'"
DEST_PATH_WITH_SPECIAL_CHARACTERS = "Escape me too [' , ']"
ESCAPED_DEST_PATH_WITH_SPECIAL_CHARACTERS = (
"'Escape me too ['\"'\"' , '\"'\"']'"
)
def test_isdir(ssh_connection):
assert ssh_connection.isdir(here)
assert not ssh_connection.isdir(__file__)
def test_exists(ssh_connection):
assert not ssh_connection.exists("/path/to/non/existent/file")
assert ssh_connection.exists(__file__)
def test_isfile(ssh_connection):
assert ssh_connection.isfile(__file__)
assert not ssh_connection.isfile(here)
def test_makedirs(tmp_path, ssh_connection):
tmp = tmp_path.absolute().as_posix()
path = posixpath.join(tmp, "dir", "subdir")
ssh_connection.makedirs(path)
assert os.path.isdir(path)
def test_remove_dir(tmp_path, ssh_connection):
dpath = tmp_path / "dir"
dpath.mkdir()
(dpath / "file").write_text("file")
(dpath / "subdir").mkdir()
(dpath / "subdir" / "subfile").write_text("subfile")
ssh_connection.remove(dpath.absolute().as_posix())
assert not dpath.exists()
def test_walk(tmp_path, ssh_connection):
root_path = tmp_path
dir_path = root_path / "dir"
subdir_path = dir_path / "subdir"
dir_path.mkdir()
subdir_path.mkdir()
root_data_path = root_path / "root_data"
dir_data_path = dir_path / "dir_data"
subdir_data_path = subdir_path / "subdir_data"
with root_data_path.open("w+") as fobj:
fobj.write("")
with dir_data_path.open("w+") as fobj:
fobj.write("")
with subdir_data_path.open("w+") as fobj:
fobj.write("")
entries = [
dir_path,
subdir_path,
root_data_path,
dir_data_path,
subdir_data_path,
]
expected = {entry.absolute().as_posix() for entry in entries}
paths = set()
for root, dirs, files in ssh_connection.walk(
root_path.absolute().as_posix()
):
for entry in dirs + files:
paths.add(posixpath.join(root, entry))
assert paths == expected
@pytest.mark.skipif(
get_fs_type(tempfile.gettempdir())[0] not in ["xfs", "apfs", "btrfs"],
reason="Reflinks only work in specified file systems",
)
def test_reflink(tmp_dir, ssh_connection):
tmp_dir.gen("foo", "foo content")
ssh_connection.reflink("foo", "link")
assert filecmp.cmp("foo", "link")
assert not System.is_symlink("link")
assert not System.is_hardlink("link")
@pytest.mark.skipif(
platform.system() == "Windows",
reason="sftp symlink is not supported on Windows",
)
def test_symlink(tmp_dir, ssh_connection):
tmp_dir.gen("foo", "foo content")
ssh_connection.symlink("foo", "link")
assert System.is_symlink("link")
@pytest.mark.skipif(
platform.system() == "Windows",
reason="hardlink is temporarily not supported on Windows",
)
def test_hardlink(tmp_dir, ssh_connection):
tmp_dir.gen("foo", "foo content")
ssh_connection.hardlink("foo", "link")
assert System.is_hardlink("link")
@pytest.mark.skipif(
platform.system() == "Windows",
reason="copy is temporarily not supported on Windows",
)
def test_copy(tmp_dir, ssh_connection):
tmp_dir.gen("foo", "foo content")
ssh_connection.copy("foo", "link")
assert filecmp.cmp("foo", "link")
def test_move(tmp_dir, ssh_connection):
tmp_dir.gen("foo", "foo content")
ssh_connection.move("foo", "copy")
assert os.path.exists("copy")
assert not os.path.exists("foo")
@pytest.mark.parametrize(
"uname,md5command", [("Linux", "md5sum"), ("Darwin", "md5")]
)
def test_escapes_filepaths_for_md5_calculation(
ssh_connection, uname, md5command, mocker
):
fake_md5 = "x" * 32
uname_mock = mocker.PropertyMock(return_value=uname)
mocker.patch.object(SSHConnection, "uname", new_callable=uname_mock)
ssh_connection.execute = mocker.Mock(return_value=fake_md5)
ssh_connection.md5(SRC_PATH_WITH_SPECIAL_CHARACTERS)
ssh_connection.execute.assert_called_with(
f"{md5command} {ESCAPED_SRC_PATH_WITH_SPECIAL_CHARACTERS}"
)
def test_escapes_filepaths_for_copy(ssh_connection, mocker):
ssh_connection.execute = mocker.Mock()
ssh_connection.copy(
SRC_PATH_WITH_SPECIAL_CHARACTERS, DEST_PATH_WITH_SPECIAL_CHARACTERS
)
ssh_connection.execute.assert_called_with(
f"cp {ESCAPED_SRC_PATH_WITH_SPECIAL_CHARACTERS} "
+ f"{ESCAPED_DEST_PATH_WITH_SPECIAL_CHARACTERS}"
)
@pytest.mark.parametrize(
"uname,cp_command", [("Linux", "cp --reflink"), ("Darwin", "cp -c")]
)
def test_escapes_filepaths_for_reflink(
ssh_connection, uname, cp_command, mocker
):
uname_mock = mocker.PropertyMock(return_value=uname)
mocker.patch.object(SSHConnection, "uname", new_callable=uname_mock)
ssh_connection.execute = mocker.Mock()
ssh_connection.reflink(
SRC_PATH_WITH_SPECIAL_CHARACTERS, DEST_PATH_WITH_SPECIAL_CHARACTERS
)
ssh_connection.execute.assert_called_with(
f"{cp_command} "
+ f"{ESCAPED_SRC_PATH_WITH_SPECIAL_CHARACTERS} "
+ f"{ESCAPED_DEST_PATH_WITH_SPECIAL_CHARACTERS}"
)
def test_escapes_filepaths_for_hardlink(ssh_connection, mocker):
ssh_connection.execute = mocker.Mock()
ssh_connection.hardlink(
SRC_PATH_WITH_SPECIAL_CHARACTERS, DEST_PATH_WITH_SPECIAL_CHARACTERS
)
ssh_connection.execute.assert_called_with(
f"ln {ESCAPED_SRC_PATH_WITH_SPECIAL_CHARACTERS} "
+ f"{ESCAPED_DEST_PATH_WITH_SPECIAL_CHARACTERS}"
)
tests/unit/remote/ssh/test_pool.py
-20
View File
@@ -1,20 +0,0 @@
import pytest
from dvc.fs.pool import get_connection
from dvc.fs.ssh.connection import SSHConnection
from tests.remotes.ssh import TEST_SSH_KEY_PATH, TEST_SSH_USER
def test_doesnt_swallow_errors(ssh_server):
class MyError(Exception):
pass
with pytest.raises(MyError):
with get_connection(
SSHConnection,
host=ssh_server.host,
port=ssh_server.port,
username=TEST_SSH_USER,
key_filename=TEST_SSH_KEY_PATH,
):
raise MyError
tests/unit/remote/ssh/test_ssh.py
-225
View File
@@ -1,225 +0,0 @@
import getpass
import os
import sys
from unittest.mock import mock_open, patch
import pytest
from dvc.fs.ssh import SSHFileSystem
from dvc.system import System
def test_get_kwargs_from_urls():
user = "test"
host = "123.45.67.89"
port = 1234
path = "/path/to/dir"
# URL ssh://[user@]host.xz[:port]/path
url = f"ssh://{user}@{host}:{port}{path}"
assert SSHFileSystem._get_kwargs_from_urls(url) == {
"user": user,
"host": host,
"port": port,
}
# SCP-like URL ssh://[user@]host.xz:/absolute/path
url = f"ssh://{user}@{host}:{path}"
assert SSHFileSystem._get_kwargs_from_urls(url) == {
"user": user,
"host": host,
}
def test_init():
fs = SSHFileSystem(user="test", host="12.34.56.78", port="1234")
assert fs.user == "test"
assert fs.host == "12.34.56.78"
assert fs.port == "1234"
mock_ssh_config = """
Host example.com
User ubuntu
HostName 1.2.3.4
Port 1234
IdentityFile ~/.ssh/not_default.key
"""
if sys.version_info.major == 3:
builtin_module_name = "builtins"
else:
builtin_module_name = "__builtin__"
@pytest.mark.parametrize(
"config,expected_host",
[
({"host": "example.com"}, "1.2.3.4"),
({"host": "not_in_ssh_config.com"}, "not_in_ssh_config.com"),
],
)
@patch("os.path.exists", return_value=True)
@patch(
f"{builtin_module_name}.open",
new_callable=mock_open,
read_data=mock_ssh_config,
)
def test_ssh_host_override_from_config(
mock_file, mock_exists, config, expected_host
):
fs = SSHFileSystem(**config)
mock_exists.assert_called_with(SSHFileSystem.ssh_config_filename())
mock_file.assert_called_with(SSHFileSystem.ssh_config_filename())
assert fs.host == expected_host
@pytest.mark.parametrize(
"config,expected_user",
[
({"host": "example.com", "user": "test1"}, "test1"),
({"host": "example.com"}, "ubuntu"),
({"host": "not_in_ssh_config.com", "user": "test1"}, "test1"),
({"host": "not_in_ssh_config.com"}, getpass.getuser()),
],
)
@patch("os.path.exists", return_value=True)
@patch(
f"{builtin_module_name}.open",
new_callable=mock_open,
read_data=mock_ssh_config,
)
def test_ssh_user(mock_file, mock_exists, config, expected_user):
fs = SSHFileSystem(**config)
mock_exists.assert_called_with(SSHFileSystem.ssh_config_filename())
mock_file.assert_called_with(SSHFileSystem.ssh_config_filename())
assert fs.user == expected_user
@pytest.mark.parametrize(
"config,expected_port",
[
({"host": "example.com"}, 1234),
({"host": "example.com", "port": 4321}, 4321),
({"host": "not_in_ssh_config.com"}, SSHFileSystem.DEFAULT_PORT),
({"host": "not_in_ssh_config.com", "port": 2222}, 2222),
],
)
@patch("os.path.exists", return_value=True)
@patch(
f"{builtin_module_name}.open",
new_callable=mock_open,
read_data=mock_ssh_config,
)
def test_ssh_port(mock_file, mock_exists, config, expected_port):
fs = SSHFileSystem(**config)
mock_exists.assert_called_with(SSHFileSystem.ssh_config_filename())
mock_file.assert_called_with(SSHFileSystem.ssh_config_filename())
assert fs.port == expected_port
@pytest.mark.parametrize(
"config,expected_keyfile",
[
(
{"host": "example.com", "keyfile": "dvc_config.key"},
"dvc_config.key",
),
(
{"host": "example.com"},
os.path.expanduser("~/.ssh/not_default.key"),
),
(
{
"host": "not_in_ssh_config.com",
"keyfile": "dvc_config.key",
},
"dvc_config.key",
),
({"host": "not_in_ssh_config.com"}, None),
],
)
@patch("os.path.exists", return_value=True)
@patch(
f"{builtin_module_name}.open",
new_callable=mock_open,
read_data=mock_ssh_config,
)
def test_ssh_keyfile(mock_file, mock_exists, config, expected_keyfile):
fs = SSHFileSystem(**config)
mock_exists.assert_called_with(SSHFileSystem.ssh_config_filename())
mock_file.assert_called_with(SSHFileSystem.ssh_config_filename())
assert fs.keyfile == expected_keyfile
@pytest.mark.parametrize(
"config,expected_gss_auth",
[
({"host": "example.com", "gss_auth": True}, True),
({"host": "example.com", "gss_auth": False}, False),
({"host": "not_in_ssh_config.com"}, False),
],
)
@patch("os.path.exists", return_value=True)
@patch(
f"{builtin_module_name}.open",
new_callable=mock_open,
read_data=mock_ssh_config,
)
def test_ssh_gss_auth(mock_file, mock_exists, config, expected_gss_auth):
fs = SSHFileSystem(**config)
mock_exists.assert_called_with(SSHFileSystem.ssh_config_filename())
mock_file.assert_called_with(SSHFileSystem.ssh_config_filename())
assert fs.gss_auth == expected_gss_auth
@pytest.mark.parametrize(
"config,expected_allow_agent",
[
({"host": "example.com"}, True),
({"host": "not_in_ssh_config.com"}, True),
({"host": "example.com", "allow_agent": True}, True),
({"host": "example.com", "allow_agent": False}, False),
],
)
@patch("os.path.exists", return_value=True)
@patch(
f"{builtin_module_name}.open",
new_callable=mock_open,
read_data=mock_ssh_config,
)
def test_ssh_allow_agent(mock_file, mock_exists, config, expected_allow_agent):
fs = SSHFileSystem(**config)
mock_exists.assert_called_with(SSHFileSystem.ssh_config_filename())
mock_file.assert_called_with(SSHFileSystem.ssh_config_filename())
assert fs.allow_agent == expected_allow_agent
def test_hardlink_optimization(tmp_dir, dvc, ssh):
from dvc.fs import get_cloud_fs
cls, config, path_info = get_cloud_fs(dvc, **ssh.config)
fs = cls(**config)
assert isinstance(fs, SSHFileSystem)
from_info = path_info / "empty"
to_info = path_info / "link"
with fs.open(from_info, "wb"):
pass
if os.name == "nt":
link_path = "c:" + to_info.path.replace("/", "\\")
else:
link_path = to_info.path
fs.hardlink(from_info, to_info)
assert not System.is_hardlink(link_path)