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)