Merge pull request #26 from 4Catalyzer/upload_info

Add get_upload_info

flask_annex/base.py

@@ -25,3 +25,6 @@ def save_file(self, key, in_file):
 
     def send_file(self, key):
         raise NotImplementedError()
+
+    def get_upload_info(self, key):
+        raise NotImplementedError()

flask_annex/file.py

@@ -91,3 +91,6 @@ def send_file(self, key):
             as_attachment=True,
             attachment_filename=os.path.basename(key),
         )
+
+    def get_upload_info(self, key):
+        raise NotImplementedError("file annex does not support upload info")

flask_annex/s3.py

@@ -82,5 +82,38 @@ def send_file(self, key):
                 'Bucket': self._bucket_name,
                 'Key': key,
             },
+            ExpiresIn=self._url_expires_in,
         )
         return flask.redirect(url)
+
+    def get_upload_info(self, key):
+        fields = {}
+        conditions = []
+
+        content_type = mimetypes.guess_type(key)[0]
+        if content_type:
+            fields['Content-Type'] = content_type
+
+        max_content_length = flask.current_app.config['MAX_CONTENT_LENGTH']
+        if max_content_length is not None:
+            conditions.append(
+                ('content-length-range', 0, max_content_length),
+            )
+
+        # Boto doesn't automatically add fields to conditions.
+        for field_key, field_value in fields.items():
+            conditions.append({field_key: field_value})
+
+        post_info = self._client.generate_presigned_post(
+            Bucket=self._bucket_name,
+            Key=key,
+            Fields=fields,
+            Conditions=conditions,
+            ExpiresIn=self._url_expires_in,
+        )
+
+        return {
+            'method': 'POST',
+            'url': post_info['url'],
+            'data': post_info['fields'],
+        }

tests/test_file.py

@@ -33,6 +33,10 @@ def test_send_file(self, app, annex):
         assert response.status_code == 200
         assert response.mimetype == 'application/json'
 
+    def test_get_upload_info(self, annex):
+        with pytest.raises(NotImplementedError):
+            annex.get_upload_info('foo/qux.txt')
+
 
 class TestFileAnnexFromEnv(TestFileAnnex):
     @pytest.fixture

tests/test_s3.py

@@ -1,4 +1,6 @@
+import base64
 from io import BytesIO
+import json
 
 import pytest
 
@@ -26,6 +28,26 @@ def bucket_name():
         yield bucket.name
 
 
+def get_policy(upload_info):
+    return json.loads(
+        base64.urlsafe_b64decode(
+            upload_info['data']['policy'].encode(),
+        ).decode(),
+    )
+
+
+def get_condition(conditions, key):
+    for condition in conditions:
+        if isinstance(condition, list):
+            if condition[0] == key:
+                return condition[1:]
+        else:
+            if key in condition:
+                return condition[key]
+
+    raise KeyError()
+
+
 # -----------------------------------------------------------------------------
 
 
@@ -44,6 +66,28 @@ def test_send_file(self, app, annex):
 
         assert response.status_code == 302
 
+    def test_get_upload_info(self, app, annex):
+        with app.app_context():
+            upload_info = annex.get_upload_info('foo/qux.txt')
+            assert upload_info['method'] == 'POST'
+            assert upload_info['url'] == \
+                'https://flask-annex.s3.amazonaws.com/'
+            assert upload_info['data']['key'] == 'foo/qux.txt'
+            assert upload_info['data']['Content-Type'] == 'text/plain'
+
+            conditions = get_policy(upload_info)['conditions']
+            assert get_condition(conditions, 'bucket') == 'flask-annex'
+            assert get_condition(conditions, 'key') == 'foo/qux.txt'
+            assert get_condition(conditions, 'Content-Type') == 'text/plain'
+
+    def test_get_upload_info_max_content_length(self, app, annex):
+        app.config['MAX_CONTENT_LENGTH'] = 100
+
+        with app.app_context():
+            upload_info = annex.get_upload_info('foo/qux.txt')
+            conditions = get_policy(upload_info)['conditions']
+            assert get_condition(conditions, 'content-length-range')[1] == 100
+
 
 class TestS3AnnexFromEnv(TestS3Annex):
     @pytest.fixture