Added update establishment image urls command

apps/establishment/management/commands/update_establishment_image_urls.py

@@ -0,0 +1,46 @@
+import math
+
+from django.core.management.base import BaseCommand
+from celery import group
+from establishment.models import Establishment
+from establishment.tasks import update_establishment_image_urls
+
+
+class Command(BaseCommand):
+    help = """
+    Updating image links for establishments.
+    Run command ./manage.py update_establishment_image_urls
+    """
+
+    def add_arguments(self, parser):
+        parser.add_argument(
+            '--bucket_size',
+            type=int,
+            help='Size of one basket to update'
+        )
+
+    def handle(self, *args, **kwargs):
+        bucket_size = kwargs.pop('bucket_size') if kwargs.get('bucket_size') else 128
+
+        objects = Establishment.objects.all()
+        objects_size = objects.count()
+        summary_tasks = math.ceil(objects_size / bucket_size)
+
+        tasks = []
+
+        for index in range(0, objects_size, bucket_size):
+            bucket = objects[index: index + bucket_size]
+
+            task = update_establishment_image_urls.s(
+                (index + bucket_size) / bucket_size, summary_tasks,
+                bucket.values_list('id', flat=True)
+            )
+
+            tasks.append(task)
+
+        self.stdout.write(self.style.WARNING(f'Created all celery update tasks.\n'))
+
+        job = group(*tasks)
+        job.delay()
+
+        self.stdout.write(self.style.WARNING(f'Done all celery update tasks.\n'))

apps/establishment/tasks.py

@@ -1,17 +1,15 @@
 """Establishment app tasks."""
 import logging
 
+import requests
 from celery import shared_task
 from celery.schedules import crontab
 from celery.task import periodic_task
-
-from django.core import management
-from django_elasticsearch_dsl.management.commands import search_index
-
 from django_elasticsearch_dsl.registries import registry
+
 from establishment import models
+from establishment.models import Establishment
 from location.models import Country
-from search_indexes.documents.establishment import EstablishmentDocument
 
 logger = logging.getLogger(__name__)
 
@@ -28,6 +26,7 @@ def recalculate_price_levels_by_country(country_id):
             establishment.recalculate_price_level(low_price=country.low_price,
                                                   high_price=country.high_price)
 
+
 # @periodic_task(run_every=crontab(minute=59))
 # def rebuild_establishment_indices():
 #     management.call_command(search_index.Command(), action='populate', models=[models.Establishment.__name__],
@@ -50,3 +49,31 @@ def recalculation_public_mark(establishment_id):
     establishment = models.Establishment.objects.get(id=establishment_id)
     establishment.recalculate_public_mark()
     establishment.recalculate_toque_number()
+
+
+@shared_task
+def update_establishment_image_urls(part_number: int, summary_tasks: int, bucket_ids: list):
+    queryset = Establishment.objects.filter(id__in=bucket_ids)
+
+    for establishment in queryset:
+        for data in [
+            ('image_url', establishment.image_url),
+            ('preview_image_url', establishment.preview_image_url)
+        ]:
+            attr, url = data
+
+            if establishment.image_url is not None:
+                try:
+                    response = requests.get(url, allow_redirects=True)
+
+                    if response.status_code != 200:
+                        setattr(establishment, attr, None)
+
+                except (
+                        requests.exceptions.ConnectionError,
+                        requests.exceptions.ConnectTimeout
+                ):
+                    setattr(establishment, attr, None)
+
+    logger.info(f'The {part_number}th part of the image update '
+                f'from {summary_tasks} parts was completed')