Fix: implement lazy pagination iterator for catalog list methods (#3365)

Co-authored-by: Yuya Ebihara <ebyhry@gmail.com>

Author: 2 people

pyiceberg/catalog/__init__.py

@@ -25,6 +25,7 @@
 from collections.abc import Callable
 from dataclasses import dataclass
 from enum import Enum
+from collections.abc import Iterator
 from typing import (
     TYPE_CHECKING,
     Any,
@@ -607,42 +608,42 @@ def drop_namespace(self, namespace: str | Identifier) -> None:
         """
 
     @abstractmethod
-    def list_tables(self, namespace: str | Identifier) -> list[Identifier]:
+    def list_tables(self, namespace: str | Identifier) -> Iterator[Identifier]:
         """List tables under the given namespace in the catalog.
 
         Args:
             namespace (str | Identifier): Namespace identifier to search.
 
         Returns:
-            List[Identifier]: list of table identifiers.
+            Iterator[Identifier]: an iterator of table identifiers.
 
         Raises:
             NoSuchNamespaceError: If a namespace with the given name does not exist.
         """
 
     @abstractmethod
-    def list_namespaces(self, namespace: str | Identifier = ()) -> list[Identifier]:
+    def list_namespaces(self, namespace: str | Identifier = ()) -> Iterator[Identifier]:
         """List namespaces from the given namespace. If not given, list top-level namespaces from the catalog.
 
         Args:
             namespace (str | Identifier): Namespace identifier to search.
 
         Returns:
-            List[Identifier]: a List of namespace identifiers.
+            Iterator[Identifier]: an iterator of namespace identifiers.
 
         Raises:
             NoSuchNamespaceError: If a namespace with the given name does not exist.
         """
 
     @abstractmethod
-    def list_views(self, namespace: str | Identifier) -> list[Identifier]:
+    def list_views(self, namespace: str | Identifier) -> Iterator[Identifier]:
         """List views under the given namespace in the catalog.
 
         Args:
             namespace (str | Identifier): Namespace identifier to search.
 
         Returns:
-            List[Identifier]: list of table identifiers.
+            Iterator[Identifier]: an iterator of view identifiers.
 
         Raises:
             NoSuchNamespaceError: If a namespace with the given name does not exist.

pyiceberg/catalog/bigquery_metastore.py

@@ -17,6 +17,7 @@
 from __future__ import annotations
 
 import json
+from collections.abc import Iterator
 from typing import TYPE_CHECKING, Any
 
 from google.api_core.exceptions import NotFound
@@ -252,7 +253,7 @@ def drop_namespace(self, namespace: str | Identifier) -> None:
             raise NoSuchNamespaceError(f"Namespace {namespace} does not exist.") from e
 
     @override
-    def list_tables(self, namespace: str | Identifier) -> list[Identifier]:
+    def list_tables(self, namespace: str | Identifier) -> Iterator[Identifier]:
         database_name = self.identifier_to_database(namespace)
         iceberg_tables: list[Identifier] = []
         try:
@@ -264,18 +265,18 @@ def list_tables(self, namespace: str | Identifier) -> list[Identifier]:
                 iceberg_tables.append((database_name, bq_table_list_item.table_id))
         except NotFound:
             raise NoSuchNamespaceError(f"Namespace (dataset) '{database_name}' not found.") from None
-        return iceberg_tables
+        return iter(iceberg_tables)
 
     @override
-    def list_namespaces(self, namespace: str | Identifier = ()) -> list[Identifier]:
+    def list_namespaces(self, namespace: str | Identifier = ()) -> Iterator[Identifier]:
         # Since this catalog only supports one-level namespaces, it always returns an empty list unless
         # passed an empty namespace to list all namespaces within the catalog.
         if namespace:
             raise NoSuchNamespaceError(f"Namespace (dataset) '{namespace}' not found.") from None
 
         # List top-level datasets
         datasets_iterator = self.client.list_datasets()
-        return [(dataset.dataset_id,) for dataset in datasets_iterator]
+        return iter([(dataset.dataset_id,) for dataset in datasets_iterator])
 
     @override
     def register_table(self, identifier: str | Identifier, metadata_location: str, overwrite: bool = False) -> Table:
@@ -314,7 +315,7 @@ def register_table(self, identifier: str | Identifier, metadata_location: str, o
         return self.load_table(identifier=identifier)
 
     @override
-    def list_views(self, namespace: str | Identifier) -> list[Identifier]:
+    def list_views(self, namespace: str | Identifier) -> Iterator[Identifier]:
         raise NotImplementedError
 
     @override

pyiceberg/catalog/dynamodb.py

@@ -15,6 +15,7 @@
 #  specific language governing permissions and limitations
 #  under the License.
 import uuid
+from collections.abc import Iterator
 from time import time
 from typing import (
     TYPE_CHECKING,
@@ -396,8 +397,11 @@ def drop_namespace(self, namespace: str | Identifier) -> None:
         database_name = self.identifier_to_database(namespace, NoSuchNamespaceError)
         table_identifiers = self.list_tables(namespace=database_name)
 
-        if len(table_identifiers) > 0:
+        try:
+            next(table_identifiers)
             raise NamespaceNotEmptyError(f"Database {database_name} is not empty")
+        except StopIteration:
+            pass
 
         try:
             self._delete_dynamo_item(
@@ -409,14 +413,14 @@ def drop_namespace(self, namespace: str | Identifier) -> None:
             raise NoSuchNamespaceError(f"Database does not exist: {database_name}") from e
 
     @override
-    def list_tables(self, namespace: str | Identifier) -> list[Identifier]:
+    def list_tables(self, namespace: str | Identifier) -> Iterator[Identifier]:
         """List Iceberg tables under the given namespace in the catalog.
 
         Args:
             namespace (str | Identifier): Namespace identifier to search.
 
         Returns:
-            List[Identifier]: list of table identifiers.
+            Iterator[Identifier]: an iterator of table identifiers.
         """
         database_name = self.identifier_to_database(namespace, NoSuchNamespaceError)
 
@@ -451,20 +455,20 @@ def list_tables(self, namespace: str | Identifier) -> list[Identifier]:
 
                 table_identifiers.append(self.identifier_to_tuple(identifier_col))
 
-        return table_identifiers
+        return iter(table_identifiers)
 
     @override
-    def list_namespaces(self, namespace: str | Identifier = ()) -> list[Identifier]:
+    def list_namespaces(self, namespace: str | Identifier = ()) -> Iterator[Identifier]:
         """List top-level namespaces from the catalog.
 
         We do not support hierarchical namespace.
 
         Returns:
-            List[Identifier]: a List of namespace identifiers.
+            Iterator[Identifier]: an iterator of namespace identifiers.
         """
         # Hierarchical namespace is not supported. Return an empty list
         if namespace:
-            return []
+            return iter([])
 
         paginator = self.dynamodb.get_paginator("query")
 
@@ -494,7 +498,7 @@ def list_namespaces(self, namespace: str | Identifier = ()) -> list[Identifier]:
                 namespace_col = _dict[DYNAMODB_COL_NAMESPACE]
                 database_identifiers.append(self.identifier_to_tuple(namespace_col))
 
-        return database_identifiers
+        return iter(database_identifiers)
 
     @override
     def load_namespace_properties(self, namespace: str | Identifier) -> Properties:
@@ -565,7 +569,7 @@ def create_view(
         raise NotImplementedError
 
     @override
-    def list_views(self, namespace: str | Identifier) -> list[Identifier]:
+    def list_views(self, namespace: str | Identifier) -> Iterator[Identifier]:
         raise NotImplementedError
 
     @override

pyiceberg/catalog/glue.py

@@ -17,6 +17,7 @@
 
 
 import logging
+from collections.abc import Iterator
 from typing import (
     TYPE_CHECKING,
     Any,
@@ -860,14 +861,14 @@ def drop_namespace(self, namespace: str | Identifier) -> None:
         self.glue.delete_database(Name=database_name)
 
     @override
-    def list_tables(self, namespace: str | Identifier) -> list[Identifier]:
+    def list_tables(self, namespace: str | Identifier) -> Iterator[Identifier]:
         """List Iceberg tables under the given namespace in the catalog.
 
         Args:
             namespace (str | Identifier): Namespace identifier to search.
 
         Returns:
-            List[Identifier]: list of table identifiers.
+            Iterator[Identifier]: an iterator of table identifiers.
 
         Raises:
             NoSuchNamespaceError: If a namespace with the given name does not exist, or the identifier is invalid.
@@ -889,18 +890,18 @@ def list_tables(self, namespace: str | Identifier) -> list[Identifier]:
 
         except self.glue.exceptions.EntityNotFoundException as e:
             raise NoSuchNamespaceError(f"Database does not exist: {database_name}") from e
-        return [(database_name, table["Name"]) for table in table_list if self.__is_iceberg_table(table)]
+        return iter([(database_name, table["Name"]) for table in table_list if self.__is_iceberg_table(table)])
 
     @override
-    def list_namespaces(self, namespace: str | Identifier = ()) -> list[Identifier]:
+    def list_namespaces(self, namespace: str | Identifier = ()) -> Iterator[Identifier]:
         """List namespaces from the given namespace. If not given, list top-level namespaces from the catalog.
 
         Returns:
-            List[Identifier]: a List of namespace identifiers.
+            Iterator[Identifier]: an iterator of namespace identifiers.
         """
         # Hierarchical namespace is not supported. Return an empty list
         if namespace:
-            return []
+            return iter([])
 
         database_list: list[DatabaseTypeDef] = []
         next_token: str | None = None
@@ -912,7 +913,7 @@ def list_namespaces(self, namespace: str | Identifier = ()) -> list[Identifier]:
             if not next_token:
                 break
 
-        return [self.identifier_to_tuple(database["Name"]) for database in database_list]
+        return iter([self.identifier_to_tuple(database["Name"]) for database in database_list])
 
     @override
     def load_namespace_properties(self, namespace: str | Identifier) -> Properties:
@@ -982,7 +983,7 @@ def create_view(
         raise NotImplementedError
 
     @override
-    def list_views(self, namespace: str | Identifier) -> list[Identifier]:
+    def list_views(self, namespace: str | Identifier) -> Iterator[Identifier]:
         raise NotImplementedError
 
     @override

pyiceberg/catalog/hive.py

@@ -18,6 +18,7 @@
 import logging
 import socket
 import time
+from collections.abc import Iterator
 from types import TracebackType
 from typing import (
     TYPE_CHECKING,
@@ -479,7 +480,7 @@ def register_table(self, identifier: str | Identifier, metadata_location: str, o
         return self._convert_hive_into_iceberg(hive_table)
 
     @override
-    def list_views(self, namespace: str | Identifier) -> list[Identifier]:
+    def list_views(self, namespace: str | Identifier) -> Iterator[Identifier]:
         raise NotImplementedError
 
     @override
@@ -760,7 +761,7 @@ def drop_namespace(self, namespace: str | Identifier) -> None:
             raise NoSuchNamespaceError(f"Database does not exists: {database_name}") from e
 
     @override
-    def list_tables(self, namespace: str | Identifier) -> list[Identifier]:
+    def list_tables(self, namespace: str | Identifier) -> Iterator[Identifier]:
         """List Iceberg tables under the given namespace in the catalog.
 
         When the database doesn't exist, it will just return an empty list.
@@ -769,34 +770,34 @@ def list_tables(self, namespace: str | Identifier) -> list[Identifier]:
             namespace: Database to list.
 
         Returns:
-            List[Identifier]: list of table identifiers.
+            Iterator[Identifier]: an iterator of table identifiers.
 
         Raises:
             NoSuchNamespaceError: If a namespace with the given name does not exist, or the identifier is invalid.
         """
         database_name = self.identifier_to_database(namespace, NoSuchNamespaceError)
         with self._client as open_client:
-            return [
+            return iter([
                 (database_name, table.tableName)
                 for table in open_client.get_table_objects_by_name(
                     dbname=database_name, tbl_names=open_client.get_all_tables(db_name=database_name)
                 )
                 if table.parameters.get(TABLE_TYPE, "").lower() == ICEBERG
-            ]
+            ])
 
     @override
-    def list_namespaces(self, namespace: str | Identifier = ()) -> list[Identifier]:
+    def list_namespaces(self, namespace: str | Identifier = ()) -> Iterator[Identifier]:
         """List namespaces from the given namespace. If not given, list top-level namespaces from the catalog.
 
         Returns:
-            List[Identifier]: a List of namespace identifiers.
+            Iterator[Identifier]: an iterator of namespace identifiers.
         """
         # Hierarchical namespace is not supported. Return an empty list
         if namespace:
-            return []
+            return iter([])
 
         with self._client as open_client:
-            return list(map(self.identifier_to_tuple, open_client.get_all_databases()))
+            return iter(list(map(self.identifier_to_tuple, open_client.get_all_databases())))
 
     @override
     def load_namespace_properties(self, namespace: str | Identifier) -> Properties:

pyiceberg/catalog/noop.py

@@ -16,6 +16,7 @@
 #  under the License.
 from __future__ import annotations
 
+from collections.abc import Iterator
 from typing import (
     TYPE_CHECKING,
 )
@@ -124,11 +125,11 @@ def drop_namespace(self, namespace: str | Identifier) -> None:
         raise NotImplementedError
 
     @override
-    def list_tables(self, namespace: str | Identifier) -> list[Identifier]:
+    def list_tables(self, namespace: str | Identifier) -> Iterator[Identifier]:
         raise NotImplementedError
 
     @override
-    def list_namespaces(self, namespace: str | Identifier = ()) -> list[Identifier]:
+    def list_namespaces(self, namespace: str | Identifier = ()) -> Iterator[Identifier]:
         raise NotImplementedError
 
     @override
@@ -142,7 +143,7 @@ def update_namespace_properties(
         raise NotImplementedError
 
     @override
-    def list_views(self, namespace: str | Identifier) -> list[Identifier]:
+    def list_views(self, namespace: str | Identifier) -> Iterator[Identifier]:
         raise NotImplementedError
 
     @override