Repository: RaRe-Technologies/smart_open
Pull request #719 (2022-08-25 19:57:53)

Reconnect inactive sftp clients automatically
Fix: #714

Also fix dependency installation for ssh backend


Kache: fyi @mpenkov
Kache: @mpenkov sorry about that, ready now

Patch:
diff --git a/setup.py b/setup.py
index 6c08ade..5bd4621 100644
--- a/setup.py
+++ b/setup.py
@@ -40,13 +40,13 @@ aws_deps = ['boto3']
 gcs_deps = ['google-cloud-storage>=1.31.0']
 azure_deps = ['azure-storage-blob', 'azure-common', 'azure-core']
 http_deps = ['requests']
+ssh_deps = ['paramiko']
 
-all_deps = aws_deps + gcs_deps + azure_deps + http_deps
+all_deps = aws_deps + gcs_deps + azure_deps + http_deps + ssh_deps
 tests_require = all_deps + [
     'moto[server]',
     'responses',
     'boto3',
-    'paramiko',
     'pytest',
     'pytest-rerunfailures'
 ]
@@ -79,6 +79,7 @@ setup(
         'all': all_deps,
         'http': http_deps,
         'webhdfs': http_deps,
+        'ssh': ssh_deps,
     },
     python_requires=">=3.6,<4.0",
 
diff --git a/smart_open/compression.py b/smart_open/compression.py
index fe02246..08469e7 100644
--- a/smart_open/compression.py
+++ b/smart_open/compression.py
@@ -62,6 +62,7 @@ def register_compressor(ext, callback):
     """
     if not (ext and ext[0] == '.'):
         raise ValueError('ext must be a string starting with ., not %r' % ext)
+    ext = ext.lower()
     if ext in _COMPRESSOR_REGISTRY:
         logger.warning('overriding existing compression handler for %r', ext)
     _COMPRESSOR_REGISTRY[ext] = callback
@@ -103,24 +104,23 @@ def _handle_gzip(file_obj, mode):
     return result
 
 
-def compression_wrapper(file_obj, mode, compression):
+def compression_wrapper(file_obj, mode, compression=INFER_FROM_EXTENSION, filename=None):
     """
-    This function will wrap the file_obj with an appropriate
-    [de]compression mechanism based on the specified extension.
+    Wrap `file_obj` with an appropriate [de]compression mechanism based on its file extension.
 
-    file_obj must either be a filehandle object, or a class which behaves
-    like one. It must have a .name attribute.
+    If the filename extension isn't recognized, simply return the original `file_obj` unchanged.
 
-    If the filename extension isn't recognized, will simply return the original
-    file_obj.
+    `file_obj` must either be a filehandle object, or a class which behaves like one.
+
+    If `filename` is specified, it will be used to extract the extension.
+    If not, the `file_obj.name` attribute is used as the filename.
 
     """
     if compression == NO_COMPRESSION:
         return file_obj
     elif compression == INFER_FROM_EXTENSION:
         try:
-            filename = file_obj.name
-            filename.upper()  # make sure this thing is a string
+            filename = (filename or file_obj.name).lower()
         except (AttributeError, TypeError):
             logger.warning(
                 'unable to transparently decompress %r because it '
diff --git a/smart_open/ssh.py b/smart_open/ssh.py
index fa762eb..4b31893 100644
--- a/smart_open/ssh.py
+++ b/smart_open/ssh.py
@@ -25,7 +25,11 @@ Similarly, from a command line::
 import getpass
 import logging
 import urllib.parse
-import warnings
+
+try:
+    import paramiko
+except ImportError:
+    MISSING_DEPS = True
 
 import smart_open.utils
 
@@ -74,29 +78,15 @@ def open_uri(uri, mode, transport_params):
     return open(uri_path, mode, transport_params=transport_params, **parsed_uri)
 
 
-def _connect(hostname, username, port, password, transport_params):
-    try:
-        import paramiko
-    except ImportError:
-        warnings.warn(
-            'paramiko missing, opening SSH/SCP/SFTP paths will be disabled. '
-            '`pip install paramiko` to suppress'
-        )
-        raise
-
-    key = (hostname, username)
-    ssh = _SSH.get(key)
-    if ssh is None:
-        ssh = _SSH[key] = paramiko.client.SSHClient()
-        ssh.load_system_host_keys()
-        ssh.set_missing_host_key_policy(paramiko.AutoAddPolicy())
-        kwargs = transport_params.get('connect_kwargs', {}).copy()
-        # if 'key_filename' is present in transport_params, then I do not
-        #   overwrite the credentials.
-        if 'key_filename' not in kwargs:
-            kwargs.setdefault('password', password)
-        kwargs.setdefault('username', username)
-        ssh.connect(hostname, port, **kwargs)
+def _connect_ssh(hostname, username, port, password, transport_params):
+    ssh = paramiko.SSHClient()
+    ssh.load_system_host_keys()
+    ssh.set_missing_host_key_policy(paramiko.AutoAddPolicy())
+    kwargs = transport_params.get('connect_kwargs', {}).copy()
+    if 'key_filename' not in kwargs:
+        kwargs.setdefault('password', password)
+    kwargs.setdefault('username', username)
+    ssh.connect(hostname, port, **kwargs)
     return ssh
 
 
@@ -142,8 +132,30 @@ def open(path, mode='r', host=None, user=None, password=None, port=DEFAULT_PORT,
     if not transport_params:
         transport_params = {}
 
-    conn = _connect(host, user, port, password, transport_params)
-    sftp_client = conn.get_transport().open_sftp_client()
+    key = (host, user)
+
+    attempts = 2
+    for attempt in range(attempts):
+        try:
+            ssh = _SSH[key]
+        except KeyError:
+            ssh = _SSH[key] = _connect_ssh(host, user, port, password, transport_params)
+
+        try:
+            transport = ssh.get_transport()
+            sftp_client = transport.open_sftp_client()
+            break
+        except paramiko.SSHException as ex:
+            connection_timed_out = ex.args and ex.args[0] == 'SSH session not active'
+            if attempt == attempts - 1 or not connection_timed_out:
+                raise
+
+            #
+            # Try again.  Delete the connection from the cache to force a
+            # reconnect in the next attempt.
+            #
+            del _SSH[key]
+
     fobj = sftp_client.open(path, mode)
     fobj.name = path
     return fobj


Test patch:
diff --git a/integration-tests/test_ssh.py b/integration-tests/test_ssh.py
new file mode 100644
index 0000000..252a6fb
--- /dev/null
+++ b/integration-tests/test_ssh.py
@@ -0,0 +1,42 @@
+# -*- coding: utf-8 -*-
+#
+# Copyright (C) 2022 Radim Rehurek <me@radimrehurek.com>
+#
+# This code is distributed under the terms and conditions
+# from the MIT License (MIT).
+#
+
+import os
+import tempfile
+import pytest
+
+import smart_open
+import smart_open.ssh
+
+
+def explode(*args, **kwargs):
+    raise RuntimeError("this function should never have been called")
+
+
+@pytest.mark.skipif("SMART_OPEN_SSH" not in os.environ, reason="this test only works on the dev machine")
+def test():
+    with smart_open.open("ssh://misha@localhost/Users/misha/git/smart_open/README.rst") as fin:
+        readme = fin.read()
+
+    assert 'smart_open — utils for streaming large files in Python' in readme
+
+    #
+    # Ensure the cache is being used
+    #
+    assert ('localhost', 'misha') in smart_open.ssh._SSH
+
+    try:
+        connect_ssh = smart_open.ssh._connect_ssh
+        smart_open.ssh._connect_ssh = explode
+
+        with smart_open.open("ssh://misha@localhost/Users/misha/git/smart_open/howto.md") as fin:
+            howto = fin.read()
+
+        assert 'How-to Guides' in howto
+    finally:
+        smart_open.ssh._connect_ssh = connect_ssh
diff --git a/smart_open/tests/test_compression.py b/smart_open/tests/test_compression.py
new file mode 100644
index 0000000..4b93f63
--- /dev/null
+++ b/smart_open/tests/test_compression.py
@@ -0,0 +1,39 @@
+# -*- coding: utf-8 -*-
+#
+# Copyright (C) 2020 Radim Rehurek <me@radimrehurek.com>
+#
+# This code is distributed under the terms and conditions
+# from the MIT License (MIT).
+#
+import io
+import gzip
+import pytest
+
+import smart_open.compression
+
+
+plain = 'доброе утро планета!'.encode()
+
+
+def label(thing, name):
+    setattr(thing, 'name', name)
+    return thing
+
+
+@pytest.mark.parametrize(
+    'fileobj,compression,filename',
+    [
+        (io.BytesIO(plain), 'disable', None),
+        (io.BytesIO(plain), 'disable', ''),
+        (io.BytesIO(plain), 'infer_from_extension', 'file.txt'),
+        (io.BytesIO(plain), 'infer_from_extension', 'file.TXT'),
+        (io.BytesIO(plain), '.unknown', ''),
+        (io.BytesIO(gzip.compress(plain)), 'infer_from_extension', 'file.gz'),
+        (io.BytesIO(gzip.compress(plain)), 'infer_from_extension', 'file.GZ'),
+        (label(io.BytesIO(gzip.compress(plain)), 'file.gz'), 'infer_from_extension', ''),
+        (io.BytesIO(gzip.compress(plain)), '.gz', 'file.gz'),
+    ]
+)
+def test_compression_wrapper_read(fileobj, compression, filename):
+    wrapped = smart_open.compression.compression_wrapper(fileobj, 'rb', compression, filename)
+    assert wrapped.read() == plain
diff --git a/smart_open/tests/test_ssh.py b/smart_open/tests/test_ssh.py
index 18a8f16..a4b91eb 100644
--- a/smart_open/tests/test_ssh.py
+++ b/smart_open/tests/test_ssh.py
@@ -4,6 +4,8 @@ import logging
 import unittest
 from unittest import mock
 
+from paramiko import SSHException
+
 import smart_open.ssh
 
 
@@ -12,8 +14,8 @@ def mock_ssh(func):
         smart_open.ssh._SSH.clear()
         return func(*args, **kwargs)
 
-    return mock.patch("paramiko.client.SSHClient.get_transport")(
-        mock.patch("paramiko.client.SSHClient.connect")(wrapper)
+    return mock.patch("paramiko.SSHClient.get_transport")(
+        mock.patch("paramiko.SSHClient.connect")(wrapper)
     )
 
 
@@ -49,6 +51,23 @@ class SSHOpen(unittest.TestCase):
         )
         mock_connect.assert_called_with("some-host", 22, username="user", key_filename="key")
 
+    @mock_ssh
+    def test_reconnect_after_session_timeout(self, mock_connect, get_transp_mock):
+        mock_sftp = get_transp_mock().open_sftp_client()
+        get_transp_mock().open_sftp_client.reset_mock()
+
+        def mocked_open_sftp():
+            if len(mock_connect.call_args_list) < 2:  # simulate timeout until second connect()
+                yield SSHException('SSH session not active')
+            while True:
+                yield mock_sftp
+
+        get_transp_mock().open_sftp_client.side_effect = mocked_open_sftp()
+
+        smart_open.open("ssh://user:pass@some-host/")
+        mock_connect.assert_called_with("some-host", 22, username="user", password="pass")
+        mock_sftp.open.assert_called_once()
+
 
 if __name__ == "__main__":
     logging.basicConfig(format="%(asctime)s : %(levelname)s : %(message)s", level=logging.DEBUG)