feat(eng-1375): Support large file uploads (#143)

* feat(eng-1375): Support large file uploads

* update changelog

Author: chrisimcevoy

CHANGELOG.md

@@ -9,6 +9,10 @@ and this project adheres to [Semantic Versioning](http://semver.org/spec/v2.0.0.
 
 ## [Unreleased]
 
+### Added
+
+- Added support for large file uploads ([#143](https://github.com/cloudsmith-io/cloudsmith-cli/pull/143))
+
 ### Fixed
 
 - Removed more unused dependencies relating to python 2.7 compatibility ([#142](https://github.com/cloudsmith-io/cloudsmith-cli/pull/142))

cloudsmith_cli/cli/commands/push.py

@@ -1,5 +1,5 @@
 """CLI/Commands - Push packages."""
-
+import math
 import os
 import time
 from datetime import datetime
@@ -9,6 +9,8 @@
 from ...core import utils
 from ...core.api.exceptions import ApiException
 from ...core.api.files import (
+    CHUNK_SIZE,
+    multi_part_upload_file,
     request_file_upload,
     upload_file as api_upload_file,
     validate_request_file_upload,
@@ -56,6 +58,10 @@ def upload_file(ctx, opts, owner, repo, filepath, skip_errors, md5_checksum):
     filename = click.format_filename(filepath)
     basename = os.path.basename(filename)
 
+    filesize = utils.get_file_size(filepath=filename)
+    projected_chunks = math.floor(filesize / CHUNK_SIZE) + 1
+    is_multi_part_upload = projected_chunks > 1
+
     click.echo(
         "Requesting file upload for %(filename)s ... "
         % {"filename": click.style(basename, bold=True)},
@@ -68,33 +74,58 @@ def upload_file(ctx, opts, owner, repo, filepath, skip_errors, md5_checksum):
     ):
         with maybe_spinner(opts):
             identifier, upload_url, upload_fields = request_file_upload(
-                owner=owner, repo=repo, filepath=filename, md5_checksum=md5_checksum
+                owner=owner,
+                repo=repo,
+                filepath=filename,
+                md5_checksum=md5_checksum,
+                is_multi_part_upload=is_multi_part_upload,
             )
 
     click.secho("OK", fg="green")
 
     context_msg = "Failed to upload file!"
     with handle_api_exceptions(ctx, opts=opts, context_msg=context_msg):
-        filesize = utils.get_file_size(filepath=filename)
-
         label = f"Uploading {click.style(basename, bold=True)}:"
 
-        with click.progressbar(
-            length=filesize,
-            label=label,
-            fill_char=click.style("#", fg="green"),
-            empty_char=click.style("-", fg="red"),
-        ) as pb:
-
-            def progress_callback(monitor):
-                pb.update(monitor.bytes_read)
-
-            api_upload_file(
-                upload_url=upload_url,
-                upload_fields=upload_fields,
-                filepath=filename,
-                callback=progress_callback,
-            )
+        if not is_multi_part_upload:
+            # We can upload the whole file in one go.
+            with click.progressbar(
+                length=filesize,
+                label=label,
+                fill_char=click.style("#", fg="green"),
+                empty_char=click.style("-", fg="red"),
+            ) as pb:
+
+                def progress_callback(monitor):
+                    pb.update(monitor.bytes_read)
+
+                api_upload_file(
+                    upload_url=upload_url,
+                    upload_fields=upload_fields,
+                    filepath=filename,
+                    callback=progress_callback,
+                )
+        else:
+            # The file is sufficiently large that we need to upload in chunks.
+            with click.progressbar(
+                length=projected_chunks,
+                label=label,
+                fill_char=click.style("#", fg="green"),
+                empty_char=click.style("-", fg="red"),
+            ) as pb:
+
+                def progress_callback():
+                    pb.update(1)
+
+                multi_part_upload_file(
+                    opts=opts,
+                    upload_url=upload_url,
+                    owner=owner,
+                    repo=repo,
+                    filepath=filename,
+                    callback=progress_callback,
+                    upload_id=identifier,
+                )
 
     return identifier
 

cloudsmith_cli/cli/tests/commands/test_package_commands.py

@@ -0,0 +1,80 @@
+import json
+import time
+
+import pytest
+
+from ....core.api.files import CHUNK_SIZE
+from ...commands.delete import delete
+from ...commands.list_ import list_
+from ...commands.push import push
+from ...commands.status import status
+from ..utils import random_str
+
+
+@pytest.mark.usefixtures("set_api_key_env_var", "set_api_host_env_var")
+@pytest.mark.parametrize(
+    "filesize",
+    [
+        1,  # A tiny file that will be uploaded using python bindings as a one-off post.
+        CHUNK_SIZE * 3,  # A large file that will be uploaded to S3 in multiple chunks.
+    ],
+)
+def test_push_and_delete_raw_package(
+    runner, organization, tmp_repository, tmp_path, filesize
+):
+    # List packages again - should be empty.
+    org_repo = f'{organization}/{tmp_repository["slug"]}'
+    result = runner.invoke(
+        list_, args=["pkgs", org_repo, "-F", "json"], catch_exceptions=False
+    )
+    data = json.loads(result.output)["data"]
+    assert len(data) == 0
+
+    # Create a file of the requested size.
+    pkg_file = tmp_path / f"{random_str()}.txt"
+    with open(pkg_file, "wb") as f:
+        # Fill the file with null bytes.
+        f.truncate(filesize)
+
+    # Push it to cloudsmith as a raw package using the push command.
+    runner.invoke(
+        push, args=["raw", org_repo, str(pkg_file.resolve())], catch_exceptions=False
+    )
+
+    # List packages, check that it is there.
+    result = runner.invoke(
+        list_, args=["pkgs", org_repo, "-F", "json"], catch_exceptions=False
+    )
+    data = json.loads(result.output)["data"]
+    assert len(data) == 1
+    small_file_data = data[0]
+    assert small_file_data["filename"] == pkg_file.name
+
+    # Wait for the package to sync.
+    org_repo_package = f"{org_repo}/{small_file_data['slug']}"
+    for _ in range(10):
+        time.sleep(5)
+        result = runner.invoke(status, args=[org_repo_package], catch_exceptions=False)
+        if "Fully Synchronised" in result.output:
+            break
+    else:
+        raise TimeoutError("Test timed out waiting for package sync")
+
+    # Delete the package.
+    runner.invoke(delete, args=["-y", org_repo_package], catch_exceptions=False)
+
+    # Wait for package deletion to take effect.
+    for _ in range(10):
+        time.sleep(5)
+        result = runner.invoke(status, args=[org_repo_package], catch_exceptions=False)
+        if "status: 404 - Not Found" in result.output:
+            break
+    else:
+        raise TimeoutError("Test timed out waiting for package deletion")
+
+    # List packages again - should be empty.
+    result = runner.invoke(
+        list_, args=["pkgs", org_repo, "-F", "json"], catch_exceptions=False
+    )
+    data = json.loads(result.output)["data"]
+    assert len(data) == 0

cloudsmith_cli/cli/tests/conftest.py

@@ -3,7 +3,9 @@
 import click.testing
 import pytest
 
-from cloudsmith_cli.cli.tests.utils import random_str
+from ...core.api.init import initialise_api
+from ...core.api.repos import create_repo, delete_repo
+from .utils import random_str
 
 
 def _get_env_var_or_skip(key):
@@ -47,10 +49,9 @@ def organization():
 
 
 @pytest.fixture()
-def tmp_repository(organization):
+def tmp_repository(organization, api_host, api_key):
     """Yield a temporary repository."""
-    from ...core.api.repos import create_repo, delete_repo
-
+    initialise_api(host=api_host, key=api_key)
     repo_data = create_repo(organization, {"name": random_str()})
     yield repo_data
     delete_repo(organization, repo_data["slug"])

cloudsmith_cli/core/api/files.py

@@ -13,6 +13,8 @@
 from .exceptions import ApiException, catch_raise_api_exception
 from .init import get_api_client
 
+CHUNK_SIZE = 1024 * 1024 * 100
+
 
 def get_files_api():
     """Get the files API client."""
@@ -35,16 +37,24 @@ def validate_request_file_upload(owner, repo, filepath, md5_checksum=None):
     return md5_checksum
 
 
-def request_file_upload(owner, repo, filepath, md5_checksum=None):
+def request_file_upload(
+    owner, repo, filepath, md5_checksum=None, is_multi_part_upload=False
+):
     """Request a new package file upload (for creating packages)."""
     client = get_files_api()
     md5_checksum = md5_checksum or calculate_file_md5(filepath)
 
+    method = "put_parts" if is_multi_part_upload else "post"
+
     with catch_raise_api_exception():
         data, _, headers = client.files_create_with_http_info(
             owner=owner,
             repo=repo,
-            data={"filename": os.path.basename(filepath), "md5_checksum": md5_checksum},
+            data={
+                "filename": os.path.basename(filepath),
+                "md5_checksum": md5_checksum,
+                "method": method,
+            },
         )
 
     # pylint: disable=no-member
@@ -82,3 +92,40 @@ def upload_file(upload_url, upload_fields, filepath, callback=None):
         raise ApiException(
             resp.status_code, headers=exc.response.headers, body=exc.response.content
         )
+
+
+def multi_part_upload_file(
+    opts, upload_url, owner, repo, filepath, callback, upload_id
+):
+    with open(filepath, "rb") as f:
+        chunk_number = 1
+        session = create_requests_session()
+        headers = {"X-Api-Key": opts.api_key}
+        while chunk := f.read(CHUNK_SIZE):
+            resp = session.put(
+                upload_url,
+                headers=headers,
+                data=chunk,
+                params={
+                    "upload_id": upload_id,
+                    "part_number": chunk_number,
+                },
+            )
+            try:
+                resp.raise_for_status()
+            except requests.RequestException as exc:
+                raise ApiException(
+                    resp.status_code,
+                    headers=exc.response.headers,
+                    body=exc.response.content,
+                )
+            callback()
+            chunk_number += 1
+
+    api = get_files_api()
+    api.files_complete(
+        owner,
+        repo,
+        identifier=upload_id,
+        data={"upload_id": upload_id, "complete": True},
+    )