import contextlib import tempfile import boto3 import botocore.exceptions from flask import request, send_file from .base import BaseStorage class Storage(BaseStorage): def __init__(self, app): super().__init__(app) self.client = boto3.resource('s3', **self.app.config['S3_CONFIG']) def _get_object_key(self, file_hash, user_id): return '{}_{}'.format(file_hash, user_id) 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']) key = self._get_object_key(file_hash, user.id if user else 0) obj = bucket.upload_fileobj(Fileobj=uploaded_file.stream, Key=key) size = uploaded_file.content_length if not size: obj = self.client.ObjectSummary(self.app.config['S3_BUCKET'], key) size = obj.size try: return self.add_file(file_hash, uploaded_file.filename, size, user, ip) except: obj.delete() raise def get_file(self, f, thumb=False): 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'] data = obj.get(**kwargs) rv = send_file(data['Body'], attachment_filename=f.filename) rv.headers['Content-Length'] = data['ContentLength'] rv.headers['Accept-Ranges'] = data['AcceptRanges'] if 'ContentRange' in data: rv.headers['Content-Range'] = data['ContentRange'] rv.status_code = 206 return rv 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): obj = self.client.Object(self.app.config['S3_BUCKET'], self.get_object_key(f)) with tempfile.NamedTemporaryFile() as f: obj.download_fileobj(f) 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)