Merge branch 'dev'

README.md

@@ -1,20 +1,16 @@
 # Summary
 
-At the moment, the package is mainly just a wrapper around FTP and SFTP libraries. It's used to upload files to VIP's GDI system. But could be expanded on with SQL functionality or if VIP ever added a genuine API in the future.
-
-As is, the point of the package is largely just defining how to upload a file to VIP in one place so it doesn't need to be rewritten every time a repo involving VIP is created and can instead just be imported.
-
-There are also pandera models defined for an order and an invoice.
+At the moment, the package is mainly just a wrapper around FTP and SFTP libraries. It's used to upload and download files from VIP's GDI/GDI2 system. But could be expanded on with SQL functionality or if VIP ever added a genuine API in the future.
 
 ## Installation
 
 `pip install vipwrap`
 
 ## Usage
 
-Currently consists of one function.
+Contains two modules, `gdi` and `models`. `gdi` covers the actual connection to VIP's GDI servers and the downloading/uploading of files. Largely just a wrapper arounf ftp libraries. Nothing VIP-specific happening in it. `models` contains two pandera models that describe order or sales history data. Has not been fully built out. At the moment, can be used for dataframe validation.
 
-### Send File to VIP
+### Send File to GDI
 
 Takes standard parameters expected of uploading an existing local file to a remote location via FTP or SFTP.
 
@@ -25,5 +21,20 @@ Takes standard parameters expected of uploading an existing local file to a remo
 | port | int | the SFTP server port, not used with FTP |
 | user | str | username to authenticate with |
 | password | str | password to authenticate with |
-| folder | str | The base folder location to upload the file to |
+| folder | str | The base folder location to upload the file to, usually "/in/" or "/TO_GDI/" |
 | file | IO[str] | File stream being uploaded. Usually the output of an open() function |
+
+### Download Files from GDI
+
+Takes standard paramers expected for connecting to a remote FTP/SFTP server as well as the file string to search for. This is being used to mass-download files that start with a given string.
+
+| parameter | type | description |
+| - | - | - |
+| ftp_method | str | 'ftp' or 'sftp' |
+| host | str | the FTP server host |
+| port | int | the FTP/SFTP server port |
+| user | str | username to authenticate with |
+| password | str | password to authenticate with |
+| folder | str | the folder to download the files from, usually "/out/" or "/FROM_VIP/" |
+| file_string | str | string in file name to search for |
+| delete_after_download | bool | whether to delete the files from the remote location once downloaded. |

setup.py

@@ -8,7 +8,7 @@
 
 setup(
     name="vipwrap",
-    version="0.1.1",
+    version="0.1.2",
     description="Python package that allows interacting with VIP.",
     long_description=long_description,
     long_description_content_type="text/markdown",

tests/test_download.py

@@ -0,0 +1,19 @@
+import sys
+
+sys.path.insert(0, "")
+
+from vipwrap import gdi
+
+
+def test_download_ftp():
+    # downloaded_files = gdi.download_files_from_gdi(
+    #     "ftp", "localhost", 10021, "myuser", "mypass", "/", "test", False
+    # )
+    downloaded_files = gdi.download_files_from_gdi(
+        "sftp", "localhost", 10022, "foo", "pass", "/upload/", "test", False
+    )
+    pass
+
+
+if __name__ == "__main__":
+    test_download_ftp()

vipwrap/__init__.py

@@ -2,4 +2,4 @@
 Initialization file for vipwrap package
 """
 
-from . import models, upload_handling
+from . import gdi, models

vipwrap/gdi.py

@@ -0,0 +1,218 @@
+"""
+Module handles the actual uploading of files to VIP's GDI/GDI2 server.
+The module is largely just a wrapper around the paramiko and ftplib libraries
+for SFTP and FTP uploads, respectively.
+"""
+
+import os
+import time
+from ftplib import FTP
+from typing import IO, Literal
+
+import paramiko
+
+
+def connect_sftp(host: str, port: int, user: str, password: str):
+    """
+    Connect to SFTP server
+    """
+    ssh = paramiko.SSHClient()
+    ssh.set_missing_host_key_policy(paramiko.AutoAddPolicy())
+    ssh.connect(
+        host,
+        port=port,
+        username=user,
+        password=password,
+        timeout=30,
+    )
+    return ssh
+
+
+def upload_sftp(
+    host: str, port: int, user: str, password: str, folder: str, file: IO[str]
+):
+    """
+    Upload file to SFTP server
+    """
+
+    with connect_sftp(host, port, user, password) as ssh:
+        with ssh.open_sftp() as sftp:
+            sftp.get_channel().settimeout(60)  # type: ignore
+            remote_path = folder + file.name
+            sftp.put(file.name, remote_path)
+            local_file_size = os.path.getsize(file.name)
+            remote_file_size = sftp.stat(remote_path).st_size
+            print("Local file size: ", local_file_size)
+            print("Remote file size: ", remote_file_size)
+            if local_file_size != remote_file_size:
+                raise ValueError(
+                    f"Error, file sizes do not match: {local_file_size} vs {remote_file_size}"
+                )
+
+
+def download_sftp(
+    host: str,
+    port: int,
+    user: str,
+    password: str,
+    folder: str,
+    file_string: str,
+    download_after_delete: bool = False,
+) -> list[str]:
+    """
+    Downloads all files in a folder on the VIP GDI server whose filenames start
+    with the given string.
+    """
+    downloaded_files = []
+    with connect_sftp(host, port, user, password) as ssh:
+        with ssh.open_sftp() as sftp:
+            sftp.get_channel().settimeout(60)  # type: ignore
+            files = sftp.listdir(folder)
+            for f in files:
+                if f.startswith(file_string):
+                    remote_path = folder + f
+                    local_path = f
+                    sftp.get(remote_path, local_path)
+                    downloaded_files.append(os.path.abspath(local_path))
+                    if download_after_delete:
+                        sftp.remove(remote_path)
+    return downloaded_files
+
+
+def download_ftp(
+    host: str,
+    port: int,
+    user: str,
+    password: str,
+    folder: str,
+    file_string: str,
+    download_after_delete: bool = False,
+) -> list[str]:
+    """
+    Downloads all files in a folder on the VIP GDI server whose filenames start
+    with the given string.
+    """
+    ftp = FTP()
+    downloaded_files = []
+    try:
+        ftp.connect(host, port)
+        ftp.login(user, password)
+        files = ftp.nlst(folder)
+        for f in files:
+            filename = f.lstrip("/")
+            if filename.startswith(file_string):
+                with open(filename, "wb") as local_file:
+                    ftp.retrbinary("RETR " + filename, local_file.write)
+                downloaded_files.append(os.path.abspath(filename))
+                if download_after_delete:
+                    ftp.delete(filename)
+    finally:
+        ftp.quit()
+
+    return downloaded_files
+
+
+def delete_sftp(
+    host: str, port: int, user: str, password: str, folder: str, filename: str
+) -> None:
+    """
+    Deletes a file in a folder on the VIP GDI server whose filename
+    matches the given string.
+    """
+    with connect_sftp(host, port, user, password) as ssh:
+        with ssh.open_sftp() as sftp:
+            sftp.get_channel().settimeout(60)  # type: ignore
+            files = sftp.listdir(folder)
+            for f in files:
+                if f == filename:
+                    remote_path = folder + f
+                    sftp.remove(remote_path)
+
+
+def upload_ftp(host: str, user: str, password: str, folder: str, file: IO[str]):
+    """
+    Upload file to FTP server
+    """
+
+    with FTP(host) as ftp:
+        ftp.login(user, password)
+        with open(file.name, "rb") as f:
+            ftp.storbinary("STOR " + folder + file.name, f)
+
+
+def send_file_to_gdi(
+    ftp_method: Literal["sftp", "ftp"],
+    host: str,
+    port: int,
+    user: str,
+    password: str,
+    folder: str,
+    file: IO[str],
+) -> None:
+    """
+    Sends file to VIP GDI server via FTP or SFTP, depending on which method specified.
+    The transfer is retried 3 times if it fails.
+
+    ftp_method: whether to upload via SFTP or FTP
+    host: the hostname of the FTP/SFTP server
+    port: the port number of the SFTP server
+    user: the username for the FTP/SFTP server
+    password: the password for the FTP/SFTP server
+    folder: the folder to upload the file to
+    file: file object to upload
+    """
+
+    for attempt in range(3):
+        try:
+            if ftp_method == "sftp":
+                upload_sftp(host, port, user, password, folder, file)
+            elif ftp_method == "ftp":
+                upload_ftp(host, user, password, folder, file)
+            else:
+                print("Error, invalid FTP method")
+                raise ValueError("Error, invalid FTP method")
+            break
+        except Exception as e:
+            print(f"Attempt {attempt + 1}: Error sending file to VIP: {e}")
+            if attempt < 2:
+                time.sleep(60)  # Wait before retrying
+            else:
+                raise  # Reraise the exception or handle it as needed
+
+
+def download_files_from_gdi(
+    ftp_method: Literal["sftp", "ftp"],
+    host: str,
+    port: int,
+    user: str,
+    password: str,
+    folder: str,
+    file_string: str,
+    delete_after_download: bool = False,
+) -> list[str]:
+    """
+    Downloads all files in a folder on the VIP GDI server whose filenames start
+    with the given string.
+
+    Returns a list of paths to the downloaded files.
+
+    ftp_method: whether to download via SFTP or FTP
+    host: the hostname of the FTP/SFTP server
+    port: the port number of the SFTP
+    user: the username for the FTP/SFTP server
+    password: the password for the FTP/SFTP server
+    folder: the folder to download the files from
+    file_string: the string that the filenames must start with
+    delete_after_download: whether to delete the files from the server after downloading
+    """
+    if ftp_method == "sftp":
+        return download_sftp(
+            host, port, user, password, folder, file_string, delete_after_download
+        )
+    elif ftp_method == "ftp":
+        return download_ftp(
+            host, port, user, password, folder, file_string, delete_after_download
+        )
+    else:
+        print("Error, invalid FTP method")
+        raise ValueError("Error, invalid FTP method")