Skip to content

Commit 47e8d15

Browse files
authored
[7.16] Always set 'sort' and 'scroll' for scan helper
1 parent 7b741fb commit 47e8d15

File tree

4 files changed

+126
-12
lines changed

4 files changed

+126
-12
lines changed

elasticsearch/_async/helpers.py

Lines changed: 12 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -26,6 +26,7 @@
2626
from ..exceptions import NotFoundError, TransportError
2727
from ..helpers.actions import (
2828
_ActionChunker,
29+
_add_helper_meta_to_kwargs,
2930
_process_bulk_chunk_error,
3031
_process_bulk_chunk_success,
3132
expand_action,
@@ -332,7 +333,9 @@ async def async_scan(
332333
)
333334
334335
"""
335-
scroll_kwargs = scroll_kwargs or {}
336+
scroll_kwargs = scroll_kwargs.copy() if scroll_kwargs else {}
337+
scroll_kwargs["scroll"] = scroll
338+
_add_helper_meta_to_kwargs(scroll_kwargs, "s")
336339

337340
if not preserve_order:
338341
query = query.copy() if query else {}
@@ -356,9 +359,11 @@ async def async_scan(
356359
search_kwargs = kwargs.copy()
357360
if query:
358361
search_kwargs.update(query)
359-
resp = await client.search(
360-
scroll=scroll, size=size, request_timeout=request_timeout, **search_kwargs
361-
)
362+
search_kwargs["scroll"] = scroll
363+
search_kwargs["size"] = size
364+
search_kwargs["request_timeout"] = request_timeout
365+
_add_helper_meta_to_kwargs(search_kwargs, "s")
366+
resp = await client.search(**search_kwargs)
362367
scroll_id = resp.get("_scroll_id")
363368

364369
try:
@@ -390,9 +395,9 @@ async def async_scan(
390395
shards_total,
391396
),
392397
)
393-
resp = await client.scroll(
394-
scroll_id=scroll_id, scroll=scroll, **scroll_kwargs
395-
)
398+
399+
scroll_kwargs["scroll_id"] = scroll_id
400+
resp = await client.scroll(**scroll_kwargs)
396401
scroll_id = resp.get("_scroll_id")
397402

398403
finally:

elasticsearch/helpers/actions.py

Lines changed: 10 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -543,7 +543,8 @@ def scan(
543543
)
544544
545545
"""
546-
scroll_kwargs = scroll_kwargs or {}
546+
scroll_kwargs = scroll_kwargs.copy() if scroll_kwargs else {}
547+
scroll_kwargs["scroll"] = scroll
547548
_add_helper_meta_to_kwargs(scroll_kwargs, "s")
548549

549550
if not preserve_order:
@@ -568,9 +569,11 @@ def scan(
568569
search_kwargs = kwargs.copy()
569570
if query:
570571
search_kwargs.update(query)
571-
resp = client.search(
572-
scroll=scroll, size=size, request_timeout=request_timeout, **search_kwargs
573-
)
572+
search_kwargs["scroll"] = scroll
573+
search_kwargs["size"] = size
574+
search_kwargs["request_timeout"] = request_timeout
575+
_add_helper_meta_to_kwargs(search_kwargs, "s")
576+
resp = client.search(**search_kwargs)
574577
scroll_id = resp.get("_scroll_id")
575578

576579
try:
@@ -602,7 +605,9 @@ def scan(
602605
shards_total,
603606
),
604607
)
605-
resp = client.scroll(scroll_id=scroll_id, scroll=scroll, **scroll_kwargs)
608+
609+
scroll_kwargs["scroll_id"] = scroll_id
610+
resp = client.scroll(**scroll_kwargs)
606611
scroll_id = resp.get("_scroll_id")
607612

608613
finally:

test_elasticsearch/test_async/test_server/test_helpers.py

Lines changed: 53 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -756,6 +756,59 @@ async def test_scan_auth_kwargs_favor_scroll_kwargs_option(
756756
}
757757
assert async_client.scroll.call_args[1]["sort"] == "asc"
758758

759+
async def test_scan_duplicate_parameters(self, async_client):
760+
with patch.object(async_client, "search") as search_mock, patch.object(
761+
async_client, "scroll"
762+
) as scroll_mock, patch.object(
763+
async_client, "clear_scroll"
764+
) as clear_scroll_mock:
765+
search_mock.return_value = MockResponse(
766+
{
767+
"_scroll_id": "scroll_id",
768+
"_shards": {"successful": 5, "total": 5, "skipped": 0},
769+
"hits": {"hits": [{"field": "value"}]},
770+
}
771+
)
772+
scroll_mock.return_value = MockResponse(
773+
{
774+
"_scroll_id": "scroll_id",
775+
"_shards": {"successful": 5, "total": 5, "skipped": 0},
776+
"hits": {"hits": []},
777+
}
778+
)
779+
clear_scroll_mock.return_value = MockResponse({"acknowledged": True})
780+
data = [
781+
x
782+
async for x in helpers.async_scan(
783+
async_client,
784+
index="test_index",
785+
size=10,
786+
query={"size": 1},
787+
scroll_kwargs={"scroll": "10m", "rest_total_hits_as_int": True},
788+
)
789+
]
790+
791+
assert data == [{"field": "value"}]
792+
search_mock.assert_called_with(
793+
index="test_index",
794+
size=10,
795+
sort="_doc",
796+
scroll="5m",
797+
request_timeout=None,
798+
params={"__elastic_client_meta": (("h", "s"),)},
799+
)
800+
scroll_mock.assert_called_with(
801+
scroll="5m",
802+
rest_total_hits_as_int=True,
803+
params={"__elastic_client_meta": (("h", "s"),)},
804+
scroll_id="scroll_id",
805+
)
806+
clear_scroll_mock.assert_called_with(
807+
scroll_id="scroll_id",
808+
ignore=(404,),
809+
params={"__elastic_client_meta": (("h", "s"),)},
810+
)
811+
759812

760813
@pytest.fixture(scope="function")
761814
async def reindex_setup(async_client):

test_elasticsearch/test_server/test_helpers.py

Lines changed: 51 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -568,6 +568,57 @@ def test_scan_auth_kwargs_favor_scroll_kwargs_option(self):
568568
)
569569
self.assertEqual(client_mock.scroll.call_args[1]["sort"], "asc")
570570

571+
def test_scan_duplicate_parameters(self):
572+
with patch.object(self.client, "search") as search_mock, patch.object(
573+
self.client, "scroll"
574+
) as scroll_mock, patch.object(
575+
self.client, "clear_scroll"
576+
) as clear_scroll_mock:
577+
search_mock.return_value = {
578+
"_scroll_id": "scroll_id",
579+
"_shards": {"successful": 5, "total": 5, "skipped": 0},
580+
"hits": {"hits": [{"field": "value"}]},
581+
}
582+
583+
scroll_mock.return_value = {
584+
"_scroll_id": "scroll_id",
585+
"_shards": {"successful": 5, "total": 5, "skipped": 0},
586+
"hits": {"hits": []},
587+
}
588+
589+
clear_scroll_mock.return_value = {"acknowledged": True}
590+
data = [
591+
x
592+
for x in helpers.scan(
593+
self.client,
594+
index="test_index",
595+
size=10,
596+
query={"size": 1},
597+
scroll_kwargs={"scroll": "10m", "rest_total_hits_as_int": True},
598+
)
599+
]
600+
601+
assert data == [{"field": "value"}]
602+
search_mock.assert_called_with(
603+
index="test_index",
604+
size=10,
605+
sort="_doc",
606+
scroll="5m",
607+
request_timeout=None,
608+
params={"__elastic_client_meta": (("h", "s"),)},
609+
)
610+
scroll_mock.assert_called_with(
611+
scroll="5m",
612+
rest_total_hits_as_int=True,
613+
params={"__elastic_client_meta": (("h", "s"),)},
614+
scroll_id="scroll_id",
615+
)
616+
clear_scroll_mock.assert_called_with(
617+
scroll_id="scroll_id",
618+
ignore=(404,),
619+
params={"__elastic_client_meta": (("h", "s"),)},
620+
)
621+
571622
@patch("elasticsearch.helpers.actions.logger")
572623
def test_logger(self, logger_mock):
573624
bulk = []

0 commit comments

Comments
 (0)