From 808a2c13cbbc263157a98241fc501eaf2ed2f270 Mon Sep 17 00:00:00 2001
From: Austin Weisgrau <austinweisgrau@gmail.com>
Date: Mon, 5 Feb 2024 15:53:13 -0800
Subject: [PATCH] Log progress when putting file through SFTP

---
 parsons/sftp/sftp.py | 42 +++++++++++++++++++++++++++++++++---------
 1 file changed, 33 insertions(+), 9 deletions(-)

diff --git a/parsons/sftp/sftp.py b/parsons/sftp/sftp.py
index 4b4301357b..e1b441477a 100644
--- a/parsons/sftp/sftp.py
+++ b/parsons/sftp/sftp.py
@@ -1,12 +1,13 @@
+import logging
+import re
 from contextlib import contextmanager
 from stat import S_ISDIR, S_ISREG
-import logging
+
 import paramiko
-import re
 
-from parsons.utilities import files as file_utilities
 from parsons.etl import Table
 from parsons.sftp.utilities import connect
+from parsons.utilities import files as file_utilities
 
 logger = logging.getLogger(__name__)
 
@@ -245,7 +246,22 @@ def get_table(self, remote_path, connection=None):
 
         return Table.from_csv(self.get_file(remote_path, connection=connection))
 
-    def put_file(self, local_path, remote_path, connection=None):
+    def _convert_bytes_to_megabytes(self, size_in_bytes: int) -> int:
+        result = int(size_in_bytes / (1024 * 1024))
+        return result
+
+    def _progress(self, transferred: int, to_be_transferred: int) -> None:
+        """Return progress every 5 MB"""
+        if self._convert_bytes_to_megabytes(transferred) % 5 != 0:
+            return
+        logger.info(
+            f"Transferred: {self._convert_bytes_to_megabytes(transferred)} MB \t"
+            f"out of: {self._convert_bytes_to_megabytes(to_be_transferred)} MB"
+        )
+
+    def put_file(
+        self, local_path: str, remote_path: str, connection=None, verbose: bool = True
+    ) -> None:
         """
         Put a file on the SFTP server
 
@@ -256,11 +272,18 @@ def put_file(self, local_path, remote_path, connection=None):
                 The remote path of the new file
             connection: obj
                 An SFTP connection object
+            verbose: bool
+                Log progress every 5MB. Defaults to True.
         """
+        if verbose:
+            callback = self._progress
+        else:
+            callback = None
         if connection:
-            connection.put(local_path, remote_path)
-        with self.create_connection() as connection:
-            connection.put(local_path, remote_path)
+            connection.put(local_path, remote_path, callback=callback)
+        else:
+            with self.create_connection() as connection:
+                connection.put(local_path, remote_path, callback=callback)
 
     def remove_file(self, remote_path, connection=None):
         """
@@ -275,8 +298,9 @@ def remove_file(self, remote_path, connection=None):
 
         if connection:
             connection.remove(remote_path)
-        with self.create_connection() as connection:
-            connection.remove(remote_path)
+        else:
+            with self.create_connection() as connection:
+                connection.remove(remote_path)
 
     def get_file_size(self, remote_path, connection=None):
         """