diff options
Diffstat (limited to 'fbin/file_storage/s3.py')
| -rw-r--r-- | fbin/file_storage/s3.py | 32 |
1 files changed, 28 insertions, 4 deletions
diff --git a/fbin/file_storage/s3.py b/fbin/file_storage/s3.py index 2f0b87b..e81f8b4 100644 --- a/fbin/file_storage/s3.py +++ b/fbin/file_storage/s3.py @@ -2,6 +2,7 @@ import contextlib import tempfile import boto3 +import botocore.exceptions from flask import request, send_file from .base import BaseStorage @@ -14,8 +15,11 @@ class Storage(BaseStorage): def _get_object_key(self, file_hash, user_id): return '{}_{}'.format(file_hash, user_id) - def get_object_key(self, f): - return self._get_object_key(f.hash, f.user_id if f.user_id else 0) + def get_object_key(self, f, thumb=False): + key = self._get_object_key(f.hash, f.user_id if f.user_id else 0) + if thumb: + key += '_thumb' + return key def store_file(self, uploaded_file, file_hash, user, ip): bucket = self.client.Bucket(self.app.config['S3_BUCKET']) @@ -27,8 +31,13 @@ class Storage(BaseStorage): size = obj.size return self.add_file(file_hash, uploaded_file.filename, size, user, ip) - def get_file(self, f): - obj = self.client.Object(self.app.config['S3_BUCKET'], self.get_object_key(f)) + def get_file(self, f, thumb=True): + key = self.get_object_key(f, thumb=thumb) + if thumb: + bucket = self.app.config['S3_THUMB_BUCKET']) + else: + bucket = self.app.config['S3_BUCKET'] + obj = self.client.Object(bucket, key) kwargs = {} if 'Range' in request.headers: kwargs['Range'] = request.headers['Range'] @@ -44,6 +53,8 @@ class Storage(BaseStorage): def delete_file(self, f): obj = self.client.Object(self.app.config['S3_BUCKET'], self.get_object_key(f)) obj.delete() + obj = self.client.Object(self.app.config['S3_BUCKET'], self.get_object_key(f, thumb=True)) + obj.delete() @contextlib.contextmanager def temp_file(self, f): @@ -53,3 +64,16 @@ class Storage(BaseStorage): f.seek(0) yield f + def get_thumbnail(self, f): + try: + return self.get_file(f, thumb=True) + except botocore.exceptions.ClientError as e: + if e.response['Error']['Code'] == 'NoSuchKey': + # If thumbnail does not exist, just return None. + return + raise + + def store_thumbnail(self, f, stream): + bucket = self.client.Bucket(self.app.config['S3_THUMB_BUCKET'])) + key = self.get_object_key(f, thumb=True) + obj = bucket.upload_fileobj(Fileobj=stream, Key=key) |