Skip to content

Commit

Permalink
Log progress when putting file through SFTP (#985)
Browse files Browse the repository at this point in the history
  • Loading branch information
austinweisgrau authored Feb 15, 2024
1 parent 7dfdc34 commit c9869be
Showing 1 changed file with 33 additions and 9 deletions.
42 changes: 33 additions & 9 deletions parsons/sftp/sftp.py
Original file line number Diff line number Diff line change
@@ -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__)

Expand Down Expand Up @@ -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
Expand All @@ -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):
"""
Expand All @@ -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):
"""
Expand Down

0 comments on commit c9869be

Please sign in to comment.