actions/opensearch-pyd
1
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

Scan refactor (#924)

Browse Source 
* 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
This commit is contained in:
Slam
2019-04-05 09:31:43 -06:00
committed by Nick Lang
parent cf1e946323
commit 2181611a00
2 changed files with 142 additions and 20 deletions
Download Patch File Download Diff File Expand all files Collapse all files
elasticsearch/helpers/actions.py
+10 -20
View File
@@ -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
+132
View File
@@ -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):