Add pagination support for list_namespaces

Author: vishnuprakaz

pyiceberg/catalog/rest/__init__.py

@@ -350,6 +350,7 @@ def _parse_endpoints(cls, v: list[str] | None) -> set[Endpoint] | None:
 
 class ListNamespaceResponse(IcebergBaseModel):
     namespaces: list[Identifier] = Field()
+    next_page_token: str | None = Field(default=None, alias="next-page-token")
 
 
 class NamespaceResponse(IcebergBaseModel):
@@ -1243,19 +1244,31 @@ def drop_namespace(self, namespace: str | Identifier) -> None:
     def list_namespaces(self, namespace: str | Identifier = ()) -> list[Identifier]:
         self._check_endpoint(Capability.V1_LIST_NAMESPACES)
         namespace_tuple = self.identifier_to_tuple(namespace)
-        response = self._session.get(
-            self.url(
-                f"{Endpoints.list_namespaces}?parent={self._encode_namespace_path(namespace_tuple)}"
-                if namespace_tuple
-                else Endpoints.list_namespaces
-            ),
+        url = (
+            self.url(f"{Endpoints.list_namespaces}?parent={self._encode_namespace_path(namespace_tuple)}")
+            if namespace_tuple
+            else self.url(Endpoints.list_namespaces)
         )
-        try:
-            response.raise_for_status()
-        except HTTPError as exc:
-            _handle_non_200_response(exc, {404: NoSuchNamespaceError})
 
-        return ListNamespaceResponse.model_validate_json(response.text).namespaces
+        namespaces: list[Identifier] = []
+        page_token: str | None = None
+
+        while True:
+            params = {"pageToken": page_token} if page_token else None
+            response = self._session.get(url, params=params)
+            try:
+                response.raise_for_status()
+            except HTTPError as exc:
+                _handle_non_200_response(exc, {404: NoSuchNamespaceError})
+
+            parsed = ListNamespaceResponse.model_validate_json(response.text)
+            namespaces.extend(parsed.namespaces)
+
+            if not parsed.next_page_token:
+                break
+            page_token = parsed.next_page_token
+
+        return namespaces
 
     @retry(**_RETRY_ARGS)
     @override

tests/catalog/test_rest.py

@@ -841,6 +841,50 @@ def test_list_namespace_with_parent_404(rest_mock: Mocker) -> None:
         RestCatalog("rest", uri=TEST_URI, token=TEST_TOKEN).list_namespaces(("some_namespace",))
 
 
+def test_list_namespaces_paginated_200(rest_mock: Mocker) -> None:
+    rest_mock.get(
+        f"{TEST_URI}v1/namespaces",
+        json={"namespaces": [["default"], ["examples"]], "next-page-token": "page2token"},
+        status_code=200,
+        request_headers=TEST_HEADERS,
+    )
+    rest_mock.get(
+        f"{TEST_URI}v1/namespaces?pageToken=page2token",
+        json={"namespaces": [["fokko"], ["system"]]},
+        status_code=200,
+        request_headers=TEST_HEADERS,
+    )
+
+    assert RestCatalog("rest", uri=TEST_URI, token=TEST_TOKEN).list_namespaces() == [
+        ("default",),
+        ("examples",),
+        ("fokko",),
+        ("system",),
+    ]
+
+
+def test_list_namespaces_paginated_200_none_next_page_token(rest_mock: Mocker) -> None:
+    rest_mock.get(
+        f"{TEST_URI}v1/namespaces",
+        json={"namespaces": [["default"], ["examples"]], "next-page-token": "page2token"},
+        status_code=200,
+        request_headers=TEST_HEADERS,
+    )
+    rest_mock.get(
+        f"{TEST_URI}v1/namespaces?pageToken=page2token",
+        json={"namespaces": [["fokko"], ["system"]], "next-page-token": None},
+        status_code=200,
+        request_headers=TEST_HEADERS,
+    )
+
+    assert RestCatalog("rest", uri=TEST_URI, token=TEST_TOKEN).list_namespaces() == [
+        ("default",),
+        ("examples",),
+        ("fokko",),
+        ("system",),
+    ]
+
+
 @pytest.mark.filterwarnings(
     "ignore:Deprecated in 0.8.0, will be removed in 1.0.0. "
     "Iceberg REST client is missing the OAuth2 server URI:DeprecationWarning"