Extensible file uploaders for emails

explorer/app_settings.py

@@ -99,6 +99,14 @@
 EXPLORER_DATA_EXPORTERS = getattr(
     settings, 'EXPLORER_DATA_EXPORTERS', DEFAULT_EXPORTERS
 )
+DEFAULT_UPLOADERS = [
+    ('FileUploader', 'explorer.uploaders.FileUploader'),
+    ('S3Uploader', 'explorer.uploaders.S3Uploader'),
+]
+EXPLORER_DATA_UPLOADERS = getattr(
+    settings, 'EXPLORER_DATA_UPLOADERS', DEFAULT_UPLOADERS
+)
+EXPLORER_DEFAULT_UPLOADER = 'FileUploader'
 CSV_DELIMETER = getattr(settings, "EXPLORER_CSV_DELIMETER", ",")
 
 # API access
@@ -123,6 +131,9 @@
 # Async task related. Note that the EMAIL_HOST settings must be set up for
 # email to work.
 ENABLE_TASKS = getattr(settings, "EXPLORER_TASKS_ENABLED", False)
+SNAPSHOTS_FILE_PREFIX = getattr(settings, "EXPLORER_SNAPSHOTS_FILE_PREFIX", "report-")
+SNAPSHOTS_UPLOAD_DIR = getattr(settings, "EXPLORER_SNAPSHOTS_UPLOAD_DIR", None)
+SNAPSHOTS_DOWNLOAD_DIR = getattr(settings, "EXPLORER_SNAPSHOTS_DOWNLOAD_DIR", None)
 S3_ACCESS_KEY = getattr(settings, "EXPLORER_S3_ACCESS_KEY", None)
 S3_SECRET_KEY = getattr(settings, "EXPLORER_S3_SECRET_KEY", None)
 S3_BUCKET = getattr(settings, "EXPLORER_S3_BUCKET", None)

explorer/models.py

@@ -9,7 +9,7 @@
 
 from explorer import app_settings
 from explorer.utils import (
-    extract_params, get_params_for_url, get_s3_bucket, get_valid_connection, passes_blacklist, s3_url,
+    extract_params, get_params_for_url, get_valid_connection, passes_blacklist,
     shared_dict_update, swap_params,
 )
 
@@ -162,12 +162,14 @@ def shared(self):
     @property
     def snapshots(self):
         if app_settings.ENABLE_TASKS:
-            b = get_s3_bucket()
+            from explorer.uploaders import S3Uploader
+            s3 = S3Uploader()
+            b = s3.get_s3_bucket()
             objects = b.objects.filter(Prefix=f'query-{self.id}/snap-')
             objects_s = sorted(objects, key=lambda k: k.last_modified)
             return [
                 SnapShot(
-                    s3_url(b, o.key),
+                    s3.s3_url(b, o.key),
                     o.last_modified
                 ) for o in objects_s
             ]

explorer/tasks.py

@@ -1,4 +1,3 @@
-import io
 import random
 import string
 from datetime import date, datetime, timedelta
@@ -9,13 +8,12 @@
 from explorer import app_settings
 from explorer.exporters import get_exporter_class
 from explorer.models import Query, QueryLog
+from explorer.uploaders import get_uploader_class
 
 
 if app_settings.ENABLE_TASKS:
     from celery import shared_task
     from celery.utils.log import get_task_logger
-
-    from explorer.utils import s3_upload
     logger = get_task_logger(__name__)
 else:
     import logging
@@ -33,13 +31,15 @@ def execute_query(query_id, email_address):
               [email_address])
 
     exporter = get_exporter_class('csv')(q)
+    uploader = get_uploader_class(app_settings.EXPLORER_DEFAULT_UPLOADER)
     random_part = ''.join(
         random.choice(
             string.ascii_uppercase + string.digits
         ) for _ in range(20)
     )
+    filename = f"{app_settings.SNAPSHOTS_FILE_PREFIX}{date.today().strftime('%Y%m%d%H%M%S')}-{random_part}"
     try:
-        url = s3_upload(f'{random_part}.csv', convert_csv_to_bytesio(exporter))
+        url = uploader().upload(filename, exporter)
         subj = f'[SQL Explorer] Report "{q.title}" is ready'
         msg = f'Download results:\n\r{url}'
     except Exception as e:
@@ -49,27 +49,19 @@ def execute_query(query_id, email_address):
     send_mail(subj, msg, app_settings.FROM_EMAIL, [email_address])
 
 
-# I am sure there is a much more efficient way to do this but boto3 expects a binary file basically
-def convert_csv_to_bytesio(csv_exporter):
-    csv_file_io = csv_exporter.get_file_output()
-    csv_file_io.seek(0)
-    csv_data: str = csv_file_io.read()
-    bio = io.BytesIO(bytes(csv_data, 'utf-8'))
-    return bio
-
-
 @shared_task
 def snapshot_query(query_id):
     try:
         logger.info(f"Starting snapshot for query {query_id}...")
         q = Query.objects.get(pk=query_id)
         exporter = get_exporter_class('csv')(q)
-        k = 'query-{}/snap-{}.csv'.format(
+        uploader = get_uploader_class(app_settings.EXPLORER_DEFAULT_UPLOADER)
+        filename = 'query-{}/snap-{}.csv'.format(
             q.id,
             date.today().strftime('%Y%m%d-%H:%M:%S')
         )
         logger.info(f"Uploading snapshot for query {query_id} as {k}...")
-        url = s3_upload(k, convert_csv_to_bytesio(exporter))
+        url = uploader().upload(filename, exporter)
         logger.info(
             f"Done uploading snapshot for query {query_id}. URL: {url}"
         )

explorer/uploaders.py

@@ -0,0 +1,58 @@
+import io
+
+from django.utils.module_loading import import_string
+
+from explorer import app_settings
+
+
+def get_uploader_class(uploader):
+    class_str = dict(getattr(app_settings, 'EXPLORER_DATA_UPLOADERS'))[uploader]
+    return import_string(class_str)
+
+class BaseUploader:
+
+    def prepare_file(self, data):
+        csv_file_io = data.get_file_output()
+        csv_file_io.seek(0)
+        csv_data = csv_file_io.read()
+        bio = io.BytesIO(bytes(csv_data, 'utf-8'))
+        return bio
+
+    def upload(self):
+        raise NotImplementedError
+
+class FileUploader(BaseUploader):
+    def upload(self, key, data):
+        from zipfile import ZipFile
+        import os
+        import shutil
+        full_filename = os.path.join(app_settings.SNAPSHOTS_UPLOAD_DIR, f"{key}.zip")
+        download_url = os.path.join(app_settings.SNAPSHOTS_DOWNLOAD_DIR, f"{key}.zip")
+        with ZipFile(full_filename, 'w') as zfh:
+            zfh.writestr(f"{key}.csv", self.prepare_file(data).getvalue())
+        return download_url
+
+class S3Uploader(BaseUploader):
+    def get_s3_bucket(self):
+        import boto3
+        kwargs = {
+            'aws_access_key_id': app_settings.S3_ACCESS_KEY,
+            'aws_secret_access_key': app_settings.S3_SECRET_KEY,
+            'region_name': app_settings.S3_REGION
+        }
+        if app_settings.S3_ENDPOINT_URL:
+            kwargs['endpoint_url'] = app_settings.S3_ENDPOINT_URL
+        s3 = boto3.resource('s3', **kwargs)
+        return s3.Bucket(name=app_settings.S3_BUCKET)
+
+    def s3_url(self, bucket, key):
+        url = bucket.meta.client.generate_presigned_url(
+            ClientMethod='get_object',
+            Params={'Bucket': app_settings.S3_BUCKET, 'Key': key},
+            ExpiresIn=app_settings.S3_LINK_EXPIRATION)
+        return url
+
+    def upload(self, key, data):
+        bucket = self.get_s3_bucket()
+        bucket.upload_fileobj(self.prepare_file(data), key, ExtraArgs={'ContentType': "text/csv"})
+        return self.s3_url(bucket, key)

explorer/utils.py

@@ -191,33 +191,6 @@ def get_valid_connection(alias=None):
     return connections[alias]
 
 
-def get_s3_bucket():
-    import boto3
-    kwargs = {
-        'aws_access_key_id': app_settings.S3_ACCESS_KEY,
-        'aws_secret_access_key': app_settings.S3_SECRET_KEY,
-        'region_name': app_settings.S3_REGION
-    }
-    if app_settings.S3_ENDPOINT_URL:
-        kwargs['endpoint_url'] = app_settings.S3_ENDPOINT_URL
-    s3 = boto3.resource('s3', **kwargs)
-    return s3.Bucket(name=app_settings.S3_BUCKET)
-
-
-def s3_upload(key, data):
-    bucket = get_s3_bucket()
-    bucket.upload_fileobj(data, key, ExtraArgs={'ContentType': "text/csv"})
-    return s3_url(bucket, key)
-
-
-def s3_url(bucket, key):
-    url = bucket.meta.client.generate_presigned_url(
-        ClientMethod='get_object',
-        Params={'Bucket': app_settings.S3_BUCKET, 'Key': key},
-        ExpiresIn=app_settings.S3_LINK_EXPIRATION)
-    return url
-
-
 def is_xls_writer_available():
     try:
         import xlsxwriter  # noqa