From 9b4b6e5b31de0bbada079c57eca71eeab105d969 Mon Sep 17 00:00:00 2001 From: Batuhan Taskaya Date: Wed, 28 Jul 2021 12:49:53 +0300 Subject: [PATCH] ssh: migrate to sshfs (#6295) --- dvc/fs/fsspec_wrapper.py | 13 +- dvc/fs/ssh.py | 125 ++++++++ dvc/fs/ssh/__init__.py | 263 ----------------- dvc/fs/ssh/connection.py | 356 ----------------------- dvc/objects/db/__init__.py | 4 - dvc/objects/db/ssh.py | 90 ------ dvc/utils/fs.py | 19 ++ setup.py | 4 +- tests/remotes/ssh.py | 84 ++++-- tests/unit/fs/test_ssh.py | 153 ++++++++++ tests/unit/fs/test_tree.py | 8 +- tests/unit/remote/ssh/__init__.py | 0 tests/unit/remote/ssh/test_connection.py | 199 ------------- tests/unit/remote/ssh/test_pool.py | 20 -- tests/unit/remote/ssh/test_ssh.py | 225 -------------- 15 files changed, 366 insertions(+), 1197 deletions(-) create mode 100644 dvc/fs/ssh.py delete mode 100644 dvc/fs/ssh/__init__.py delete mode 100644 dvc/fs/ssh/connection.py delete mode 100644 dvc/objects/db/ssh.py create mode 100644 tests/unit/fs/test_ssh.py delete mode 100644 tests/unit/remote/ssh/__init__.py delete mode 100644 tests/unit/remote/ssh/test_connection.py delete mode 100644 tests/unit/remote/ssh/test_pool.py delete mode 100644 tests/unit/remote/ssh/test_ssh.py diff --git a/dvc/fs/fsspec_wrapper.py b/dvc/fs/fsspec_wrapper.py index 74021714a..1e0056399 100644 --- a/dvc/fs/fsspec_wrapper.py +++ b/dvc/fs/fsspec_wrapper.py @@ -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), ) diff --git a/dvc/fs/ssh.py b/dvc/fs/ssh.py new file mode 100644 index 000000000..799e11d77 --- /dev/null +++ b/dvc/fs/ssh.py @@ -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"] = [ + "aes128-gcm@openssh.com", + "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) diff --git a/dvc/fs/ssh/__init__.py b/dvc/fs/ssh/__init__.py deleted file mode 100644 index 2fca225c2..000000000 --- a/dvc/fs/ssh/__init__.py +++ /dev/null @@ -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, - ) diff --git a/dvc/fs/ssh/connection.py b/dvc/fs/ssh/connection.py deleted file mode 100644 index 465172b27..000000000 --- a/dvc/fs/ssh/connection.py +++ /dev/null @@ -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}") diff --git a/dvc/objects/db/__init__.py b/dvc/objects/db/__init__.py index ece8e45b4..e0d450b7d 100644 --- a/dvc/objects/db/__init__.py +++ b/dvc/objects/db/__init__.py @@ -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) diff --git a/dvc/objects/db/ssh.py b/dvc/objects/db/ssh.py deleted file mode 100644 index 1461fef06..000000000 --- a/dvc/objects/db/ssh.py +++ /dev/null @@ -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) diff --git a/dvc/utils/fs.py b/dvc/utils/fs.py index fd880ed07..33f49a358 100644 --- a/dvc/utils/fs.py +++ b/dvc/utils/fs.py @@ -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) diff --git a/setup.py b/setup.py index 075e221da..ea73d5851 100644 --- a/setup.py +++ b/setup.py @@ -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 = ( diff --git a/tests/remotes/ssh.py b/tests/remotes/ssh.py index 71a051a1d..1eac36a71 100644 --- a/tests/remotes/ssh.py +++ b/tests/remotes/ssh.py @@ -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], ) diff --git a/tests/unit/fs/test_ssh.py b/tests/unit/fs/test_ssh.py new file mode 100644 index 000000000..474a589de --- /dev/null +++ b/tests/unit/fs/test_ssh.py @@ -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 diff --git a/tests/unit/fs/test_tree.py b/tests/unit/fs/test_tree.py index c5f8769cb..1254c5bed 100644 --- a/tests/unit/fs/test_tree.py +++ b/tests/unit/fs/test_tree.py @@ -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") diff --git a/tests/unit/remote/ssh/__init__.py b/tests/unit/remote/ssh/__init__.py deleted file mode 100644 index e69de29bb..000000000 diff --git a/tests/unit/remote/ssh/test_connection.py b/tests/unit/remote/ssh/test_connection.py deleted file mode 100644 index f77d48abd..000000000 --- a/tests/unit/remote/ssh/test_connection.py +++ /dev/null @@ -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}" - ) diff --git a/tests/unit/remote/ssh/test_pool.py b/tests/unit/remote/ssh/test_pool.py deleted file mode 100644 index c2e0087d0..000000000 --- a/tests/unit/remote/ssh/test_pool.py +++ /dev/null @@ -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 diff --git a/tests/unit/remote/ssh/test_ssh.py b/tests/unit/remote/ssh/test_ssh.py deleted file mode 100644 index 72e9d5692..000000000 --- a/tests/unit/remote/ssh/test_ssh.py +++ /dev/null @@ -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)