Merge pull request #162 from open-craft/jill/es_agg_search_size

Make ElasticSearch aggregate search results size configurable

Author: pomegranited

analytics_data_api/v0/models.py

@@ -407,11 +407,14 @@ def get_course_metadata(cls, course_id):
             }
         }
         """
+        # Use the configured default page size to set the number of aggregate search results.
+        page_size = getattr(settings, 'AGGREGATE_PAGE_SIZE', 10)
+
         search = cls.search()
         search.query = Q('bool', must=[Q('term', course_id=course_id)])
-        search.aggs.bucket('enrollment_modes', 'terms', field='enrollment_mode')
-        search.aggs.bucket('segments', 'terms', field='segments')
-        search.aggs.bucket('cohorts', 'terms', field='cohort')
+        search.aggs.bucket('enrollment_modes', 'terms', field='enrollment_mode', size=page_size)
+        search.aggs.bucket('segments', 'terms', field='segments', size=page_size)
+        search.aggs.bucket('cohorts', 'terms', field='cohort', size=page_size)
         response = search.execute()
         # Build up the map of aggregation name to count
         aggregations = {

analytics_data_api/v0/tests/views/test_learners.py

@@ -14,6 +14,7 @@
 
 from django.conf import settings
 from django.core import management
+from django.test import override_settings
 
 from analyticsdataserver.tests import TestCaseWithAuthentication
 from analytics_data_api.constants import engagement_events
@@ -769,6 +770,29 @@ def test_cohorts(self, course_id, cohorts):
         )
         self.assert_response_matches(self._get(course_id), 200, expected)
 
+    @ddt.data(
+        (CourseSamples.course_ids[0], [], {}),
+        (CourseSamples.course_ids[1], ['Yellow'], {'Yellow': 1}),
+        (CourseSamples.course_ids[1], ['Yellow', 'Green'], {'Yellow': 1, 'Green': 1}),
+        (CourseSamples.course_ids[0], ['Red', 'Red', 'yellow team', 'green'], {'Red': 2, 'green': 1}),
+    )
+    @ddt.unpack
+    @override_settings(AGGREGATE_PAGE_SIZE=2)
+    def test_cohorts_page_size(self, course_id, cohorts, expected_cohorts):
+        """ Ensure that the AGGREGATE_PAGE_SIZE sets the max number of cohorts returned."""
+
+        self.create_learners([
+            {'username': 'user_{}'.format(i), 'course_id': course_id, 'cohort': cohort}
+            for i, cohort in enumerate(cohorts)
+        ])
+        expected = self.get_expected_json(
+            course_id=course_id,
+            segments={'disengaging': 0, 'struggling': 0, 'highly_engaged': 0, 'inactive': 0, 'unenrolled': 0},
+            enrollment_modes={'honor': len(cohorts)} if cohorts else {},
+            cohorts=expected_cohorts,
+        )
+        self.assert_response_matches(self._get(course_id), 200, expected)
+
     @property
     def empty_engagement_ranges(self):
         """ Returns the engagement ranges where all fields are set to None. """

analyticsdataserver/settings/base.py

@@ -310,6 +310,7 @@
 # Warning: using 0 or None for these can alter the structure of the REST response.
 DEFAULT_PAGE_SIZE = 25
 MAX_PAGE_SIZE = 100
+AGGREGATE_PAGE_SIZE = 10
 
 ########## END ANALYTICS DATA API CONFIGURATION