From 2181611a008a8cbd4068bfda6269289d69e9807a Mon Sep 17 00:00:00 2001 From: Slam <3lnc.slam@gmail.com> Date: Fri, 5 Apr 2019 18:31:43 +0300 Subject: [PATCH] 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 | 30 ++-- .../test_server/test_helpers.py | 132 ++++++++++++++++++ 2 files changed, 142 insertions(+), 20 deletions(-) diff --git a/elasticsearch/helpers/actions.py b/elasticsearch/helpers/actions.py index 19dd1a89..cf7f16aa 100644 --- a/elasticsearch/helpers/actions.py +++ b/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,)) diff --git a/test_elasticsearch/test_server/test_helpers.py b/test_elasticsearch/test_server/test_helpers.py index 1e307d51..81fc4f22 100644 --- a/test_elasticsearch/test_server/test_helpers.py +++ b/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):