Scan refactor (elastic#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

Author: 0bsearch

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 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
- )
-
+ while scroll_id and resp['hits']['hits']:
  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 @@ def test_errors_are_collected_properly(self):
 
 
 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 @@ def test_all_documents_are_read(self):
  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):