Extra fields for establishments

apps/establishment/models.py

@@ -14,7 +14,7 @@ from django.contrib.postgres.fields import ArrayField
 from django.core.exceptions import ValidationError
 from django.core.validators import MinValueValidator, MaxValueValidator
 from django.db import models
-from django.db.models import When, Case, F, ExpressionWrapper, Subquery, Q
+from django.db.models import When, Case, F, ExpressionWrapper, Subquery, Q, Prefetch
 from django.utils import timezone
 from django.utils.translation import gettext_lazy as _
 from phonenumber_field.modelfields import PhoneNumberField
@@ -23,6 +23,7 @@ from timezone_field import TimeZoneField
 from collection.models import Collection
 from location.models import Address
 from main.models import Award, Currency
+from tag.models import Tag
 from review.models import Review
 from utils.models import (ProjectBaseMixin, TJSONField, URLImageMixin,
                           TranslatedFieldsMixin, BaseAttributes, GalleryModelMixin,
@@ -321,6 +322,13 @@ class EstablishmentQuerySet(models.QuerySet):
         """Exclude countries."""
         return self.exclude(address__city__country__in=countries)
 
+    def with_certain_tag_category_related(self, index_name, attr_name):
+        """Includes extra tags."""
+        return self.prefetch_related(
+            Prefetch('tags', queryset=Tag.objects.filter(category__index_name=index_name),
+                     to_attr=attr_name)
+        )
+
 
 class Establishment(GalleryModelMixin, ProjectBaseMixin, URLImageMixin,
                     TranslatedFieldsMixin, HasTagsMixin, FavoritesMixin):

apps/establishment/serializers/common.py

@@ -320,12 +320,16 @@ class EstablishmentListRetrieveSerializer(EstablishmentBaseSerializer):
 
     address = AddressDetailSerializer()
     schedule = ScheduleRUDSerializer(many=True, allow_null=True)
+    restaurant_category = TagBaseSerializer(read_only=True, many=True, allow_null=True)
+    restaurant_cuisine = TagBaseSerializer(read_only=True, many=True, allow_null=True)
 
     class Meta(EstablishmentBaseSerializer.Meta):
         """Meta class."""
 
         fields = EstablishmentBaseSerializer.Meta.fields + [
             'schedule',
+            'restaurant_category',
+            'restaurant_cuisine',
         ]
 
 

apps/establishment/views/web.py

@@ -35,7 +35,9 @@ class EstablishmentListView(EstablishmentMixinView, generics.ListAPIView):
 
     def get_queryset(self):
         return super().get_queryset().with_schedule() \
-            .with_extended_address_related().with_currency_related()
+            .with_extended_address_related().with_currency_related() \
+            .with_certain_tag_category_related('category', 'restaurant_category') \
+            .with_certain_tag_category_related('cuisine', 'restaurant_cuisine') \
 
 
 class EstablishmentRetrieveView(EstablishmentMixinView, generics.RetrieveAPIView):

apps/search_indexes/documents/establishment.py

@@ -49,6 +49,22 @@ class EstablishmentDocument(Document):
             'value': fields.KeywordField(),
         },
         multi=True)
+    restaurant_category = fields.ObjectField(
+        properties={
+            'id': fields.IntegerField(attr='id'),
+            'label': fields.ObjectField(attr='label_indexing',
+                                        properties=OBJECT_FIELD_PROPERTIES),
+            'value': fields.KeywordField(),
+        },
+        multi=True)
+    restaurant_cuisine = fields.ObjectField(
+        properties={
+            'id': fields.IntegerField(attr='id'),
+            'label': fields.ObjectField(attr='label_indexing',
+                                        properties=OBJECT_FIELD_PROPERTIES),
+            'value': fields.KeywordField(),
+        },
+        multi=True)
     visible_tags = fields.ObjectField(
         properties={
             'id': fields.IntegerField(attr='id'),
@@ -142,4 +158,6 @@ class EstablishmentDocument(Document):
         )
 
     def get_queryset(self):
-        return super().get_queryset().with_es_related()
+        return super().get_queryset().with_es_related() \
+            .with_certain_tag_category_related('category', 'restaurant_category') \
+            .with_certain_tag_category_related('cuisine', 'restaurant_cuisine')

apps/search_indexes/serializers.py

@@ -229,6 +229,8 @@ class EstablishmentDocumentSerializer(InFavoritesMixin, DocumentSerializer):
     establishment_subtypes = EstablishmentTypeSerializer(many=True)
     address = AddressDocumentSerializer(allow_null=True)
     tags = TagsDocumentSerializer(many=True, source='visible_tags')
+    restaurant_category = TagsDocumentSerializer(many=True)
+    restaurant_cuisine = TagsDocumentSerializer(many=True)
     schedule = ScheduleDocumentSerializer(many=True, allow_null=True)
 
     class Meta:
@@ -247,6 +249,8 @@ class EstablishmentDocumentSerializer(InFavoritesMixin, DocumentSerializer):
             'preview_image',
             'address',
             'tags',
+            'restaurant_category',
+            'restaurant_cuisine',
             'schedule',
             'works_noon',
             'works_evening',