From a87dea670dc7a6a27cc2cb83f4c90b0129e9998b Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Timo=20G=C3=BCnther?=
 <18397625+TimoGuenther@users.noreply.github.com>
Date: Wed, 24 May 2023 23:39:57 +0200
Subject: [PATCH] [sftp] add closing contextmanager (#1253)

This ensures resources can be cleaned up when used standalone. Guidance is also added to the docs.
---
 docs/backends/sftp.rst           | 15 +++++++++++++++
 setup.cfg                        |  2 +-
 storages/backends/sftpstorage.py |  9 ++++++++-
 tests/test_sftp.py               | 12 ++++++++++++
 4 files changed, 36 insertions(+), 2 deletions(-)

diff --git a/docs/backends/sftp.rst b/docs/backends/sftp.rst
index a83130a65..d09303564 100644
--- a/docs/backends/sftp.rst
+++ b/docs/backends/sftp.rst
@@ -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("")
diff --git a/setup.cfg b/setup.cfg
index a37261788..e916cd0c6 100644
--- a/setup.cfg
+++ b/setup.cfg
@@ -47,7 +47,7 @@ google =
 libcloud =
     apache-libcloud
 sftp =
-    paramiko >= 1.10.0
+    paramiko >= 1.15.0
 
 [flake8]
 exclude =
diff --git a/storages/backends/sftpstorage.py b/storages/backends/sftpstorage.py
index 34a5755c2..3f1395230 100644
--- a/storages/backends/sftpstorage.py
+++ b/storages/backends/sftpstorage.py
@@ -15,6 +15,7 @@
 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
@@ -22,7 +23,7 @@
 
 
 @deconstructible
-class SFTPStorage(BaseStorage):
+class SFTPStorage(ClosingContextManager, BaseStorage):
     def __init__(self, **settings):
         super().__init__(**settings)
         self._host = self.host
@@ -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):
@@ -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"""
diff --git a/tests/test_sftp.py b/tests/test_sftp.py
index 703a6b272..28525c4d1 100644
--- a/tests/test_sftp.py
+++ b/tests/test_sftp.py
@@ -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)