Revised file-upload logic, client can now upload multiple files OR large files. Fixes #23

Author: bschroeter

meorg_client/cli.py

@@ -4,7 +4,6 @@
 import meorg_client.utilities as mcu
 import os
 import sys
-from inspect import getmembers
 import getpass
 from pathlib import Path
 import json
@@ -21,20 +20,19 @@ def _get_client():
     # Get the dev-mode flag from the environment, better than passing the dev flag everywhere.
     dev_mode = os.getenv("MEORG_DEV_MODE", "0") == "1"
 
-    credentials = mcu.get_user_data_filepath('credentials.json')
-    credentials_dev = mcu.get_user_data_filepath('credentials-dev.json')
+    credentials = mcu.get_user_data_filepath("credentials.json")
+    credentials_dev = mcu.get_user_data_filepath("credentials-dev.json")
 
     # In dev mode and the configuration file exists
     if dev_mode and credentials_dev.is_file():
-        credentials = mcu.load_user_data('credentials-dev.json')
-    
+        credentials = mcu.load_user_data("credentials-dev.json")
+
     # In dev mode and it doesn't (i.e. Actions)
     elif dev_mode and not credentials_dev.is_file():
         credentials = dict(
-            email=os.getenv('MEORG_EMAIL'),
-            password=os.getenv('MEORG_PASSWORD')
+            email=os.getenv("MEORG_EMAIL"), password=os.getenv("MEORG_PASSWORD")
         )
-    
+
     # Production credentials
     else:
         credentials = mcu.load_user_data("credentials.json")
@@ -65,11 +63,10 @@ def _call(func, **kwargs):
     try:
         return func(**kwargs)
     except Exception as ex:
-        
         click.echo(ex.msg, err=True)
 
         # Bubble up the exception
-        if os.getenv('MEORG_DEV_MODE') == '1':
+        if os.getenv("MEORG_DEV_MODE") == "1":
             raise
 
         sys.exit(1)
@@ -86,7 +83,7 @@ def cli():
     pass
 
 
-@click.command('list')
+@click.command("list")
 def list_endpoints():
     """
     List the available endpoints for the server.
@@ -107,27 +104,8 @@ def list_endpoints():
             click.echo(out)
 
 
-@click.command('status')
-@click.argument("id")
-def file_status(id):
-    """
-    Check the file status based on the job ID from file-upload.
-
-    Prints the true file ID or a status.
-    """
-    client = _get_client()
-    response_data = _call(client.get_file_status, id=id).get("data")
-
-    # If the file is complete (transferred to object store), get the true ID
-    if response_data.get("status") == "complete":
-        file_id = response_data.get("files")[0].get("file")
-        click.echo(file_id)
-    else:
-        click.echo("Pending")
-
-
-@click.command('upload')
-@click.argument("file_path")
+@click.command("upload")
+@click.argument("file_path", nargs=-1)
 def file_upload(file_path):
     """
     Upload a file to the server.
@@ -137,12 +115,13 @@ def file_upload(file_path):
     client = _get_client()
 
     # Upload the file, get the job ID
-    response = _call(client.upload_file, file_path=file_path)
-    job_id = response.get("data").get("jobId")
-    click.echo(job_id)
+    response = _call(client.upload_file, file_path=list(file_path))
+    files = response.get("data").get("files")
+    for f in files:
+        click.echo(f.get("file"))
 
 
-@click.command('list')
+@click.command("list")
 @click.argument("id")
 def file_list(id):
     """
@@ -157,7 +136,7 @@ def file_list(id):
         click.echo(f)
 
 
-@click.command('attach')
+@click.command("attach")
 @click.argument("file_id")
 @click.argument("output_id")
 def file_attach(file_id, output_id):
@@ -166,12 +145,12 @@ def file_attach(file_id, output_id):
     """
     client = _get_client()
 
-    response = _call(client.attach_files_to_model_output, id=output_id, files=[file_id])
+    _ = _call(client.attach_files_to_model_output, id=output_id, files=[file_id])
 
     click.echo("SUCCESS")
 
 
-@click.command('start')
+@click.command("start")
 @click.argument("id")
 def analysis_start(id):
     """
@@ -188,7 +167,7 @@ def analysis_start(id):
         click.echo(analysis_id)
 
 
-@click.command('status')
+@click.command("status")
 @click.argument("id")
 def analysis_status(id):
     """
@@ -250,22 +229,25 @@ def initialise(dev=False):
 
 
 # Add groups for nested subcommands
-@click.group('endpoints', help='API endpoint commands.')
+@click.group("endpoints", help="API endpoint commands.")
 def cli_endpoints():
     pass
 
-@click.group('file', help='File commands.')
+
+@click.group("file", help="File commands.")
 def cli_file():
     pass
 
-@click.group('analysis', help='Analysis commands.')
+
+@click.group("analysis", help="Analysis commands.")
 def cli_analysis():
     pass
 
+
 # Add file commands
 cli_file.add_command(file_list)
 cli_file.add_command(file_upload)
-cli_file.add_command(file_status)
+# cli_file.add_command(file_status)
 cli_file.add_command(file_attach)
 
 # Add endpoint commands

meorg_client/client.py

@@ -9,6 +9,7 @@
 import meorg_client.endpoints as endpoints
 import meorg_client.exceptions as mx
 import mimetypes as mt
+import io
 
 
 class Client:
@@ -213,45 +214,67 @@ def logout(self):
             self.headers.pop("X-User-Id", None)
             self.headers.pop("X-Auth-Token", None)
 
-    def get_file_status(self, id: str) -> Union[dict, requests.Response]:
-        """Get the file status.
+    def upload_file(
+        self,
+        file_path: Union[str, list],
+        file_obj: Union[io.BufferedReader, list] = None,
+    ) -> Union[dict, requests.Response]:
+        """Upload a file.
 
         Parameters
         ----------
-        id : str
-            Job ID of the file.
+        file_path : str or list
+            Path to the file.
+        file_obj : io.BufferedReader, optional
+            File object (handle) to allow direct supply of file object, by default None
 
         Returns
         -------
         Union[dict, requests.Response]
             Response from ME.org.
         """
-        return self._make_request(
-            method=mcc.HTTP_GET, endpoint=endpoints.FILE_STATUS, url_params=dict(id=id)
-        )
 
-    def upload_file(self, file_path: str) -> Union[dict, requests.Response]:
-        """Upload a file.
+        payload = list()
 
-        Parameters
-        ----------
-        file_path : str
-            Path to the file.
+        # Cast as list for iterative upload
+        if not isinstance(file_path, list):
+            file_path = [file_path]
 
-        Returns
-        -------
-        Union[dict, requests.Response]
-            Response from ME.org.
-        """
-        # Get the filename and extension
-        filename = os.path.basename(file_path)
-        ext = filename.split(".")[-1]
+        # Payload assembly
+        if file_obj is not None:
+            # Cast as a list for iterative upload
+            if not isinstance(file_obj, list):
+                file_objs = [file_obj]
 
-        # Get the MIME type (raises a KeyError if it is unknown)
-        mimetype = mt.types_map[f".{ext}"]
+            if len(file_objs) != len(file_path):
+                raise ValueError("Supplied file paths and file objects do not match")
 
-        # Assemble the file payload
-        payload = dict(file=(filename, open(file_path, "rb"), mimetype))
+            for ix, file_obj in enumerate(file_objs):
+                if not isinstance(file_obj, io.BufferedReader) and not isinstance(
+                    file_obj, io.BytesIO
+                ):
+                    raise TypeError(
+                        f"Supplied file object {ix} is not an io.BufferedReader or io.BytesIO."
+                    )
+
+                _file_path = file_path[ix]
+
+                # Get the filename and extension
+                filename = os.path.basename(_file_path)
+                ext = filename.split(".")[-1]
+                mimetype = mt.types_map[f".{ext}"]
+
+                payload.append(("file", (filename, file_obj, mimetype)))
+
+        else:
+            for _file_path in file_path:
+                # Get the filename and extension
+                filename = os.path.basename(_file_path)
+                ext = filename.split(".")[-1]
+                mimetype = mt.types_map[f".{ext}"]
+                file_obj = open(_file_path, "rb")
+
+                payload.append(("file", (filename, file_obj, mimetype)))
 
         return self._make_request(
             method=mcc.HTTP_POST,

meorg_client/endpoints.py

@@ -9,7 +9,7 @@
 
 # Files
 FILE_LIST = "modeloutput/{id}/files"
-FILE_UPLOAD = "files"
+FILE_UPLOAD = "upload"
 FILE_STATUS = "files/status/{id}"
 
 # Analysis

meorg_client/tests/test_cli.py

@@ -27,23 +27,38 @@ def test_file_upload(runner):
     assert result.exit_code == 0
 
     # Add the job_id to the store for the next test
-    store.set("job_id", result.output.strip())
+    store.set("file_id", result.output.strip())
 
     # Let it wait for a short while, allow the server to transfer to object store.
     time.sleep(5)
 
 
-def test_file_status(runner):
-    """Test file-status via CLI."""
+def test_file_multiple(runner):
+    """Test file-upload via CLI."""
 
-    # Get the file ID based on the job ID
-    job_id = store.get("job_id")
-    result = runner.invoke(cli.file_status, [job_id])
+    # Upload a tiny test file
+    filepath = os.path.join(mu.get_installed_data_root(), "test/test.txt")
+    result = runner.invoke(cli.file_upload, [filepath, filepath])
     assert result.exit_code == 0
-    assert result.output != "Pending"
 
-    # Add file_id to the store for the next test
-    store.set("file_id", result.output.strip())
+    # Add the job_id to the store for the next test
+    store.set("file_ids", result.output.strip())
+
+    # Let it wait for a short while, allow the server to transfer to object store.
+    time.sleep(5)
+
+
+# def test_file_status(runner):
+#     """Test file-status via CLI."""
+
+#     # Get the file ID based on the job ID
+#     job_id = store.get("job_id")
+#     result = runner.invoke(cli.file_status, [job_id])
+#     assert result.exit_code == 0
+#     assert result.output != "Pending"
+
+#     # Add file_id to the store for the next test
+#     store.set("file_id", result.output.strip())
 
 
 def test_file_list(runner):

meorg_client/tests/test_client.py

@@ -4,6 +4,7 @@
 from meorg_client.client import Client
 import meorg_client.utilities as mu
 from conftest import store
+import io
 
 
 def _get_authenticated_client():
@@ -55,14 +56,19 @@ def test_upload_file(client):
     store.set("file_upload", response)
 
 
-def test_file_status(client):
-    # Get the response
-    job_id = store.get("file_upload").get("data").get("jobId")
+def test_upload_file_multiple(client):
+    """Test the uploading of a file."""
+    # Upload the file.
+    filepath = os.path.join(mu.get_installed_data_root(), "test/test.txt")
 
-    response = client.get_file_status(job_id)
+    # Upload the file
+    response = client.upload_file([filepath, filepath])
+
+    # Make sure it worked
     assert client.success()
 
-    store.set("file_status", response)
+    # Store the response.
+    store.set("file_upload_multiple", response)
 
 
 def test_file_list(client):
@@ -73,7 +79,8 @@ def test_file_list(client):
 
 def test_attach_files_to_model_output(client):
     # Get the file id from the job id
-    file_id = store.get("file_status").get("data").get("files")[0].get("file")
+    # file_id = store.get("file_status").get("data").get("files")[0].get("file")
+    file_id = store.get("file_upload").get("data").get("files")[0].get("file")
 
     # Attach it to the model output
     _ = client.attach_files_to_model_output(client._model_output_id, [file_id])
@@ -98,3 +105,22 @@ def test_logout(client):
     """Test logout."""
     client.logout()
     assert "X-Auth-Token" not in client.headers.keys()
+
+
+def test_upload_file_large(client):
+    """Test the uploading of a large-ish file."""
+
+    # Create an in-memory 10mb file
+    size = 10000000
+    with io.BytesIO() as buffer:
+        buffer.write(bytearray(os.urandom(size)))
+        buffer.seek(0)
+
+        # Upload the file.
+        filepath = os.path.join(mu.get_installed_data_root(), "test/test.txt")
+
+        # Upload the file
+        _ = client.upload_file(file_path=filepath, file_obj=buffer)
+
+        # Make sure it worked
+        assert client.success()

meorg_client/utilities.py

@@ -60,6 +60,7 @@ def get_user_data_filepath(filename):
     """Get the filepath to the user file."""
     return Path.home() / ".meorg" / filename
 
+
 def load_user_data(filename):
     """Load data from the user's home directory."""
     filepath = get_user_data_filepath(filename)