diff --git a/sshfs/pools/base.py b/sshfs/pools/base.py index 9363800..1510a29 100644 --- a/sshfs/pools/base.py +++ b/sshfs/pools/base.py @@ -1,5 +1,6 @@ import asyncio from contextlib import AsyncExitStack, suppress +from typing import Optional from asyncssh.misc import ChannelOpenError @@ -15,9 +16,10 @@ def __init__( self, client, *, - max_channels=None, - timeout=MAX_TIMEOUT, - unsafe_terminate=True, + max_channels: Optional[int] = None, + timeout: Optional[float] = MAX_TIMEOUT, + unsafe_terminate: bool = True, + sftp_client_kwargs: Optional[dict] = None, **kwargs, ): self.client = client @@ -38,6 +40,8 @@ def __init__( self.unsafe_terminate = unsafe_terminate self._stack = AsyncExitStack() + self.sftp_client_kwargs = sftp_client_kwargs or {} + async def _maybe_new_channel(self): # If there is no hard limit or the limit is not hit yet # try to create a new channel @@ -47,7 +51,7 @@ async def _maybe_new_channel(self): ): try: return await self._stack.enter_async_context( - self.client.start_sftp_client() + self.client.start_sftp_client(**self.sftp_client_kwargs) ) except ChannelOpenError: # If we can't create any more channels, then change diff --git a/sshfs/spec.py b/sshfs/spec.py index 77a1e11..b47dcf0 100644 --- a/sshfs/spec.py +++ b/sshfs/spec.py @@ -5,6 +5,7 @@ import weakref from contextlib import AsyncExitStack, suppress from datetime import datetime, timezone +from typing import Optional import asyncssh from asyncssh.sftp import SFTPOpUnsupported @@ -40,6 +41,7 @@ def __init__( host, *, pool_type=SFTPSoftChannelPool, + sftp_client_kwargs: Optional[dict] = None, **kwargs, ): """ @@ -51,15 +53,21 @@ def __init__( SSH host to connect. **kwargs: Any Any option that will be passed to either the top level - `AsyncFileSystem` or the `asyncssh.connect`. + `AsyncFileSystem` or the `asyncssh.connect`. `timeout` + limits the initial connection setup. pool_type: sshfs.pools.base.BaseSFTPChannelPool Pool manager to use (when doing concurrent operations together, pool managers offer the flexibility of prioritizing channels and deciding which to use). + sftp_client_kwargs: Optional[dict] + Options passed to `SSHClientConnection.start_sftp_client` + (e.g. env, send_env, path_encoding, path_errors, + sftp_version). """ super().__init__(self, **kwargs) + _timeout = kwargs.pop("timeout", None) max_sessions = kwargs.pop("max_sessions", _DEFAULT_MAX_SESSIONS) if max_sessions <= _SHELL_CHANNELS: raise ValueError( @@ -67,6 +75,7 @@ def __init__( ) _client_args = kwargs.copy() _client_args.setdefault("known_hosts", None) + sftp_client_kwargs = sftp_client_kwargs or {} self._stack = AsyncExitStack() self.active_executors = 0 @@ -74,7 +83,11 @@ def __init__( host, pool_type, max_sftp_channels=max_sessions - _SHELL_CHANNELS, - **_client_args, + # timeout is consumed by the sync() machinery underneath + # sync_wrapper, it never reaches _connect + timeout=_timeout, + connect_args=_client_args, + sftp_client_kwargs=sftp_client_kwargs, ) weakref.finalize( self, self._finalize, self.loop, self._pool, self._stack @@ -95,13 +108,22 @@ def _get_kwargs_from_urls(urlpath): @wrap_exceptions async def _connect( - self, host, pool_type, max_sftp_channels, **client_args + self, + host, + pool_type, + max_sftp_channels, + connect_args, + sftp_client_kwargs, ): self._client_lock = asyncio.Semaphore(_SHELL_CHANNELS) - _raw_client = asyncssh.connect(host, **client_args) + _raw_client = asyncssh.connect(host, **connect_args) client = await self._stack.enter_async_context(_raw_client) - pool = pool_type(client, max_channels=max_sftp_channels) + pool = pool_type( + client, + max_channels=max_sftp_channels, + sftp_client_kwargs=sftp_client_kwargs, + ) return client, pool connect = sync_wrapper(_connect) diff --git a/tests/test_sftp_pools.py b/tests/test_sftp_pools.py index 39375b9..8e7af33 100644 --- a/tests/test_sftp_pools.py +++ b/tests/test_sftp_pools.py @@ -38,11 +38,13 @@ class FakeSSHClient: def __init__(self, max_channels=None): self.counter = 0 self.max_channels = max_channels + self.received_sftp_client_kwargs = None @asynccontextmanager - async def start_sftp_client(self): + async def start_sftp_client(self, **kwargs): from asyncssh.misc import ChannelOpenError + self.received_sftp_client_kwargs = kwargs if self.max_channels is not None and self.counter >= self.max_channels: raise ChannelOpenError(None, None) @@ -55,6 +57,19 @@ def fake_client(): yield FakeSSHClient() +@all_queues +@pytest.mark.asyncio +async def test_pool_forwards_sftp_client_kwargs(queue_type, fake_client): + pool = queue_type( + fake_client, + sftp_client_kwargs={"sftp_version": 3}, + poll=False, + ) + async with pool.get() as channel: + assert channel.no == 1 + assert fake_client.received_sftp_client_kwargs == {"sftp_version": 3} + + @pytest.mark.asyncio async def test_pool_soft_queue_caching(fake_client): pool = SFTPSoftChannelPool(fake_client, poll=False) diff --git a/tests/test_sshfs.py b/tests/test_sshfs.py index ebe351b..9bd3242 100644 --- a/tests/test_sshfs.py +++ b/tests/test_sshfs.py @@ -100,6 +100,43 @@ def test_fsspec_url_parsing(ssh_server, remote_dir, user="user"): } +def test_sftp_client_kwargs(ssh_server, base_remote_dir, user="user"): + fs = SSHFileSystem( + host=ssh_server.host, + port=ssh_server.port, + username=user, + client_keys=[USERS[user]], + sftp_client_kwargs={"sftp_version": 3}, + ) + assert fs._pool.sftp_client_kwargs == {"sftp_version": 3} + + file = posixpath.join(base_remote_dir, "sftp_client_kwargs_probe") + fs.touch(file) + assert fs.exists(file) + + +def test_sftp_client_kwargs_path_encoding( + ssh_server, base_remote_dir, user="user" +): + # path_encoding=None makes asyncssh deliver remote paths as raw + # bytes instead of str, for servers with non-UTF-8 file names (#39). + fs = SSHFileSystem( + host=ssh_server.host, + port=ssh_server.port, + username=user, + client_keys=[USERS[user]], + sftp_client_kwargs={"path_encoding": None}, + ) + + directory = Path(base_remote_dir) / "path_encoding_probe" + directory.mkdir() + (directory / "data.txt").write_bytes(b"data") + + encoded = str(directory).encode() + assert fs.ls(encoded) == [encoded + b"/data.txt"] + assert fs.cat_file(encoded + b"/data.txt") == b"data" + + def test_info(fs, remote_dir): fs.touch(remote_dir + "/a.txt") details = fs.info(remote_dir + "/a.txt")