Close store connections at the end of the program.

2 files changed, 18 insertions, 0 deletions

diff --git a/obnam/backend.py b/obnam/backend.py
index 09d029e0..e5d806eb 100644
--- a/obnam/backend.py
+++ b/obnam/backend.py
@@ -226,6 +226,11 @@ class SftpBackend(Backend):
             logging.debug("Opening sftp client")
             self.sftp_client = self.sftp_transport.open_sftp_client()
     
+    def close(self):
+        """Close the connection, if any."""
+        if self.sftp_transport:
+            self.sftp_transport.close()
+    
     def sftp_makedirs(self, dirname, mode=0777):
         """Create dirname, if it doesn't exist, and all its parents, too"""
         stack = []
@@ -315,6 +320,9 @@ class SftpBackend(Backend):
 
 class FileBackend(Backend):
 
+    def close(self):
+        pass
+
     def really_upload_block(self, block_id, block):
         dir_full = os.path.join(self.path, os.path.dirname(block_id))
         if not os.path.isdir(dir_full):
diff --git a/obnam/store.py b/obnam/store.py
index 905e66f8..4672af1c 100644
--- a/obnam/store.py
+++ b/obnam/store.py
@@ -36,6 +36,16 @@ class Store:
         self._context = context
         self._host = None
 
+    def close(self):
+        """Close connection to the store.
+        
+        You must not use this store instance for anything after
+        closing it.
+        
+        """
+        
+        self._context.be.close()
+
     def get_host_block(self):
         """Return current host block, or None if one is not known.