Fetch and store thumbnails via storage modules

This will allow us to remotely store thumbnails in case of S3. For S3
the thumb bucket is configurable to allow these to be stored separately.
The S3 key for thumbnails does not conflict with files, so these can be
stored in the same bucket if needed.

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)