arkindex/project/elastic.py

@@ -34,7 +34,7 @@ class ESQuerySet(object):
             index=self.es_index,
             body={
                 "from": value.start,
-                "size": value.stop - value.start,
+                "size": min(value.stop - value.start, settings.ES_RESULTS_LIMIT),
                 "_source": self._source,
                 "query": self.query,
                 "aggs": self.aggs,
@@ -48,13 +48,13 @@ class ESQuerySet(object):
         return self.count()
 
     def count(self):
-        return self.elastic.count(
+        return min(self.elastic.count(
             index=self.es_index,
             body={
                 "query": self.query,
             },
             doc_type=self.es_type,
-        )['count']
+        )['count'], settings.ES_RESULTS_LIMIT)
 
     def __iter__(self):
         return self[0:len(self)]

arkindex/project/settings.py

@@ -214,6 +214,8 @@ REST_FRAMEWORK = {
 ELASTIC_SEARCH_HOSTS = [
     os.environ.get('ES_HOST', 'localhost'),
 ]
+# The Scroll API is required to go over 10K results
+ES_RESULTS_LIMIT = 10000
 ES_INDEX_TRANSCRIPTIONS = 'transcriptions'
 ES_INDEX_ACTS = 'acts'
 

arkindex/project/tests/__init__.py

@@ -57,6 +57,7 @@ class RedisMockMixin(object):
     def tearDown(self):
         for p in self.patches:
             p.stop()
+        self.messages.stop()
         self.redis.flushall()
 
 

arkindex/project/tests/test_elastic.py

+from unittest import TestCase
+from unittest.mock import patch, MagicMock
+from arkindex.project.elastic import ESQuerySet
+
+
+class TestESQuerySet(TestCase):
+
+    @classmethod
+    def setUpClass(cls):
+        cls.es_index = 'es_index'
+        cls.es_type = 'es_type'
+        cls.query = {
+            "some": "query"
+        }
+        cls.aggs = {
+            "some": "aggs"
+        }
+        cls.sort = {
+            "some sort of": "sorts"
+        }
+        cls.post_process_mock = MagicMock()
+
+        cls.es_mock = patch('arkindex.project.elastic.Elasticsearch').start()
+
+        cls.esqs = ESQuerySet(
+            es_index=cls.es_index,
+            es_type=cls.es_type,
+            post_process=cls.post_process_mock,
+            query=cls.query,
+            aggs=cls.aggs,
+            sort=cls.sort,
+        )
+
+    def setUp(self):
+        self.es_mock.reset_mock()
+        self.es_mock().reset_mock()
+        self.post_process_mock.reset_mock()
+
+    def test_count(self):
+        """
+        Check the ESQuerySet supports the count() and __len__ methods and calls ES correctly
+        """
+        self.es_mock().count.return_value = {'count': 42}
+
+        self.assertEqual(self.esqs.count(), 42)
+        self.assertEqual(len(self.esqs), 42)
+
+        self.assertEqual(self.es_mock().count.call_count, 2)
+        args, kwargs = self.es_mock().count.call_args
+        self.assertEqual(len(args), 0)
+        self.assertDictEqual(kwargs, {
+            "index": self.es_index,
+            "doc_type": self.es_type,
+            "body": {
+                "query": self.query,
+            },
+        })
+
+    def test_iter(self):
+        """
+        Check trying to iterate on a ESQuerySet fetches all the results from ES
+        """
+        self.es_mock().count.return_value = {'count': 42}
+        self.es_mock().search.return_value = {'some': 'results'}
+
+        iter(self.esqs)
+
+        self.assertEqual(self.es_mock().count.call_count, 1)
+        self.assertEqual(self.es_mock().search.call_count, 1)
+        self.assertEqual(self.post_process_mock.call_count, 1)
+
+        args, kwargs = self.es_mock().search.call_args
+        self.assertEqual(len(args), 0)
+        self.assertEqual(kwargs['body']['from'], 0)
+        self.assertEqual(kwargs['body']['size'], 42)
+
+        args, kwargs = self.post_process_mock.call_args
+        self.assertTupleEqual(args, ({'some': 'results'}, ))
+
+    def test_slice(self):
+        """
+        Check retrieving a slice from the ESQuerySEt retrieves only this slice from ES
+        """
+        self.esqs[10:15]
+
+        self.assertEqual(self.es_mock().search.call_count, 1)
+        self.assertEqual(self.post_process_mock.call_count, 1)
+
+        args, kwargs = self.es_mock().search.call_args
+        self.assertEqual(len(args), 0)
+        self.assertDictEqual(kwargs, {
+            "index": self.es_index,
+            "doc_type": self.es_type,
+            "body": {
+                "from": 10,
+                "size": 5,
+                "_source": True,
+                "query": self.query,
+                "aggs": self.aggs,
+                "sort": self.sort,
+            }
+        })
+
+    def test_single_item(self):
+        """
+        Check retrieving a single item by index from the ESQuerySet retrieves a single item from ES
+        """
+        self.esqs[10]
+
+        self.assertEqual(self.es_mock().search.call_count, 1)
+        self.assertEqual(self.post_process_mock.call_count, 1)
+
+        args, kwargs = self.es_mock().search.call_args
+        self.assertEqual(len(args), 0)
+        self.assertDictEqual(kwargs, {
+            "index": self.es_index,
+            "doc_type": self.es_type,
+            "body": {
+                "from": 10,
+                "size": 1,
+                "_source": True,
+                "query": self.query,
+                "aggs": self.aggs,
+                "sort": self.sort,
+            }
+        })
+
+    def test_large_results(self):
+        """
+        Check the ESQuerySet caps to 10,000 results to prevent errors without the scroll API
+        """
+        self.es_mock().count.return_value = {'count': 99999}
+        self.es_mock().search.return_value = {'some': 'results'}
+
+        self.assertEqual(len(self.esqs), 10000)
+        self.esqs[500:20000]
+
+        self.assertEqual(self.es_mock().count.call_count, 1)
+        self.assertEqual(self.es_mock().search.call_count, 1)
+        self.assertEqual(self.post_process_mock.call_count, 1)
+
+        args, kwargs = self.es_mock().search.call_args
+        self.assertEqual(len(args), 0)
+        self.assertEqual(kwargs['body']['from'], 500)
+        self.assertEqual(kwargs['body']['size'], 10000)