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.

3 files changed, 52 insertions, 4 deletions

diff --git a/fbin/file_storage/base.py b/fbin/file_storage/base.py
index 6f39665..9f09199 100644
--- a/fbin/file_storage/base.py
+++ b/fbin/file_storage/base.py
@@ -37,3 +37,12 @@ class BaseStorage:
         This is used internally for eg. thumbnails.'''
         raise NotImplementedError()
 
+    def get_thumbnail(self, f):
+        '''Return a file object for the specified file's thumbnail.
+
+        Subclasses can also return a flask.Response instance if required.'''
+        raise NotImplementedError()
+
+    def store_thumbnail(self, f, stream):
+        '''Store thumbnail for the specified file.'''
+        raise NotImplementedError()
diff --git a/fbin/file_storage/filesystem.py b/fbin/file_storage/filesystem.py
index 3b46e34..1259002 100644
--- a/fbin/file_storage/filesystem.py
+++ b/fbin/file_storage/filesystem.py
@@ -8,6 +8,7 @@ class Storage(BaseStorage):
     def __init__(self, app):
         super().__init__(app)
         os.makedirs(self.app.config['FILE_DIRECTORY'], exist_ok=True)
+        os.makedirs(self.app.config['THUMB_DIRECTORY'], exist_ok=True)
 
     def store_file(self, uploaded_file, file_hash, user, ip):
         size = uploaded_file.content_length
@@ -45,3 +46,17 @@ class Storage(BaseStorage):
     def temp_file(self, f):
         with open(f.get_path(), 'rb') as f:
             yield f
+
+    def get_thumbnail(self, f):
+        path = f.get_thumb_path()
+        if not os.path.exists(path):
+            return
+        return path
+
+    def store_thumbnail(self, f, stream):
+        path = f.get_thumb_path()
+        with open(path, 'wb') as f:
+            buf = stream.read(1024*10)
+            while buf:
+                f.write(buf)
+                buf = stream.read(1024*10)
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)