REST: Send empty first-page pageToken when rest-page-size is set

pyiceberg/catalog/rest/__init__.py

@@ -1119,7 +1119,7 @@ def add_headers(self, request: PreparedRequest, **kwargs: Any) -> None:  # pylin
 
                 url = str(request.url).split("?")[0]
                 query = str(parse.urlsplit(request.url).query)
-                params = dict(parse.parse_qsl(query))
+                params = dict(parse.parse_qsl(query, keep_blank_values=True))
 
                 # remove the connection header as it will be updated after signing
                 if "connection" in request.headers:
@@ -1410,10 +1410,11 @@ def list_tables(self, namespace: str | Identifier) -> list[Identifier]:
             params["pageSize"] = str(page_size)
 
         tables: list[Identifier] = []
-        page_token: str | None = None
+        # An empty pageToken opts in to pagination on the first request.
+        page_token: str | None = "" if page_size is not None else None
 
         while True:
-            if page_token:
+            if page_token is not None:
                 params["pageToken"] = page_token
             response = self._session.get(url, params=params)
             try:
@@ -1550,10 +1551,11 @@ def list_views(self, namespace: str | Identifier) -> list[Identifier]:
             params["pageSize"] = str(page_size)
 
         views: list[Identifier] = []
-        page_token: str | None = None
+        # An empty pageToken opts in to pagination on the first request.
+        page_token: str | None = "" if page_size is not None else None
 
         while True:
-            if page_token:
+            if page_token is not None:
                 params["pageToken"] = page_token
 
             response = self._session.get(url, params=params)
@@ -1672,12 +1674,13 @@ def list_namespaces(self, namespace: str | Identifier = ()) -> list[Identifier]:
             params["pageSize"] = str(page_size)
 
         namespaces: list[Identifier] = []
-        page_token: str | None = None
+        # An empty pageToken opts in to pagination on the first request.
+        page_token: str | None = "" if page_size is not None else None
 
         while True:
             if namespace_tuple:
                 params["parent"] = self._encode_namespace_path(namespace_tuple)
-            if page_token:
+            if page_token is not None:
                 params["pageToken"] = page_token
             response = self._session.get(self.url(Endpoints.list_namespaces), params=params)
 

tests/catalog/test_rest.py

@@ -579,14 +579,44 @@ def test_list_tables_page_size(rest_mock: Mocker) -> None:
     )
 
     result = RestCatalog("rest", uri=TEST_URI, token=TEST_TOKEN, **{PAGE_SIZE: "100"}).list_tables(namespace)
-    assert rest_mock.last_request.url == f"{TEST_URI}v1/namespaces/examples/tables?pageSize=100"
+    assert rest_mock.last_request.url == f"{TEST_URI}v1/namespaces/examples/tables?pageSize=100&pageToken="
 
     assert result == [
         ("examples", "table1"),
         ("examples", "table2"),
     ]
 
 
+def test_list_tables_page_size_paginated(rest_mock: Mocker) -> None:
+    namespace = "examples"
+    rest_mock.get(
+        f"{TEST_URI}v1/namespaces/{namespace}/tables?pageSize=1&pageToken=",
+        complete_qs=True,
+        json={
+            "identifiers": [{"namespace": ["examples"], "name": "table1"}],
+            "next-page-token": "page2token",
+        },
+        status_code=200,
+        request_headers=TEST_HEADERS,
+    )
+    rest_mock.get(
+        f"{TEST_URI}v1/namespaces/{namespace}/tables?pageSize=1&pageToken=page2token",
+        complete_qs=True,
+        json={
+            "identifiers": [{"namespace": ["examples"], "name": "table2"}],
+            "next-page-token": None,
+        },
+        status_code=200,
+        request_headers=TEST_HEADERS,
+    )
+
+    result = RestCatalog("rest", uri=TEST_URI, token=TEST_TOKEN, **{PAGE_SIZE: "1"}).list_tables(namespace)
+    assert result == [
+        ("examples", "table1"),
+        ("examples", "table2"),
+    ]
+
+
 def test_list_tables_200_sigv4(rest_mock: Mocker) -> None:
     namespace = "examples"
     rest_mock.get(
@@ -658,6 +688,33 @@ def test_sigv4_sign_request_with_body(rest_mock: Mocker) -> None:
     assert prepared.headers.get("x-amz-content-sha256") != EMPTY_BODY_SHA256
 
 
+def test_sigv4_sign_request_keeps_blank_query_params(rest_mock: Mocker) -> None:
+    from botocore.auth import SigV4Auth
+
+    catalog = RestCatalog(
+        "rest",
+        **{
+            "uri": TEST_URI,
+            "token": TEST_TOKEN,
+            "rest.sigv4-enabled": "true",
+            "rest.signing-region": "us-west-2",
+            "client.access-key-id": "id",
+            "client.secret-access-key": "secret",
+        },
+    )
+
+    prepared = catalog._session.prepare_request(
+        Request("GET", f"{TEST_URI}v1/namespaces", params={"pageSize": "1", "pageToken": ""})
+    )
+    adapter = catalog._session.adapters[catalog.uri]
+    assert isinstance(adapter, HTTPAdapter)
+    with mock.patch.object(SigV4Auth, "add_auth", autospec=True) as add_auth:
+        adapter.add_headers(prepared)
+
+    # The signed query must match the one sent, including the empty pageToken
+    assert add_auth.call_args.args[1].params == {"pageSize": "1", "pageToken": ""}
+
+
 def test_sigv4_adapter_default_retry_config(rest_mock: Mocker) -> None:
     catalog = RestCatalog(
         "rest",
@@ -848,7 +905,7 @@ def test_list_views_page_size(rest_mock: Mocker) -> None:
     )
 
     result = RestCatalog("rest", uri=TEST_URI, token=TEST_TOKEN, **{PAGE_SIZE: "100"}).list_views(namespace)
-    assert rest_mock.last_request.url == f"{TEST_URI}v1/namespaces/examples/views?pageSize=100"
+    assert rest_mock.last_request.url == f"{TEST_URI}v1/namespaces/examples/views?pageSize=100&pageToken="
 
     assert result == [
         ("examples", "view1"),
@@ -1082,7 +1139,7 @@ def test_list_namespaces_page_size(rest_mock: Mocker) -> None:
     )
 
     result = RestCatalog("rest", uri=TEST_URI, token=TEST_TOKEN, **{PAGE_SIZE: "100"}).list_namespaces()
-    assert rest_mock.last_request.url == f"{TEST_URI}v1/namespaces?pageSize=100"
+    assert rest_mock.last_request.url == f"{TEST_URI}v1/namespaces?pageSize=100&pageToken="
 
     assert result == [
         ("ns1",),