diff --git a/rodan-main/code/rodan/jobs/core.py b/rodan-main/code/rodan/jobs/core.py index c972d632d..d58601c70 100644 --- a/rodan-main/code/rodan/jobs/core.py +++ b/rodan-main/code/rodan/jobs/core.py @@ -7,6 +7,8 @@ import sys import tempfile import zipfile +import io, base64 +import six from celery import task, registry from celery import Task @@ -17,7 +19,6 @@ from django.db.models import Q, Case, Value, When, BooleanField import PIL from pybagit.bagit import BagIt -import six import rodan # noqa from rodan.models import ( @@ -43,7 +44,8 @@ from templated_mail.mail import BaseEmailMessage # from rodan.celery import app - +from celery.utils.log import get_task_logger +logger = get_task_logger(__name__) class create_resource(Task): name = "rodan.core.create_resource" @@ -946,26 +948,29 @@ def send_templated_email(to, email_template, context): @task(name="rodan.core.create_archive") def create_archive(resource_uuids): """ - Creates a zip archive of resosurces in memory for a user to download + Creates a zip archive of resources in memory for a user to download """ condition = Q() for uuid in resource_uuids: condition |= Q(uuid=uuid) resources = Resource.objects.filter(Q(resource_file__isnull=False) & condition) + logger.info('Number of resources to archive: %d', resources.count()) + # Don't return an empty zip file if resources.count() == 0: return None - temporary_storage = six.StringIO() - with zipfile.ZipFile(temporary_storage, "a", zipfile.ZIP_DEFLATED) as archive: + + temporary_storage = io.BytesIO() + with zipfile.ZipFile(temporary_storage, "w", zipfile.ZIP_DEFLATED) as archive: for resource in resources: if not resource.resource_file: - print("{} has no file!".format(resource.name)) + logger.warning("{} has no file!".format(resource.name)) continue - # determine a path that doesn't conflict + # Determine a path that doesn't conflict filepath = resource.name + "." + resource.resource_type.extension if filepath in archive.namelist(): - for i in six.moves.range(1, sys.maxint): + for i in range(1, sys.maxsize): filepath = resource.name + " ({}).".format(i) + resource.resource_type.extension if filepath not in archive.namelist(): break @@ -976,7 +981,7 @@ def create_archive(resource_uuids): ) temporary_storage.seek(0) - return temporary_storage + return base64.b64encode(temporary_storage.getvalue()).decode('utf-8') class test_work(Task): diff --git a/rodan-main/code/rodan/views/resource.py b/rodan-main/code/rodan/views/resource.py index ada3cc6eb..affb8ba63 100644 --- a/rodan-main/code/rodan/views/resource.py +++ b/rodan-main/code/rodan/views/resource.py @@ -4,6 +4,7 @@ import re # import urlparse import six.moves.urllib.parse +import base64, io from celery import registry from django.conf import settings @@ -408,11 +409,14 @@ def get(self, request, format=None): archive = registry.tasks['rodan.core.create_archive'] \ .si(resource_uuids).apply_async(queue="celery") - storage = archive.get() - if storage is None: + encoded_archive = archive.get() + + if encoded_archive is None: raise ValidationError({'resource_uuid': ["The specified resources must exist."]}) + + binary_data = base64.b64decode(encoded_archive) response = FileResponse( - storage, + io.BytesIO(binary_data), content_type="application/zip" ) response['Content-Disposition'] = "attachment; filename=Archive.zip"