diff --git a/apps/addons/management/commands/process_addons.py b/apps/addons/management/commands/process_addons.py new file mode 100644 index 0000000000..89c15fef5b --- /dev/null +++ b/apps/addons/management/commands/process_addons.py @@ -0,0 +1,35 @@ +from optparse import make_option + +from django.core.management.base import BaseCommand + +from addons.models import Addon +from amo.utils import chunked +from devhub.tasks import flag_binary + +from celery.messaging import establish_connection + +tasks = { + 'flag_binary': flag_binary, +} + + +class Command(BaseCommand): + """ + A generic command to run a task on *all* addons. + Add tasks to the reg dictionary. + """ + help = 'Corrects the binary status of an addon.' + + option_list = BaseCommand.option_list + ( + make_option('--task', action='store', type='string', + dest='task', help='Run task on all addons.'), + ) + + def handle(self, *args, **options): + pks = Addon.objects.values_list('pk', flat=True).order_by('id') + task = tasks.get(options.get('task')) + if not task: + raise ValueError('Unknown task: %s' % ', '.join(tasks.keys())) + with establish_connection(): + for chunk in chunked(pks, 100): + task.delay(chunk) diff --git a/apps/devhub/tasks.py b/apps/devhub/tasks.py index 4eeb9e3635..ca1d0b9451 100644 --- a/apps/devhub/tasks.py +++ b/apps/devhub/tasks.py @@ -1,3 +1,4 @@ +import json import logging import os import sys @@ -7,6 +8,7 @@ from django.conf import settings from django.core.management import call_command from celeryutils import task +from addons.models import Addon import amo from amo.decorators import write, set_modified_on from amo.utils import resize_image @@ -90,6 +92,28 @@ def run_validator(file_path, for_appversions=None): spidermonkey=settings.SPIDERMONKEY) +@task(rate_limit='1/m') +@write +def flag_binary(ids, **kw): + log.info('[%s@%s] Flagging binary addons starting with id: %s...' + % (len(ids), flag_binary.rate_limit, ids[0])) + addons = Addon.objects.filter(pk__in=ids).no_transforms() + + for addon in addons: + try: + log.info('Validating addon with id: %s' % addon.pk) + file = File.objects.filter(version__addon=addon).latest('created') + result = run_validator(file.file_path) + binary = (json.loads(result)['metadata'] + .get('contains_binary_extension', False)) + log.info('Setting binary for addon with id: %s to %s' + % (addon.pk, binary)) + addon.update(binary=binary) + except Exception, err: + log.error('Failed to run validation on addon id: %s, %s' + % (addon.pk, err)) + + @task @set_modified_on def resize_icon(src, dst, size, **kw): diff --git a/apps/devhub/tests/test_tasks.py b/apps/devhub/tests/test_tasks.py index d58d3976cf..7ee895dd57 100644 --- a/apps/devhub/tests/test_tasks.py +++ b/apps/devhub/tests/test_tasks.py @@ -10,8 +10,9 @@ import test_utils from nose.tools import eq_ from PIL import Image +from addons.models import Addon from files.models import FileUpload -from devhub.tasks import resize_icon, validator +from devhub.tasks import flag_binary, resize_icon, validator def test_resize_icon_shrink(): @@ -112,3 +113,29 @@ class TestValidator(test_utils.TestCase): validator(self.upload.pk) error = self.get_upload().task_error assert error.startswith('Traceback (most recent call last)'), error + + +class TestFlagBinary(test_utils.TestCase): + fixtures = ['base/addon_3615'] + + def setUp(self): + self.addon = Addon.objects.get(pk=3615) + self.addon.update(binary=False) + + @mock.patch('devhub.tasks.run_validator') + def test_flag_binary(self, _mock): + _mock.return_value = '{"metadata":{"contains_binary_extension": 1}}' + flag_binary([self.addon.pk]) + eq_(Addon.objects.get(pk=self.addon.pk).binary, True) + + @mock.patch('devhub.tasks.run_validator') + def test_flag_not_binary(self, _mock): + _mock.return_value = '{"metadata":{"contains_binary_extension": 0}}' + flag_binary([self.addon.pk]) + eq_(Addon.objects.get(pk=self.addon.pk).binary, False) + + @mock.patch('devhub.tasks.run_validator') + def test_flag_error(self, _mock): + _mock.side_effect = RuntimeError() + flag_binary([self.addon.pk]) + eq_(Addon.objects.get(pk=self.addon.pk).binary, False)