Rewrite the 'from' parameter from kwargs/query into 'from_'

elastic · Feb 3, 2022 · e049ebb · e049ebb
1 parent 88422cb
commit e049ebb
Show file tree

Hide file tree

Showing 4 changed files with 83 additions and 0 deletions.
diff --git a/elasticsearch/_async/helpers.py b/elasticsearch/_async/helpers.py
@@ -428,8 +428,16 @@ def pop_transport_kwargs(kw: MutableMapping[str, Any]) -> MutableMapping[str, An
     )
     client._client_meta = (("h", "s"),)
 
+    # Setting query={"from": ...} would make 'from' be used
+    # as a keyword argument instead of 'from_'. We handle that here.
+    def normalize_from_keyword(kw: MutableMapping[str, Any]) -> None:
+        if "from" in kw:
+            kw["from_"] = kw.pop("from")
+
+    normalize_from_keyword(kwargs)
     try:
         search_kwargs = query.copy() if query else {}
+        normalize_from_keyword(search_kwargs)
         search_kwargs.update(kwargs)
         search_kwargs["scroll"] = scroll
         search_kwargs["size"] = size

diff --git a/elasticsearch/helpers/actions.py b/elasticsearch/helpers/actions.py
@@ -690,8 +690,16 @@ def pop_transport_kwargs(kw: MutableMapping[str, Any]) -> Dict[str, Any]:
     )
     client._client_meta = (("h", "s"),)
 
+    # Setting query={"from": ...} would make 'from' be used
+    # as a keyword argument instead of 'from_'. We handle that here.
+    def normalize_from_keyword(kw: MutableMapping[str, Any]) -> None:
+        if "from" in kw:
+            kw["from_"] = kw.pop("from")
+
+    normalize_from_keyword(kwargs)
     try:
         search_kwargs = query.copy() if query else {}
+        normalize_from_keyword(search_kwargs)
         search_kwargs.update(kwargs)
         search_kwargs["scroll"] = scroll
         search_kwargs["size"] = size

diff --git a/test_elasticsearch/test_async/test_server/test_helpers.py b/test_elasticsearch/test_async/test_server/test_helpers.py
@@ -831,6 +831,44 @@ async def test_scan_auth_kwargs_favor_scroll_kwargs_option(
                     ]
 
 
+@pytest.mark.parametrize(
+    "scan_kwargs",
+    [
+        {"from": 1},
+        {"from_": 1},
+        {"query": {"from": 1}},
+        {"query": {"from_": 1}},
+        {"query": {"query": {"match_all": {}}}, "from": 1},
+        {"query": {"query": {"match_all": {}}}, "from_": 1},
+    ],
+)
+async def test_scan_from_keyword_is_aliased(async_client, scan_kwargs):
+    with patch.object(async_client, "options", return_value=async_client), patch.object(
+        async_client,
+        "search",
+        return_value=MockResponse(
+            ObjectApiResponse(
+                body={
+                    "_scroll_id": "dummy_id",
+                    "_shards": {"successful": 5, "total": 5},
+                    "hits": {"hits": []},
+                },
+                meta=None,
+            )
+        ),
+    ) as search_mock, patch.object(
+        async_client, "clear_scroll", return_value=MockResponse(None)
+    ):
+        [
+            x
+            async for x in helpers.async_scan(
+                async_client, index="test_index", **scan_kwargs
+            )
+        ]
+        assert search_mock.call_args[1]["from_"] == 1
+        assert "from" not in search_mock.call_args[1]
+
+
 @pytest.fixture(scope="function")
 async def reindex_setup(async_client):
     bulk = []

diff --git a/test_elasticsearch/test_server/test_helpers.py b/test_elasticsearch/test_server/test_helpers.py
@@ -799,6 +799,35 @@ def test_shards_no_skipped_field(sync_client):
         assert data == [{"search_data": 1}, {"scroll_data": 42}]
 
 
+@pytest.mark.parametrize(
+    "scan_kwargs",
+    [
+        {"from": 1},
+        {"from_": 1},
+        {"query": {"from": 1}},
+        {"query": {"from_": 1}},
+        {"query": {"query": {"match_all": {}}}, "from": 1},
+        {"query": {"query": {"match_all": {}}}, "from_": 1},
+    ],
+)
+def test_scan_from_keyword_is_aliased(sync_client, scan_kwargs):
+    with patch.object(sync_client, "options", return_value=sync_client), patch.object(
+        sync_client,
+        "search",
+        return_value=ObjectApiResponse(
+            raw={
+                "_scroll_id": "dummy_id",
+                "_shards": {"successful": 5, "total": 5},
+                "hits": {"hits": []},
+            },
+            meta=None,
+        ),
+    ) as search_mock, patch.object(sync_client, "clear_scroll"):
+        list(helpers.scan(sync_client, index="test_index", **scan_kwargs))
+        assert search_mock.call_args[1]["from_"] == 1
+        assert "from" not in search_mock.call_args[1]
+
+
 @pytest.fixture(scope="function")
 def reindex_setup(sync_client):
     bulk = []