diff --git a/elasticsearch/_async/helpers.py b/elasticsearch/_async/helpers.py index 9dfc13c3..738ffc9a 100644 --- a/elasticsearch/_async/helpers.py +++ b/elasticsearch/_async/helpers.py @@ -338,6 +338,20 @@ async def async_scan( query = query.copy() if query else {} query["sort"] = "_doc" + # Grab options that should be propagated to every + # API call within this helper instead of just 'search()' + transport_kwargs = {} + for key in ("headers", "api_key", "http_auth"): + if key in kwargs: + transport_kwargs[key] = kwargs[key] + + # If the user is using 'scroll_kwargs' we want + # to propagate there too, but to not break backwards + # compatibility we'll not override anything already given. + if scroll_kwargs is not None and transport_kwargs: + for key, val in transport_kwargs.items(): + scroll_kwargs.setdefault(key, val) + # initial search resp = await client.search( body=query, scroll=scroll, size=size, request_timeout=request_timeout, **kwargs @@ -382,6 +396,7 @@ async def async_scan( if scroll_id and clear_scroll: await client.clear_scroll( body={"scroll_id": [scroll_id]}, + **transport_kwargs, ignore=(404,), params={"__elastic_client_meta": (("h", "s"),)}, ) diff --git a/elasticsearch/helpers/actions.py b/elasticsearch/helpers/actions.py index a8c7712c..eaf535ec 100644 --- a/elasticsearch/helpers/actions.py +++ b/elasticsearch/helpers/actions.py @@ -550,6 +550,20 @@ def scan( query = query.copy() if query else {} query["sort"] = "_doc" + # Grab options that should be propagated to every + # API call within this helper instead of just 'search()' + transport_kwargs = {} + for key in ("headers", "api_key", "http_auth"): + if key in kwargs: + transport_kwargs[key] = kwargs[key] + + # If the user is using 'scroll_kwargs' we want + # to propagate there too, but to not break backwards + # compatibility we'll not override anything already given. + if scroll_kwargs is not None and transport_kwargs: + for key, val in transport_kwargs.items(): + scroll_kwargs.setdefault(key, val) + # initial search resp = client.search( body=query, scroll=scroll, size=size, request_timeout=request_timeout, **kwargs @@ -596,6 +610,7 @@ def scan( body={"scroll_id": [scroll_id]}, ignore=(404,), params={"__elastic_client_meta": (("h", "s"),)}, + **transport_kwargs ) diff --git a/test_elasticsearch/test_async/test_server/test_helpers.py b/test_elasticsearch/test_async/test_server/test_helpers.py index 46927b9c..dc25b5dd 100644 --- a/test_elasticsearch/test_async/test_server/test_helpers.py +++ b/test_elasticsearch/test_async/test_server/test_helpers.py @@ -34,6 +34,9 @@ class AsyncMock(MagicMock): async def __call__(self, *args, **kwargs): return super(AsyncMock, self).__call__(*args, **kwargs) + def __await__(self): + return self().__await__() + class FailingBulkClient(object): def __init__( @@ -419,6 +422,9 @@ class MockResponse: async def __call__(self, *args, **kwargs): return self.resp + def __await__(self): + return self().__await__() + @pytest.fixture(scope="function") async def scan_teardown(async_client): @@ -645,6 +651,105 @@ class TestScan(object): ] spy.assert_not_called() + @pytest.mark.parametrize( + "kwargs", + [ + {"api_key": ("name", "value")}, + {"http_auth": ("username", "password")}, + {"headers": {"custom", "header"}}, + ], + ) + async def test_scan_auth_kwargs_forwarded( + self, async_client, scan_teardown, kwargs + ): + ((key, val),) = kwargs.items() + + with patch.object( + async_client, + "search", + return_value=MockResponse( + { + "_scroll_id": "scroll_id", + "_shards": {"successful": 5, "total": 5, "skipped": 0}, + "hits": {"hits": [{"search_data": 1}]}, + } + ), + ) as search_mock: + with patch.object( + async_client, + "scroll", + return_value=MockResponse( + { + "_scroll_id": "scroll_id", + "_shards": {"successful": 5, "total": 5, "skipped": 0}, + "hits": {"hits": []}, + } + ), + ) as scroll_mock: + with patch.object( + async_client, "clear_scroll", return_value=MockResponse({}) + ) as clear_mock: + data = [ + x + async for x in helpers.async_scan( + async_client, index="test_index", **kwargs + ) + ] + + assert data == [{"search_data": 1}] + + for api_mock in (search_mock, scroll_mock, clear_mock): + assert api_mock.call_args[1][key] == val + + async def test_scan_auth_kwargs_favor_scroll_kwargs_option( + self, async_client, scan_teardown + ): + with patch.object( + async_client, + "search", + return_value=MockResponse( + { + "_scroll_id": "scroll_id", + "_shards": {"successful": 5, "total": 5, "skipped": 0}, + "hits": {"hits": [{"search_data": 1}]}, + } + ), + ): + with patch.object( + async_client, + "scroll", + return_value=MockResponse( + { + "_scroll_id": "scroll_id", + "_shards": {"successful": 5, "total": 5, "skipped": 0}, + "hits": {"hits": []}, + } + ), + ): + with patch.object( + async_client, "clear_scroll", return_value=MockResponse({}) + ): + data = [ + x + async for x in helpers.async_scan( + async_client, + index="test_index", + headers={"not scroll": "kwargs"}, + scroll_kwargs={ + "headers": {"scroll": "kwargs"}, + "sort": "asc", + }, + ) + ] + + assert data == [{"search_data": 1}] + + # Assert that we see 'scroll_kwargs' options used instead of 'kwargs' + assert async_client.scroll.call_args[1]["headers"] == { + "scroll": "kwargs" + } + assert async_client.scroll.call_args[1]["sort"] == "asc" + @pytest.fixture(scope="function") async def reindex_setup(async_client): diff --git a/test_elasticsearch/test_server/test_helpers.py b/test_elasticsearch/test_server/test_helpers.py index d9451062..8c4ac8ac 100644 --- a/test_elasticsearch/test_server/test_helpers.py +++ b/test_elasticsearch/test_server/test_helpers.py @@ -478,6 +478,69 @@ class TestScan(ElasticsearchTestCase): client_mock.scroll.assert_not_called() client_mock.clear_scroll.assert_not_called() + def test_scan_auth_kwargs_forwarded(self): + for key, val in { + "api_key": ("name", "value"), + "http_auth": ("username", "password"), + "headers": {"custom": "header"}, + }.items(): + with patch.object(self, "client") as client_mock: + client_mock.search.return_value = { + "_scroll_id": "scroll_id", + "_shards": {"successful": 5, "total": 5, "skipped": 0}, + "hits": {"hits": [{"search_data": 1}]}, + } + client_mock.scroll.return_value = { + "_scroll_id": "scroll_id", + "_shards": {"successful": 5, "total": 5, "skipped": 0}, + "hits": {"hits": []}, + } + client_mock.clear_scroll.return_value = {} + + data = list(helpers.scan(self.client, index="test_index", **{key: val})) + + self.assertEqual(data, [{"search_data": 1}]) + + # Assert that 'search', 'scroll' and 'clear_scroll' all + # received the extra kwarg related to authentication. + for api_mock in ( + client_mock.search, + client_mock.scroll, + client_mock.clear_scroll, + ): + self.assertEqual(api_mock.call_args[1][key], val) + + def test_scan_auth_kwargs_favor_scroll_kwargs_option(self): + with patch.object(self, "client") as client_mock: + client_mock.search.return_value = { + "_scroll_id": "scroll_id", + "_shards": {"successful": 5, "total": 5, "skipped": 0}, + "hits": {"hits": [{"search_data": 1}]}, + } + client_mock.scroll.return_value = { + "_scroll_id": "scroll_id", + "_shards": {"successful": 5, "total": 5, "skipped": 0}, + "hits": {"hits": []}, + } + client_mock.clear_scroll.return_value = {} + + data = list( + helpers.scan( + self.client, + index="test_index", + scroll_kwargs={"headers": {"scroll": "kwargs"}, "sort": "asc"}, + headers={"not scroll": "kwargs"}, + ) + ) + + self.assertEqual(data, [{"search_data": 1}]) + + # Assert that we see 'scroll_kwargs' options used instead of 'kwargs' + self.assertEqual( + client_mock.scroll.call_args[1]["headers"], {"scroll": "kwargs"} + ) + self.assertEqual(client_mock.scroll.call_args[1]["sort"], "asc") + @patch("elasticsearch.helpers.actions.logger") def test_logger(self, logger_mock): bulk = []