Skip to content

Commit

Permalink
[sftp] add closing contextmanager (#1253)
Browse files Browse the repository at this point in the history
This ensures resources can be cleaned up when used standalone. Guidance is also added to the docs.
  • Loading branch information
TimoGuenther authored May 24, 2023
1 parent 1e58135 commit a87dea6
Show file tree
Hide file tree
Showing 4 changed files with 36 additions and 2 deletions.
15 changes: 15 additions & 0 deletions docs/backends/sftp.rst
Original file line number Diff line number Diff line change
Expand Up @@ -71,3 +71,18 @@ static files) to use the sftp backend::
.. _`paramiko SSHClient.connect() documentation`: http://docs.paramiko.org/en/latest/api/client.html#paramiko.client.SSHClient.connect

.. _`Python os.chmod documentation`: http://docs.python.org/library/os.html#os.chmod


Standalone Use
--------------

If you intend to construct a storage instance not through Django but directly,
use the storage instance as a context manager to make sure the underlying SSH
connection is closed after use and no longer consumes resources.

.. code-block:: python
from storages.backends.sftpstorage import SFTPStorage
with SFTPStorage(...) as sftp:
sftp.listdir("")
2 changes: 1 addition & 1 deletion setup.cfg
Original file line number Diff line number Diff line change
Expand Up @@ -47,7 +47,7 @@ google =
libcloud =
apache-libcloud
sftp =
paramiko >= 1.10.0
paramiko >= 1.15.0

[flake8]
exclude =
Expand Down
9 changes: 8 additions & 1 deletion storages/backends/sftpstorage.py
Original file line number Diff line number Diff line change
Expand Up @@ -15,14 +15,15 @@
import paramiko
from django.core.files.base import File
from django.utils.deconstruct import deconstructible
from paramiko.util import ClosingContextManager

from storages.base import BaseStorage
from storages.utils import is_seekable
from storages.utils import setting


@deconstructible
class SFTPStorage(BaseStorage):
class SFTPStorage(ClosingContextManager, BaseStorage):
def __init__(self, **settings):
super().__init__(**settings)
self._host = self.host
Expand All @@ -35,6 +36,7 @@ def __init__(self, **settings):
self._known_host_file = self.known_host_file
self._root_path = self.root_path
self._base_url = self.base_url
self._ssh = None
self._sftp = None

def get_default_settings(self):
Expand Down Expand Up @@ -81,6 +83,11 @@ def _connect(self):
if self._ssh.get_transport():
self._sftp = self._ssh.open_sftp()

def close(self):
if self._ssh is None:
return
self._ssh.close()

@property
def sftp(self):
"""Lazy SFTP connection"""
Expand Down
12 changes: 12 additions & 0 deletions tests/test_sftp.py
Original file line number Diff line number Diff line change
Expand Up @@ -40,6 +40,18 @@ def test_connect(self, mock_ssh):
self.storage._connect()
self.assertEqual('foo', mock_ssh.return_value.connect.call_args[0][0])

@patch('paramiko.SSHClient')
def test_close_unopened(self, mock_ssh):
with self.storage:
pass
mock_ssh.return_value.close.assert_not_called()

@patch('paramiko.SSHClient')
def test_close_opened(self, mock_ssh):
with self.storage as storage:
storage._connect()
mock_ssh.return_value.close.assert_called_once_with()

def test_open(self):
file_ = self.storage._open('foo')
self.assertIsInstance(file_, sftpstorage.SFTPStorageFile)
Expand Down

0 comments on commit a87dea6

Please sign in to comment.