Scan refactor (#924)

* Adds scan test for exception & data yielded

* Adds scan test for fast route & initial search error

* Refactores scan implementation; better scroll test

* Adds tests clear_scroll & logger

elasticsearch/helpers/actions.py

@@ -427,33 +427,19 @@ def scan(
     if not preserve_order:
         query = query.copy() if query else {}
         query["sort"] = "_doc"
+
     # initial search
     resp = client.search(
         body=query, scroll=scroll, size=size, request_timeout=request_timeout, **kwargs
     )
-
     scroll_id = resp.get("_scroll_id")
-    if scroll_id is None:
-        return
 
     try:
-        first_run = True
+        while scroll_id and resp['hits']['hits']:
-        while True:
-            # if we didn't set search_type to scan initial search contains data
-            if first_run:
-                first_run = False
-            else:
-                resp = client.scroll(
-                    scroll_id,
-                    scroll=scroll,
-                    request_timeout=request_timeout,
-                    **scroll_kwargs
-                )
-
             for hit in resp["hits"]["hits"]:
                 yield hit
 
-            # check if we have any errrors
+            # check if we have any errors
             if resp["_shards"]["successful"] < resp["_shards"]["total"]:
                 logger.warning(
                     "Scroll request has only succeeded on %d shards out of %d.",
@@ -467,10 +453,14 @@ def scan(
                         % (resp["_shards"]["successful"], resp["_shards"]["total"]),
                     )
 
+            resp = client.scroll(
+                scroll_id,
+                scroll=scroll,
+                request_timeout=request_timeout,
+                **scroll_kwargs
+            )
             scroll_id = resp.get("_scroll_id")
-            # end of scroll
+
-            if scroll_id is None or not resp["hits"]["hits"]:
-                break
     finally:
         if scroll_id and clear_scroll:
             client.clear_scroll(body={"scroll_id": [scroll_id]}, ignore=(404,))

test_elasticsearch/test_server/test_helpers.py

@@ -1,4 +1,7 @@
+from mock import patch
+
 from elasticsearch import helpers, TransportError
+from elasticsearch.helpers import ScanError
 
 from . import ElasticsearchTestCase
 from ..test_cases import SkipTest
@@ -305,6 +308,24 @@ class TestBulk(ElasticsearchTestCase):
 
 
 class TestScan(ElasticsearchTestCase):
+    mock_scroll_responses = [
+        {
+            '_scroll_id': 'dummy_id',
+            '_shards': {'successful': 4, 'total': 5},
+            'hits': {'hits': [{'scroll_data': 42}]},
+        },
+        {
+            '_scroll_id': 'dummy_id',
+            '_shards': {'successful': 4, 'total': 5},
+            'hits': {'hits': []},
+        },
+    ]
+
+    @classmethod
+    def tearDownClass(cls):
+        cls.client.transport.perform_request('DELETE', '/_search/scroll/_all')
+        super(TestScan, cls).tearDownClass()
+
     def test_order_can_be_preserved(self):
         bulk = []
         for x in range(100):
@@ -338,6 +359,117 @@ class TestScan(ElasticsearchTestCase):
         self.assertEquals(set(map(str, range(100))), set(d["_id"] for d in docs))
         self.assertEquals(set(range(100)), set(d["_source"]["answer"] for d in docs))
 
+    def test_scroll_error(self):
+        bulk = []
+        for x in range(4):
+            bulk.append({"index": {"_index": "test_index", "_type": "_doc"}})
+            bulk.append({"value": x})
+        self.client.bulk(bulk, refresh=True)
+
+        with patch.object(self.client, 'scroll') as scroll_mock:
+            scroll_mock.side_effect = self.mock_scroll_responses
+            data = list(helpers.scan(
+                self.client,
+                index='test_index',
+                size=2,
+                raise_on_error=False,
+                clear_scroll=False
+            ))
+            self.assertEqual(len(data), 3)
+            self.assertEqual(data[-1], {'scroll_data': 42})
+
+            scroll_mock.side_effect = self.mock_scroll_responses
+            with self.assertRaises(ScanError):
+                data = list(helpers.scan(
+                    self.client,
+                    index='test_index',
+                    size=2,
+                    raise_on_error=True,
+                    clear_scroll=False
+                ))
+            self.assertEqual(len(data), 3)
+            self.assertEqual(data[-1], {'scroll_data': 42})
+
+    def test_initial_search_error(self):
+        with patch.object(self, 'client') as client_mock:
+            client_mock.search.return_value = {
+                '_scroll_id': 'dummy_id',
+                '_shards': {'successful': 4, 'total': 5},
+                'hits': {'hits': [{'search_data': 1}]},
+            }
+            client_mock.scroll.side_effect = self.mock_scroll_responses
+
+            data = list(helpers.scan(self.client, index='test_index', size=2, raise_on_error=False))
+            self.assertEqual(data, [{'search_data': 1}, {'scroll_data': 42}])
+
+            client_mock.scroll.side_effect = self.mock_scroll_responses
+            with self.assertRaises(ScanError):
+                data = list(
+                    helpers.scan(self.client, index='test_index', size=2, raise_on_error=True)
+                )
+                self.assertEqual(data, [{'search_data': 1}])
+                client_mock.scroll.assert_not_called()
+
+    def test_no_scroll_id_fast_route(self):
+        with patch.object(self, 'client') as client_mock:
+            client_mock.search.return_value = {'no': '_scroll_id'}
+            data = list(helpers.scan(self.client, index='test_index'))
+
+            self.assertEqual(data, [])
+            client_mock.scroll.assert_not_called()
+            client_mock.clear_scroll.assert_not_called()
+
+    @patch('elasticsearch.helpers.actions.logger')
+    def test_logger(self, logger_mock):
+        bulk = []
+        for x in range(4):
+            bulk.append({'index': {'_index': 'test_index', '_type': '_doc'}})
+            bulk.append({'value': x})
+        self.client.bulk(bulk, refresh=True)
+
+        with patch.object(self.client, 'scroll') as scroll_mock:
+            scroll_mock.side_effect = self.mock_scroll_responses
+            list(helpers.scan(
+                self.client,
+                index='test_index',
+                size=2,
+                raise_on_error=False,
+                clear_scroll=False
+            ))
+            logger_mock.warning.assert_called()
+
+            scroll_mock.side_effect = self.mock_scroll_responses
+            try:
+                list(helpers.scan(
+                    self.client,
+                    index='test_index',
+                    size=2,
+                    raise_on_error=True,
+                    clear_scroll=False
+                ))
+            except ScanError:
+                pass
+            logger_mock.warning.assert_called()
+
+    def test_clear_scroll(self):
+        bulk = []
+        for x in range(4):
+            bulk.append({'index': {'_index': 'test_index', '_type': '_doc'}})
+            bulk.append({'value': x})
+        self.client.bulk(bulk, refresh=True)
+
+        with patch.object(self.client, 'clear_scroll', wraps=self.client.clear_scroll) as spy:
+            list(helpers.scan(self.client, index='test_index', size=2))
+            spy.assert_called_once()
+
+            spy.reset_mock()
+            list(helpers.scan(self.client, index='test_index', size=2, clear_scroll=True))
+            spy.assert_called_once()
+
+            spy.reset_mock()
+            list(helpers.scan(self.client, index='test_index', size=2, clear_scroll=False))
+            spy.assert_not_called()
+
 
 class TestReindex(ElasticsearchTestCase):
     def setUp(self):