Merge pull request learningequality#5995 from rtibbles/resumable_uploads

Add opt-in resumable uploads to the file upload URL endpoint

Author: rtibbles

contentcuration/contentcuration/tests/test_storage_common.py

@@ -1,4 +1,3 @@
-import codecs
 import hashlib
 from datetime import timedelta
 from io import BytesIO
@@ -12,6 +11,7 @@
 
 from .base import StudioTestCase
 from contentcuration.models import generate_object_storage_name
+from contentcuration.utils.gcs_storage import GoogleCloudStorage
 from contentcuration.utils.storage_common import _get_gcs_presigned_put_url
 from contentcuration.utils.storage_common import determine_content_type
 from contentcuration.utils.storage_common import get_presigned_upload_url
@@ -79,7 +79,7 @@ def test_raises_error(self):
         with pytest.raises(UnknownStorageBackendError):
             get_presigned_upload_url(
                 "nice",
-                "err",
+                "d41d8cd98f00b204e9800998ecf8427e",
                 5,
                 0,
                 storage=self.STORAGE,
@@ -157,6 +157,23 @@ def test_generate_signed_url_called_with_required_arguments(self):
             content_type=mimetype,
         )
 
+    def test_create_resumable_session_pins_md5_size_and_returns_url(self):
+        storage = GoogleCloudStorage(self.client, "bucket")
+        blob = self.client.get_bucket.return_value.blob.return_value
+        blob.create_resumable_upload_session.return_value = "https://session.url"
+
+        url = storage.create_resumable_upload_session(
+            "storage/a/b/abc.jpg",
+            "d41d8cd98f00b204e9800998ecf8427e",
+            2048,
+        )
+
+        assert url == "https://session.url"
+        assert blob.md5_hash == "1B2M2Y8AsgTpgAmY7PhCfg=="  # hex checksum, b64-encoded
+        assert blob.content_type == "image/jpeg"
+        assert blob.metadata == {"declared-size": "2048"}
+        blob.create_resumable_upload_session.assert_called_once()
+
 
 class S3StoragePresignedURLUnitTestCase(StudioTestCase):
     """
@@ -177,7 +194,12 @@ def test_returns_string_if_inputs_are_valid(self):
 
         # use a real connection here as a sanity check
         ret = get_presigned_upload_url(
-            "a/b/abc.jpg", "aBc", 10, 1, storage=self.STORAGE, client=None
+            "a/b/abc.jpg",
+            "d41d8cd98f00b204e9800998ecf8427e",
+            10,
+            1,
+            storage=self.STORAGE,
+            client=None,
         )
         url = ret["uploadURL"]
 
@@ -189,19 +211,12 @@ def test_can_upload_file_to_presigned_url(self):
         """
         file_contents = b"blahfilecontents"
         file = BytesIO(file_contents)
-        # S3 expects a base64-encoded MD5 checksum
-        md5 = hashlib.md5(file_contents)
-        md5_checksum = md5.hexdigest()
-        md5_checksum_base64 = codecs.encode(
-            codecs.decode(md5_checksum, "hex"), "base64"
-        ).decode()
+        md5_checksum = hashlib.md5(file_contents).hexdigest()
 
         filename = "blahfile.jpg"
         filepath = generate_object_storage_name(md5_checksum, filename)
 
-        ret = get_presigned_upload_url(
-            filepath, md5_checksum_base64, 1000, len(file_contents)
-        )
+        ret = get_presigned_upload_url(filepath, md5_checksum, 1000, len(file_contents))
         url = ret["uploadURL"]
         content_type = ret["mimetype"]
 

contentcuration/contentcuration/tests/viewsets/test_file.py

@@ -1,4 +1,5 @@
 import uuid
+from unittest import mock
 
 from django.urls import reverse
 from le_utils.constants import content_kinds
@@ -12,6 +13,8 @@
 from contentcuration.tests.viewsets.base import generate_delete_event
 from contentcuration.tests.viewsets.base import generate_update_event
 from contentcuration.tests.viewsets.base import SyncTestMixin
+from contentcuration.viewsets.file import FileUploadURLSerializer
+from contentcuration.viewsets.file import MAX_NON_RESUMABLE_UPLOAD_SIZE
 from contentcuration.viewsets.sync.constants import CONTENTNODE
 from contentcuration.viewsets.sync.constants import FILE
 
@@ -546,6 +549,9 @@ def test_mismatched_preset_upload(self):
 
     def test_insufficient_storage(self):
         self.file["size"] = 100000000000000
+        self.file[
+            "resumable"
+        ] = True  # resumable bypasses the >500MB guard so this still exercises the quota (412) path
 
         self.client.force_authenticate(user=self.user)
         response = self.client.post(
@@ -590,6 +596,26 @@ def test_duration_zero(self):
 
         self.assertEqual(response.status_code, 400)
 
+    def test_fractional_size_rejected(self):
+        s = FileUploadURLSerializer(data={**self.file, "size": 1000.5})
+        assert not s.is_valid()
+        assert "size" in s.errors
+
+    def test_large_non_resumable_rejected(self):
+        self.file["size"] = MAX_NON_RESUMABLE_UPLOAD_SIZE + 1
+        self.client.force_authenticate(user=self.user)
+        resp = self.client.post(reverse("file-upload-url"), self.file, format="json")
+        assert resp.status_code == 400
+
+    def test_large_resumable_allowed(self):
+        self.user.disk_space = 10 * 1024 * 1024 * 1024
+        self.user.save()
+        self.file["size"] = MAX_NON_RESUMABLE_UPLOAD_SIZE + 1
+        self.file["resumable"] = True
+        self.client.force_authenticate(user=self.user)
+        resp = self.client.post(reverse("file-upload-url"), self.file, format="json")
+        assert resp.status_code == 200
+
 
 class ContentIDTestCase(SyncTestMixin, StudioAPITestCase):
     def setUp(self):
@@ -763,3 +789,60 @@ def test_content_id__thumbnails_dont_update_content_id(self):
         self.assertEqual(
             copied_node_content_id_before_upload, copied_node_content_id_after_upload
         )
+
+
+class ResumableUploadURLTestCase(StudioAPITestCase):
+    def setUp(self):
+        super(ResumableUploadURLTestCase, self).setUp()
+        self.user = testdata.user()
+        # Give user enough quota to handle resumable uploads
+        self.user.disk_space = 10 * 1024 * 1024 * 1024
+        self.user.save()
+        self.file = {
+            "size": 1000,
+            "checksum": uuid.uuid4().hex,
+            "name": "le_studio",
+            "file_format": file_formats.MP3,
+            "preset": format_presets.AUDIO,
+            "duration": 10.123,
+            "resumable": True,
+        }
+
+    @mock.patch("contentcuration.viewsets.file.default_storage")
+    def test_resumable_returns_session_when_not_stored(self, mock_storage):
+        mock_storage.get_stored_object_md5.return_value = None
+        mock_storage.create_resumable_upload_session.return_value = (
+            "https://session.url"
+        )
+        self.client.force_authenticate(user=self.user)
+        resp = self.client.post(reverse("file-upload-url"), self.file, format="json")
+        assert resp.status_code == 200
+        data = resp.json()
+        assert data["resumable"] is True
+        assert data["uploadURL"] == "https://session.url"
+        assert data["alreadyUploaded"] is False
+        assert "file" in data
+        assert data["file"]["id"]
+        mock_storage.create_resumable_upload_session.assert_called_once()
+
+    @mock.patch("contentcuration.viewsets.file.default_storage")
+    def test_resumable_skips_when_already_stored(self, mock_storage):
+        mock_storage.get_stored_object_md5.return_value = self.file["checksum"]
+        self.client.force_authenticate(user=self.user)
+        resp = self.client.post(reverse("file-upload-url"), self.file, format="json")
+        data = resp.json()
+        assert data["resumable"] is True and data["alreadyUploaded"] is True
+        assert data["uploadURL"] is None
+        assert "file" in data
+        assert data["file"]["id"]
+        mock_storage.create_resumable_upload_session.assert_not_called()
+
+    def test_resumable_falls_back_to_single_put_on_s3(self):
+        # default_storage is S3 in the test env → no resumable support
+        self.client.force_authenticate(user=self.user)
+        resp = self.client.post(reverse("file-upload-url"), self.file, format="json")
+        data = resp.json()
+        assert data["resumable"] is False
+        assert "uploadURL" in data
+        assert "file" in data
+        assert data["file"]["id"]

contentcuration/contentcuration/utils/files.py

@@ -1,5 +1,7 @@
 import base64
+import codecs
 import copy
+import mimetypes
 import os
 import re
 import tempfile
@@ -17,6 +19,12 @@
 from contentcuration.models import File
 from contentcuration.models import generate_object_storage_name
 
+
+# Do this to ensure that we infer mimetypes for files properly, specifically
+# zip file and epub files.
+# to add additional files add them to the mime.types file
+mimetypes.init([os.path.join(os.path.dirname(__file__), "mime.types")])
+
 ImageFile.LOAD_TRUNCATED_IMAGES = True
 THUMBNAIL_WIDTH = 400
 
@@ -196,3 +204,27 @@ def create_thumbnail_from_base64(
             )
     finally:
         os.close(fd)
+
+
+def determine_content_type(filename):
+    """
+    Guesses the content type of a filename. Returns the mimetype of a file.
+
+    Returns "application/octet-stream" if the type can't be guessed.
+    Raises an AssertionError if filename is not a string.
+    """
+
+    typ, _ = mimetypes.guess_type(filename)
+
+    if not typ:
+        return "application/octet-stream"
+    return typ
+
+
+def hex_to_base64(hexdigest):
+    """Convert a hex-encoded digest (e.g. an MD5 checksum) to base64."""
+    return codecs.encode(codecs.decode(hexdigest, "hex"), "base64").decode().strip()
+
+
+def base64_to_hex(b64):
+    return codecs.encode(codecs.decode(b64.encode(), "base64"), "hex").decode().strip()

contentcuration/contentcuration/utils/gcs_storage.py

@@ -11,6 +11,10 @@
 from google.cloud.storage import Client
 from google.cloud.storage.blob import Blob
 
+from .files import determine_content_type
+from contentcuration.utils.files import base64_to_hex
+from contentcuration.utils.files import hex_to_base64
+
 OLD_STUDIO_STORAGE_PREFIX = "/contentworkshop_content/"
 
 CONTENT_DATABASES_MAX_AGE = 5  # seconds
@@ -120,10 +124,6 @@ def save(self, name, fobj, max_length=None, blob_object=None):
             blob.content_encoding = "gzip"
             fobj = buffer
 
-        # determine the current file's mimetype based on the name
-        # import determine_content_type lazily in here, so we don't get into an infinite loop with circular dependencies
-        from contentcuration.utils.storage_common import determine_content_type
-
         content_type = determine_content_type(name)
 
         # force the current file to be at file location 0, to
@@ -215,6 +215,18 @@ def _is_file_empty(fobj):
             fobj.seek(current_location)
         return len(byt) == 0
 
+    def create_resumable_upload_session(self, name, md5, size):
+        blob = self.bucket.blob(name)
+        md5_b64 = hex_to_base64(md5)
+        blob.md5_hash = md5_b64.strip()
+        blob.content_type = determine_content_type(name)
+        blob.metadata = {"declared-size": str(size)}
+        return blob.create_resumable_upload_session(client=self.client)
+
+    def get_stored_object_md5(self, name):
+        blob = self.bucket.get_blob(name)
+        return base64_to_hex(blob.md5_hash) if blob is not None else None
+
 
 class CompositeGCS(Storage):
     def __init__(self):
@@ -283,3 +295,11 @@ def get_created_time(self, name):
 
     def get_modified_time(self, name):
         return self._get_readable_backend(name).get_modified_time(name)
+
+    def create_resumable_upload_session(self, name, md5_b64, size):
+        return self._get_writeable_backend().create_resumable_upload_session(
+            name, md5_b64, size
+        )
+
+    def get_stored_object_md5(self, name):
+        return self._get_readable_backend(name).get_stored_object_md5(name)

contentcuration/contentcuration/utils/storage_common.py

@@ -1,43 +1,22 @@
-import mimetypes
-import os
 from datetime import timedelta
 
 from django.conf import settings
 from django.core.files.storage import default_storage
 from django_s3_storage.storage import S3Storage
 
+from .files import determine_content_type
+from .files import hex_to_base64
 from .gcs_storage import CompositeGCS
 from .gcs_storage import GoogleCloudStorage
 
 
-# Do this to ensure that we infer mimetypes for files properly, specifically
-# zip file and epub files.
-# to add additional files add them to the mime.types file
-mimetypes.init([os.path.join(os.path.dirname(__file__), "mime.types")])
-
-
 class UnknownStorageBackendError(Exception):
     pass
 
 
-def determine_content_type(filename):
-    """
-    Guesses the content type of a filename. Returns the mimetype of a file.
-
-    Returns "application/octet-stream" if the type can't be guessed.
-    Raises an AssertionError if filename is not a string.
-    """
-
-    typ, _ = mimetypes.guess_type(filename)
-
-    if not typ:
-        return "application/octet-stream"
-    return typ
-
-
 def get_presigned_upload_url(
     filepath,
-    md5sum_b64,
+    md5_hex,
     lifetime_sec,
     content_length,
     storage=default_storage,
@@ -48,9 +27,10 @@ def get_presigned_upload_url(
     contents with the contents of your PUT request.
 
     :param: filepath: the file path inside the bucket, to the file.
-    :param: md5sum_b64: the base64 encoded md5 hash of the file. The holder of the URL will
-    have to set a Content-MD5 HTTP header matching this md5sum once it
-    initiates the download.
+    :param: md5_hex: the hex-encoded md5 hash of the file. The base64 encoding
+    that GCS requires for the Content-MD5 header is handled internally; the
+    holder of the URL must set a Content-MD5 HTTP header matching that
+    base64-encoded value once it initiates the upload.
     :param: lifetime_sec: the lifetime of the generated upload url, in seconds.
     :param: content_length: the size of the content, in bytes.
     :param: client: the storage client that will be used to gennerate the presigned URL.
@@ -67,6 +47,7 @@ def get_presigned_upload_url(
     # both storage types are having difficulties enforcing it.
 
     mimetype = determine_content_type(filepath)
+    md5sum_b64 = hex_to_base64(md5_hex)
     if isinstance(storage, (GoogleCloudStorage, CompositeGCS)):
         client = client or storage.get_client()
         bucket = settings.AWS_S3_BUCKET_NAME