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

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

Browse Source 
Co-authored-by: Seth Michael Larson <seth.larson@elastic.co>
This commit is contained in:
github-actions[bot]
2021-07-01 16:56:23 -05:00
committed by GitHub
parent 0bfac34b88
commit 29f6a1ec42
4 changed files with 198 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
elasticsearch/_async/helpers.py
+15
View File
@@ -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
+15
View File
@@ -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
+105
View File
@@ -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):
test_elasticsearch/test_server/test_helpers.py
+63
View File
@@ -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 = []