Consolidate presigned URL handling into storage arch

Author: bjester

contentcuration/contentcuration/tests/utils/test_storage.py

@@ -7,16 +7,15 @@
 import requests
 from django.core.files.storage import FileSystemStorage
 from django.test import TestCase
-from django_s3_storage.storage import S3Storage
 from mock import MagicMock
 
 from ..base import StudioTestCase
 from contentcuration.models import generate_object_storage_name
-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
 from contentcuration.utils.storage.common import UnknownStorageBackendError
-# The modules we'll test
+from contentcuration.utils.storage.dev import Storage as DevStorage
+from contentcuration.utils.storage.gcs import GoogleCloudStorage
 
 
 class MimeTypesTestCase(TestCase):
@@ -77,7 +76,10 @@ def test_raises_error(self):
         """
         with pytest.raises(UnknownStorageBackendError):
             get_presigned_upload_url(
-                "nice", "err", 5, 0, storage=self.STORAGE,
+                "nice",
+                "err",
+                5,
+                storage=self.STORAGE,
             )
 
 
@@ -90,7 +92,9 @@ class GoogleCloudStoragePresignedURLUnitTestCase(TestCase):
     """
 
     def setUp(self):
+        super().setUp()
         self.client = MagicMock()
+        self.storage = GoogleCloudStorage(self.client, "fake")
         self.generate_signed_url_method = (
             self.client.get_bucket.return_value.blob.return_value.generate_signed_url
         )
@@ -102,19 +106,15 @@ def test_that_generate_signed_url_is_called(self):
         """
         Check that we even call blob.generate_signed_url in the first place.
         """
-        bucket = "fake"
-        _get_gcs_presigned_put_url(self.client, bucket, "/object.jpg", "aBc", 0, 0)
+        get_presigned_upload_url("/object.jpg", "aBc", 0, storage=self.storage)
         self.generate_signed_url_method.assert_called_once()
 
     def test_that_we_return_a_string(self):
         """
         Check that _get_gcs_presigned_put_url returns a string.
         """
-        bucket = "fake"
-        ret = _get_gcs_presigned_put_url(
-            self.client, bucket, "/object.jpg", "aBc", 0, 0
-        )
-        assert isinstance(ret, str)
+        ret = get_presigned_upload_url("/object.jpg", "aBc", 0, storage=self.storage)
+        assert isinstance(ret["uploadURL"], str)
 
     def test_generate_signed_url_called_with_required_arguments(self):
         """
@@ -132,11 +132,9 @@ def test_generate_signed_url_called_with_required_arguments(self):
         bucket_name = "fake"
         filepath = "object.jpg"
         lifetime = 20  # seconds
-        mimetype = "doesntmatter"
+        mimetype = "image/jpeg"
 
-        _get_gcs_presigned_put_url(
-            self.client, bucket_name, filepath, content_md5, lifetime, mimetype
-        )
+        get_presigned_upload_url(filepath, content_md5, lifetime, storage=self.storage)
 
         # assert that we're creating the right object
         self.client.get_bucket.assert_called_once_with(bucket_name)
@@ -148,8 +146,8 @@ def test_generate_signed_url_called_with_required_arguments(self):
         self.generate_signed_url_method.assert_called_once_with(
             method=method,
             content_md5=content_md5,
-            expiration=lifetime_timedelta,
             content_type=mimetype,
+            expiration=lifetime_timedelta,
         )
 
 
@@ -158,11 +156,9 @@ class S3StoragePresignedURLUnitTestCase(StudioTestCase):
     Test cases for generating presigned URLs for S3 storage, i.e. Minio.
     """
 
-    STORAGE = S3Storage()
-
     def setUp(self):
-        self.client = MagicMock()
         super().setUp()
+        self.storage = DevStorage()
 
     def test_returns_string_if_inputs_are_valid(self):
         """
@@ -171,9 +167,7 @@ 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
-        )
+        ret = get_presigned_upload_url("a/b/abc.jpg", "aBc", 10, storage=self.storage)
         url = ret["uploadURL"]
 
         assert isinstance(url, str)
@@ -187,12 +181,14 @@ def test_can_upload_file_to_presigned_url(self):
         # 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_base64 = codecs.encode(
+            codecs.decode(md5_checksum, "hex"), "base64"
+        ).decode()
 
         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_base64, 1000)
         url = ret["uploadURL"]
         content_type = ret["mimetype"]
 
@@ -201,6 +197,6 @@ def test_can_upload_file_to_presigned_url(self):
             data=file,
             headers={
                 "Content-Type": content_type,
-            }
+            },
         )
         resp.raise_for_status()

contentcuration/contentcuration/utils/storage/base.py

@@ -14,6 +14,20 @@ def get_client(self):
         """
         return None
 
+    def get_presigned_put_url(
+        self, filepath, md5sum, lifetime_sec, mimetype="application/octet-stream"
+    ):
+        """
+        Creates a pre-signed URL for uploading files.
+
+        :param filepath: A string representing the destination file path inside the bucket
+        :param md5sum: A MD5 checksum of the file to be uploaded
+        :param lifetime_sec: The lifetime of the URL in seconds
+        :param mimetype: The content type of the file to be uploaded
+        :return: A pre-signed URL for uploading the file
+        """
+        raise NotImplementedError("Subclasses must implement this method")
+
 
 class CompositeStorage(Storage):
     def __init__(self):
@@ -40,7 +54,7 @@ def _get_readable_backend(self, name):
     def get_client(self):
         return self._get_writeable_backend().get_client()
 
-    def open(self, name, mode='rb'):
+    def open(self, name, mode="rb"):
         return self._get_readable_backend(name).open(name, mode)
 
     def save(self, name, content, max_length=None):
@@ -74,3 +88,10 @@ def get_created_time(self, name):
 
     def get_modified_time(self, name):
         return self._get_readable_backend(name).get_modified_time(name)
+
+    def get_presigned_put_url(
+        self, filepath, md5sum, lifetime_sec, mimetype="application/octet-stream"
+    ):
+        return self._get_writeable_backend().get_presigned_put_url(
+            filepath, md5sum, lifetime_sec, mimetype=mimetype
+        )

contentcuration/contentcuration/utils/storage/common.py

@@ -1,15 +1,10 @@
 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 .base import CompositeStorage
 from .base import Storage
-from .gcs import CompositeGCS
-from .gcs import GoogleCloudStorage
 
 
 # Do this to ensure that we infer mimetypes for files properly, specifically
@@ -38,9 +33,10 @@ def determine_content_type(filename):
 
 
 def get_presigned_upload_url(
-    filepath, md5sum_b64, lifetime_sec, content_length, storage=default_storage, client=None
+    filepath, md5sum_b64, lifetime_sec, storage=default_storage
 ):
-    """Return a presigned URL that can modify the given filepath through a PUT
+    """
+    Return a presigned URL that can modify the given filepath through a PUT
     request. Performing a PUT request on the returned URL changes the object's
     contents with the contents of your PUT request.
 
@@ -49,92 +45,22 @@ def get_presigned_upload_url(
     have to set a Content-MD5 HTTP header matching this md5sum once it
     initiates the download.
     :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.
-    This must have an API that's similar to either the GCS client or the boto3 client.
 
     :returns: a dictionary containing 2 keys:
         mimetype: the mimetype that will be required to send as part of the file upload's mimetype header
         uploadURL: the URL to upload the file to.
 
     :raises: :class:`UnknownStorageBackendError`: If the storage backend is not S3 or GCS.
     """
-
-    # Aron: note that content_length is not used right now because
-    # both storage types are having difficulties enforcing it.
-
     mimetype = determine_content_type(filepath)
-    bucket = settings.AWS_S3_BUCKET_NAME
-
-    if isinstance(storage, Storage):
-        client = client or storage.get_client()
 
-    if isinstance(storage, (GoogleCloudStorage, CompositeGCS)):
-        upload_url = _get_gcs_presigned_put_url(client, bucket, filepath, md5sum_b64, lifetime_sec, mimetype=mimetype)
-    elif isinstance(storage, (S3Storage, CompositeStorage)):
-        upload_url = _get_s3_presigned_put_url(client, bucket, filepath, md5sum_b64, lifetime_sec)
+    if isinstance(storage, (Storage, CompositeStorage)):
+        upload_url = storage.get_presigned_put_url(
+            filepath, md5sum_b64, lifetime_sec, mimetype=mimetype
+        )
     else:
         raise UnknownStorageBackendError(
             "Please ensure your storage backend is either Google Cloud Storage or S3 Storage!"
         )
 
-    return {
-        "mimetype": mimetype,
-        "uploadURL": upload_url
-    }
-
-
-def _get_gcs_presigned_put_url(gcs_client, bucket, filepath, md5sum, lifetime_sec, mimetype="application/octet-stream"):
-    bucket_obj = gcs_client.get_bucket(bucket)
-    blob_obj = bucket_obj.blob(filepath)
-
-    # ensure the md5sum doesn't have any whitespace, including newlines.
-    # We should do the same whitespace stripping as well on any client that actually
-    # uses the returned presigned url.
-    md5sum_stripped = md5sum.strip()
-
-    # convert the lifetime to a timedelta, so gcloud library will interpret the lifetime
-    # as the seconds from right now. If we use an absolute integer, it's the number of seconds
-    # from unix time
-    lifetime_timedelta = timedelta(seconds=lifetime_sec)
-
-    url = blob_obj.generate_signed_url(
-        method="PUT",
-        content_md5=md5sum_stripped,
-        content_type=mimetype,
-        expiration=lifetime_timedelta,
-    )
-
-    return url
-
-
-def _get_s3_presigned_put_url(s3_client, bucket, filepath, md5sum, lifetime_sec):
-    """
-    Creates a pre-signed URL for S3-like backends, e.g. Minio.
-
-    Note that since our production object storage backend is GCS, we do not enforce or require
-    any Content-MD5 value.
-
-    :param: s3_client: an initialized S3 client. We will use this to create the presigned PUT url.
-    :param: bucket: the bucket where the user can PUT their object.
-    :param: filepath: the file path inside the bucket that the user can PUT their object.
-    :param: md5sum: the base64-encoded MD5sum of the object the user is planning to PUT.
-        This is ignored for this function and added solely to maintain API compatibility with other
-        private presigned URL functions.
-    :param: lifetime_sec: how long before the presigned URL expires, in seconds.
-    """
-    # S3's PUT Object parameters:
-    # https://docs.aws.amazon.com/AmazonS3/latest/API/API_PutObject.html
-    method = "put_object"
-    fields = {
-        "Bucket": bucket,
-        "Key": filepath,
-    }
-
-    response = s3_client.generate_presigned_url(
-        ClientMethod=method,
-        Params=fields,
-        ExpiresIn=lifetime_sec,
-    )
-
-    return response
+    return {"mimetype": mimetype, "uploadURL": upload_url}

contentcuration/contentcuration/utils/storage/dev.py

@@ -1,3 +1,4 @@
+from django.conf import settings
 from django_s3_storage.storage import S3Storage
 from google.cloud.storage import Client
 
@@ -9,13 +10,44 @@
 class Storage(S3Storage, BaseStorage):
     def get_client(self):
         """
-        :rtype: object
+        :rtype: botocore.client.BaseClient
         """
         return self.s3_connection
 
+    def get_presigned_put_url(self, filepath, md5sum, lifetime_sec, mimetype=None):
+        """
+        Creates a pre-signed URL for development storage backends
+
+        Note that since our production object storage backend is GCS, we do not enforce or require
+        any Content-MD5 value.
+
+        :param: filepath: the file path inside the bucket that the user can PUT their object.
+        :param: md5sum: the base64-encoded MD5sum of the object the user is planning to PUT.
+            This is ignored for this function and added solely to maintain API compatibility with other
+            private presigned URL functions.
+        :param: lifetime_sec: how long before the presigned URL expires, in seconds.
+        :param: mimetype: the content type of the file to be uploaded
+        :return: A pre-signed URL for uploading the file
+        """
+        # S3's PUT Object parameters:
+        # https://docs.aws.amazon.com/AmazonS3/latest/API/API_PutObject.html
+        method = "put_object"
+        fields = {
+            "Bucket": settings.AWS_S3_BUCKET_NAME,
+            "Key": filepath,
+        }
+
+        return self.get_client().generate_presigned_url(
+            ClientMethod=method,
+            Params=fields,
+            ExpiresIn=lifetime_sec,
+        )
+
 
 class CompositeStorage(BaseCompositeStorage):
     def __init__(self):
         super(CompositeStorage, self).__init__()
         self.backends.append(Storage())
-        self.backends.append(GoogleCloudStorage(Client.create_anonymous_client(), "studio-content"))
+        self.backends.append(
+            GoogleCloudStorage(Client.create_anonymous_client(), "studio-content")
+        )