Use ES to get apps for a collection in search/featured (bug 927420)

migrations/681-add-collection-es-waffle.sql

@@ -0,0 +1,2 @@
+INSERT INTO waffle_switch_mkt (name, active, created, modified, note)
+       VALUES ('collections-use-es-for-apps', 0, NOW(), NOW(), 'Use ElasticSearch to fetch apps in Collections returned by client-side APIs');

mkt/collections/models.py

@@ -9,6 +9,7 @@ from addons.models import Category, clean_slug
 from amo.decorators import use_master
 from amo.utils import to_language
 from mkt.webapps.models import Webapp
+from mkt.webapps.tasks import index_webapps
 from translations.fields import PurifiedField, save_signal
 
 from .constants import COLLECTION_TYPES
@@ -42,6 +43,8 @@ class Collection(amo.models.ModelBase):
     can_be_hero = models.BooleanField(default=False, help_text=(
         'Indicates whether an operator shelf collection can be displayed with'
         'a hero graphic'))
+    apps = models.ManyToManyField(Webapp, through='CollectionMembership',
+                                  related_name='appcollections')
 
     objects = amo.models.ManagerBase()
     public = PublicCollectionsManager()
@@ -72,24 +75,20 @@ class Collection(amo.models.ModelBase):
                             str(self.pk / 1000),
                             'app_collection_%s.png' % (self.pk,))
 
-    def apps(self):
-        """
-        Return a list containing all apps in this collection.
-        """
-        return [a.app for a in self.collectionmembership_set.all()]
-
     def add_app(self, app, order=None):
         """
         Add an app to this collection. If specified, the app will be created
         with the specified `order`. If not, it will be added to the end of the
         collection.
         """
-        if not order:
+        if order is None:
             qs = CollectionMembership.objects.filter(collection=self)
             aggregate = qs.aggregate(models.Max('order'))['order__max']
             order = aggregate + 1 if aggregate is not None else 0
-        return CollectionMembership.objects.create(collection=self, app=app,
+        rval = CollectionMembership.objects.create(collection=self, app=app,
                                                    order=order)
+        index_webapps([app.pk])
+        return rval
 
     def remove_app(self, app):
         """
@@ -102,6 +101,7 @@ class Collection(amo.models.ModelBase):
             return False
         else:
             membership.delete()
+            index_webapps([app.pk])
             return True
 
     def reorder(self, new_order):
@@ -114,11 +114,12 @@ class Collection(amo.models.ModelBase):
         passed order. A ValueError will be raised if each app in the
         collection is not included in the ditionary.
         """
-        if set(a.pk for a in self.apps()) != set(new_order):
+        if set(a.pk for a in self.apps.all()) != set(new_order):
             raise ValueError('Not all apps included')
         for order, pk in enumerate(new_order):
             CollectionMembership.objects.get(collection=self,
                                              app_id=pk).update(order=order)
+        index_webapps(new_order)
 
     def has_curator(self, userprofile):
         """

mkt/collections/serializers.py

@@ -8,9 +8,11 @@ from rest_framework.compat import smart_text
 from rest_framework.reverse import reverse
 from tastypie.bundle import Bundle
 from tower import ugettext_lazy as _
+import waffle
 
 from django.conf import settings
-from django.core.exceptions import ObjectDoesNotExist, ValidationError
+from django.core.exceptions import (ImproperlyConfigured, ObjectDoesNotExist,
+                                    ValidationError)
 from django.core.files.base import File
 from django.core.files.storage import default_storage as storage
 
@@ -24,7 +26,8 @@ import mkt
 from addons.models import Category
 from mkt.api.fields import TranslationSerializerField
 from mkt.api.resources import AppResource
-from mkt.constants.features import FeatureProfile
+from mkt.features.utils import get_feature_profile
+from mkt.webapps.models import Webapp
 from users.models import UserProfile
 
 from .models import Collection
@@ -44,22 +47,49 @@ class CollectionMembershipField(serializers.RelatedField):
         return AppResource().full_dehydrate(bundle).data
 
     def field_to_native(self, obj, field_name):
+        if not hasattr(self, 'context') or not 'request' in self.context:
+            raise ImproperlyConfigured('Pass request in self.context when'
+                                       ' using CollectionMembershipField.')
+
+        request = self.context['request']
+
+        # If we have a search resource in the context, we should try to use ES
+        # to fetch the apps, if the waffle flag is active.
+        if ('search_resource' in self.context
+            and waffle.switch_is_active('collections-use-es-for-apps')):
+            return self.field_to_native_es(obj, request)
+
         value = get_component(obj, self.source)
 
         # Filter apps based on feature profiles.
-        if hasattr(self, 'context') and 'request' in self.context:
-            sig = self.context['request'].GET.get('pro')
-            if sig:
-                try:
-                    profile = FeatureProfile.from_signature(sig)
-                except ValueError:
-                    pass
-                else:
-                    value = value.filter(**profile.to_kwargs(
-                        prefix='app___current_version__features__has_'))
+        profile = get_feature_profile(request)
+        if profile and waffle.switch_is_active('buchets'):
+            value = value.filter(**profile.to_kwargs(
+                prefix='app___current_version__features__has_'))
 
         return [self.to_native(item) for item in value.all()]
 
+    def field_to_native_es(self, obj, request):
+        """
+        A version of field_to_native that uses ElasticSearch to fetch the apps
+        belonging to the collection instead of SQL.
+
+        Relies on a SearchResource instance in self.context['search_resource']
+        to properly rehydrate results returned by ES.
+        """
+        search_resource = self.context['search_resource']
+        profile = get_feature_profile(request)
+        region = search_resource.get_region(request)
+
+        qs = Webapp.from_search(request, region=region)
+        filters = {'collection.id': obj.pk}
+        if profile and waffle.switch_is_active('buchets'):
+            filters.update(**profile.to_kwargs(prefix='features.has_'))
+        qs = qs.filter(**filters).order_by('collection.order')
+
+        return [bundle.data
+                for bundle in search_resource.rehydrate_results(request, qs)]
+
 
 class HyperlinkedRelatedOrNullField(serializers.HyperlinkedRelatedField):
     read_only = True

mkt/collections/tests/test_models.py

@@ -53,14 +53,14 @@ class TestCollection(amo.tests.TestCase):
         eq_(added.app, self.apps[2])
         eq_(added.collection, self.collection)
 
-        eq_(self.collection.apps(), [self.apps[2], self.apps[1]])
+        eq_(self.collection.apps.all(), [self.apps[2], self.apps[1]])
 
     def test_apps(self):
         self._generate_apps()
 
-        self.assertSetEqual(self.collection.apps(), [])
+        self.assertSetEqual(self.collection.apps.all(), [])
         self._add_apps()
-        self.assertSetEqual(self.collection.apps(), self.apps)
+        self.assertSetEqual(self.collection.apps.all(), self.apps)
         eq_(list(CollectionMembership.objects.values_list('order', flat=True)),
             [0, 1, 2, 3])
 
@@ -70,7 +70,7 @@ class TestCollection(amo.tests.TestCase):
         extra_app = amo.tests.app_factory()
         added = self.collection.add_app(extra_app, order=3)
         eq_(added.order, 3)
-        self.assertSetEqual(self.collection.apps(), [extra_app])
+        self.assertSetEqual(self.collection.apps.all(), [extra_app])
         self._add_apps()
         eq_(list(CollectionMembership.objects.values_list('order', flat=True)),
             [3, 4, 5, 6, 7])

mkt/collections/tests/test_serializers.py

@@ -9,7 +9,9 @@ import amo.tests
 from addons.models import Category
 from amo.urlresolvers import reverse
 
+import mkt
 from mkt.api.resources import AppResource
+from mkt.search.api import WithFeaturedResource
 from mkt.collections.constants import (COLLECTIONS_TYPE_BASIC,
                                        COLLECTIONS_TYPE_OPERATOR)
 from mkt.constants.features import FeatureProfile
@@ -28,16 +30,23 @@ class CollectionDataMixin(object):
     }
 
 
-class TestCollectionMembershipField(CollectionDataMixin, amo.tests.TestCase):
+class BaseTestCollectionMembershipField(object):
 
     def setUp(self):
         self.collection = Collection.objects.create(**self.collection_data)
         self.app = amo.tests.app_factory()
-        self.collection.add_app(self.app)
+        self.collection.add_app(self.app, order=1)
         self.field = CollectionMembershipField()
+        self.field.context = {}
         self.membership = CollectionMembership.objects.all()[0]
-        self.request = RequestFactory()
         self.profile = FeatureProfile(apps=True).to_signature()
+        self.create_switch('buchets')
+
+    def get_request(self, query_string):
+        request = RequestFactory().get('/', query_string)
+        request.REGION = mkt.regions.WORLDWIDE
+        request.API = True
+        return request
 
     def test_to_native(self):
         resource = AppResource().full_dehydrate(Bundle(obj=self.app))
@@ -48,39 +57,75 @@ class TestCollectionMembershipField(CollectionDataMixin, amo.tests.TestCase):
             else:
                 eq_(value, resource.data[key])
 
-    def test_field_to_native(self):
+    def _field_to_native_profile(self, profile='0.0'):
+        request = self.get_request({'pro': profile, 'dev': 'firefoxos'})
         self.field.parent = self.collection
         self.field.source = 'collectionmembership_set'
-        self.field.context = {
-            'request': self.request.get('/', {'pro': self.profile})
-        }
+        self.field.context['request'] = request
 
-        # Ensure that the one app is included in the default response.
-        native = self.field.field_to_native(self.collection,
-                                           'collectionmembership_set')
-        eq_(len(native), 1)
-        eq_(int(native[0]['id']), self.app.id)
+        return self.field.field_to_native(self.collection,
+                                          'collectionmembership_set')
 
-        # ...but is not included when there is a feature profile mismatch.
+    def test_ordering(self):
+        self.app2 = amo.tests.app_factory()
+        self.collection.add_app(self.app2, order=0)
+        result = self._field_to_native_profile()
+        eq_(len(result), 2)
+        eq_(int(result[0]['id']), self.app2.id)
+        eq_(int(result[1]['id']), self.app.id)
+
+    def test_field_to_native_profile(self):
+        result = self._field_to_native_profile(self.profile)
+        eq_(len(result), 1)
+        eq_(int(result[0]['id']), self.app.id)
+
+    def test_field_to_native_profile_mismatch(self):
         self.app.current_version.features.update(has_geolocation=True)
-        native_without = self.field.field_to_native(self.collection,
-                                                    'collectionmembership_set')
-        eq_(len(native_without), 0)
+        result = self._field_to_native_profile(self.profile)
+        eq_(len(result), 0)
 
     def test_field_to_native_invalid_profile(self):
-        self.field.parent = self.collection
-        self.field.source = 'collectionmembership_set'
-        self.field.context = {
-            'request': self.request.get('/', {'pro': 'muahaha'})
-        }
-
-        # Ensure that no exceptions were raised.
-        native = self.field.field_to_native(self.collection,
-                                            'collectionmembership_set')
-
+        result = self._field_to_native_profile('muahahah')
         # Ensure that no filtering took place.
-        eq_(len(native), 1)
-        eq_(int(native[0]['id']), self.app.id)
+        eq_(len(result), 1)
+        eq_(int(result[0]['id']), self.app.id)
+
+
+class TestCollectionMembershipField(BaseTestCollectionMembershipField,
+                                    CollectionDataMixin, amo.tests.TestCase):
+    pass
+
+
+class TestCollectionMembershipFieldES(BaseTestCollectionMembershipField,
+                                      CollectionDataMixin,
+                                      amo.tests.ESTestCase):
+    """ Same tests as TestCollectionMembershipField above, but we need a
+        different setUp and more importantly, we need to force a sync refresh
+        in ES when we modify our app """
+
+    def setUp(self):
+        self.create_switch('collections-use-es-for-apps')
+        super(TestCollectionMembershipFieldES, self).setUp()
+        self.field.context['search_resource'] = WithFeaturedResource()
+        self.refresh('webapp')
+
+    def test_field_to_native_profile_mismatch(self):
+        self.app.current_version.features.update(has_geolocation=True)
+        # FIXME: a simple refresh() wasn't enough, don't we reindex apps when
+        # feature profiles change ? Investigate.
+        self.reindex(self.app.__class__, 'webapp')
+        result = self._field_to_native_profile(self.profile)
+        eq_(len(result), 0)
+
+    def test_ordering(self):
+        self.app2 = amo.tests.app_factory()
+        amo.tests.app_factory()  # Extra app not belonging to a collection.
+        self.collection.add_app(self.app2, order=0)
+        self.refresh('webapp')
+        result = self._field_to_native_profile()
+        eq_(len(result), 2)
+        eq_(int(result[0]['id']), self.app2.id)
+        eq_(int(result[1]['id']), self.app.id)
 
 
 class TestCollectionSerializer(CollectionDataMixin, amo.tests.TestCase):

mkt/collections/tests/test_views.py

@@ -754,7 +754,7 @@ class TestCollectionViewSetDuplicate(BaseCollectionViewSetTest):
         eq_(res.status_code, 201)
         new_collection = Collection.objects.get(pk=data['id'])
         ok_(new_collection.pk != self.collection.pk)
-        eq_(new_collection.apps(), self.collection.apps())
+        eq_(new_collection.apps.all(), self.collection.apps.all())
         eq_(len(data['apps']), len(self.apps))
         for order, app in enumerate(self.apps):
             eq_(int(data['apps'][order]['id']), self.apps[order].pk)
@@ -984,7 +984,7 @@ class TestCollectionViewSetReorderApps(CollectionViewSetChangeAppsMixin):
         eq_(res.status_code, 400)
         eq_(data['detail'], CollectionViewSet.exceptions['app_mismatch'])
         self.assertSetEqual([a['slug'] for a in data['apps']],
-                            [a.app_slug for a in self.collection.apps()])
+                            [a.app_slug for a in self.collection.apps.all()])
 
 
 class TestCollectionViewSetEditCollection(BaseCollectionViewSetTest):

mkt/collections/views.py

@@ -146,7 +146,7 @@ class CollectionViewSet(CORSMixin, SlugOrIdMixin, viewsets.ModelViewSet):
             return result
 
         # And now, add apps from the original collection.
-        for app in collection.apps():
+        for app in collection.apps.all():
             self.object.add_app(app)
 
         # Re-Serialize to include apps.

mkt/features/utils.py

@@ -0,0 +1,13 @@
+from mkt.constants.features import FeatureProfile
+
+
+def get_feature_profile(request):
+    profile = None
+    if request.GET.get('dev') in ('firefoxos', 'android'):
+        sig = request.GET.get('pro')
+        if sig:
+            try:
+                profile = FeatureProfile.from_signature(sig)
+            except ValueError:
+                pass
+    return profile

mkt/search/api.py

@@ -1,6 +1,5 @@
 from django.conf.urls import url
 
-import waffle
 from tastypie.authorization import ReadOnlyAuthorization
 from tastypie.throttle import BaseThrottle
 from tastypie.utils import trailing_slash
@@ -19,7 +18,7 @@ from mkt.collections.constants import (COLLECTIONS_TYPE_BASIC,
 from mkt.collections.filters import CollectionFilterSetWithFallback
 from mkt.collections.models import Collection
 from mkt.collections.serializers import CollectionSerializer
-from mkt.constants.features import FeatureProfile
+from mkt.features.utils import get_feature_profile
 from mkt.search.views import _filter_search
 from mkt.search.forms import ApiSearchForm
 from mkt.webapps.models import Webapp
@@ -50,14 +49,6 @@ class SearchResource(CORSResource, MarketplaceResource):
             raise self.form_errors(form)
         return form.cleaned_data
 
-    def get_feature_profile(self, request):
-        profile = None
-        if request.GET.get('dev') in ('firefoxos', 'android'):
-            sig = request.GET.get('pro')
-            if sig:
-                profile = FeatureProfile.from_signature(sig)
-        return profile
-
     def get_region(self, request):
         return getattr(request, 'REGION', mkt.regions.WORLDWIDE)
 
@@ -69,7 +60,7 @@ class SearchResource(CORSResource, MarketplaceResource):
 
     def apply_filters(self, request, qs, data=None):
         # Build device features profile filter.
-        profile = self.get_feature_profile(request)
+        profile = get_feature_profile(request)
 
         # Build region filter.
         region = self.get_region(request)
@@ -158,7 +149,8 @@ class WithFeaturedResource(SearchResource):
             qs = Collection.public.all()
         qs = CollectionFilterSetWithFallback(filters, queryset=qs).qs
         serializer = CollectionSerializer(qs[:limit],
-                                          context={'request': request})
+                                          context={'request': request,
+                                                   'search_resource': self})
         return serializer.data, getattr(qs, 'filter_fallback', None)
 
     def alter_list_data_to_serialize(self, request, data):

mkt/search/tests/test_api.py

@@ -488,6 +488,8 @@ class BaseFeaturedTests(BaseOAuth, ESTestCase):
         self.create_switch('buchets')
         self.cat = Category.objects.create(type=amo.ADDON_WEBAPP, slug='shiny')
         self.app = Webapp.objects.get(pk=337141)
+        AddonDeviceType.objects.create(addon=self.app,
+            device_type=DEVICE_CHOICES_IDS['firefoxos'])
         AddonCategory.objects.get_or_create(addon=self.app, category=self.cat)
         self.profile = FeatureProfile(apps=True, audio=True, fullscreen=True,
                                       geolocation=True, indexeddb=True,
@@ -509,6 +511,7 @@ class TestFeaturedCollections(BaseFeaturedTests):
             collection_type=self.col_type, category=self.cat, is_public=True,
             region=mkt.regions.US.id)
         self.qs['region'] = mkt.regions.US.slug
+        self.create_switch('collections-use-es-for-apps')
         # FIXME: mock the search part, we don't care about it
 
     def make_request(self):
@@ -525,6 +528,8 @@ class TestFeaturedCollections(BaseFeaturedTests):
 
     def test_apps_included(self):
         self.col.add_app(self.app)
+        self.refresh('webapp')
+
         res, json = self.test_added_to_results()
         eq_(len(json[self.prop_name][0]['apps']), 1)
 
@@ -533,9 +538,9 @@ class TestFeaturedCollections(BaseFeaturedTests):
         Test that the app list is passed through feature profile filtering.
         """
         self.app.current_version.features.update(has_pay=True)
-        self.reindex(Webapp, 'webapp')
-        self.qs.update({'dev': ''})
         self.col.add_app(self.app)
+        self.refresh('webapp')
+
         res, json = self.test_added_to_results()
         eq_(len(json[self.prop_name][0]['apps']), 0)
 

mkt/webapps/models.py

@@ -1047,6 +1047,13 @@ class WebappIndexer(MappingType, Indexable):
                         'type': 'string',
                         'index': 'not_analyzed'
                     },
+                    'collection': {
+                        'type': 'object',
+                        'properties': {
+                            'id': {'type': 'long'},
+                            'order': {'type': 'short'}
+                        }
+                    },
                     'content_ratings': {
                         'type': 'object',
                         'dynamic': 'true',
@@ -1204,6 +1211,8 @@ class WebappIndexer(MappingType, Indexable):
         d['app_type'] = obj.app_type_id
         d['author'] = obj.developer_name
         d['category'] = list(obj.categories.values_list('slug', flat=True))
+        d['collection'] = [{'id': cms.collection_id, 'order': cms.order}
+                           for cms in obj.collectionmembership_set.all()]
         d['content_ratings'] = content_ratings if content_ratings else None
         d['current_version'] = version.version if version else None
         d['default_locale'] = obj.default_locale