Skip to content

Commit

Permalink
fix: AWS file upload (#6707)
Browse files Browse the repository at this point in the history
* add test that reproduces the issue (fail with binary files)

* add len to UploadedFile

* aws upload: do not load to memory
  • Loading branch information
AmitAronovitch authored and iamareebjamal committed Jan 1, 2020
1 parent cf88e99 commit 4b55ca6
Show file tree
Hide file tree
Showing 2 changed files with 39 additions and 8 deletions.
23 changes: 16 additions & 7 deletions app/api/helpers/storage.py
Original file line number Diff line number Diff line change
Expand Up @@ -99,7 +99,16 @@ class UploadedFile:
def __init__(self, file_path, filename):
self.file_path = file_path
self.filename = filename
self.file = open(file_path)
self.file = open(file_path, 'rb')

def __len__(self):
position = self.file.tell()
try:
self.file.seek(0, os.SEEK_END)
last_position = self.file.tell()
finally:
self.file.seek(position)
return last_position

def save(self, new_path):
copyfile(self.file_path, new_path)
Expand Down Expand Up @@ -224,15 +233,15 @@ def upload_to_aws(bucket_name, aws_region, aws_key, aws_secret, file, key, acl='
item.delete()
# set object settings

file_data = file.read()
file_mime = magic.from_buffer(file_data, mime=True)
size = len(file_data)
sent = k.set_contents_from_string(
file_data,
file_mime = magic.from_file(file.file_path, mime=True)
size = len(file)
sent = k.set_contents_from_file(
file.file,
headers={
'Content-Disposition': 'attachment; filename=%s' % filename,
'Content-Type': '%s' % file_mime
}
},
rewind=True
)
k.set_acl(acl)
s3_url = 'https://%s.s3.amazonaws.com/' % bucket_name
Expand Down
24 changes: 23 additions & 1 deletion tests/all/unit/api/helpers/test_storage.py
Original file line number Diff line number Diff line change
@@ -1,8 +1,10 @@
"""Test file for storage functions."""
import os
from tempfile import TemporaryDirectory
import unittest
from unittest.mock import patch

from app.api.helpers.storage import create_url, generate_hash
from app.api.helpers.storage import create_url, generate_hash, UploadedFile


class TestStorageHelperValidation(unittest.TestCase):
Expand Down Expand Up @@ -65,5 +67,25 @@ def patch_settings(settings):
self.assertEqual(len(actual_output), 10)


class TestUploadedFile(unittest.TestCase):
test_data = b'\xffhellothere'

def _uploaded_file(self, folder, filename='testfile.bin'):
path = os.path.join(folder, filename)
with open(path, 'wb') as f:
f.write(self.test_data)
return UploadedFile(path, filename)

def test_len(self):
with TemporaryDirectory() as folder:
file = self._uploaded_file(folder)
self.assertEqual(len(file), len(self.test_data))

def test_read(self):
with TemporaryDirectory() as folder:
file = self._uploaded_file(folder)
self.assertEqual(file.read(), self.test_data)


if __name__ == '__main__':
unittest.main()

0 comments on commit 4b55ca6

Please sign in to comment.