[7.x] Forward auth parameters to APIs in scan helpers

Co-authored-by: Seth Michael Larson <seth.larson@elastic.co>

Author: github-actions[bot]

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"),)},
             )

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
             )
 
 

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 @@ def __init__(self, resp):
     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 @@ async def test_clear_scroll(self, async_client, scan_teardown):
             ]
             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):

test_elasticsearch/test_server/test_helpers.py

@@ -478,6 +478,69 @@ def test_no_scroll_id_fast_route(self):
             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 = []