feat!: list_databases returns an iterator (#4442)

CatalogConnection.list_databases should return an iterator so that users
don't have to manually re-implement pagination. For those who do want
more control over pagination, add num_page_results and page_token to the
Iterator object, and also allow the iterator object to be created at a
specific page token.

python/python/lancedb/__init__.py

@@ -62,7 +62,6 @@ from .namespace import (
 from .catalog import (
     AsyncCatalog,
     Catalog,
-    ListDatabasesResponse,
     connect_catalog,
     connect_catalog_async,
 )
@@ -578,7 +577,6 @@ async def connect_async(
 __all__ = [
     "Catalog",
     "AsyncCatalog",
-    "ListDatabasesResponse",
     "connect_catalog",
     "connect_catalog_async",
     "AsyncMaterializedView",

python/python/lancedb/_lancedb.pyi

@@ -946,9 +946,15 @@ class Catalog:
     async def drop_database(
         self, name: str, *, ignore_missing: bool = False
     ) -> None: ...
-    async def list_databases(
-        self, *, limit: Optional[int] = None, page_token: Optional[str] = None
-    ) -> tuple[list[str], Optional[str]]: ...
+    def list_databases(
+        self, *, page_token: Optional[str] = None, page_limit: Optional[int] = None
+    ) -> DatabaseNames: ...
+
+class DatabaseNames:
+    def num_page_results(self) -> int: ...
+    def page_token(self) -> Optional[str]: ...
+    def __aiter__(self) -> DatabaseNames: ...
+    async def __anext__(self) -> str: ...
 
 async def connect_catalog(
     endpoint: str,

python/python/lancedb/catalog.py

@@ -3,9 +3,8 @@
 
 """Remote catalogs manage databases through a server's root namespace."""
 
-from dataclasses import dataclass
 from datetime import timedelta
-from typing import Any, Optional, Union
+from typing import Any, AsyncIterator, Iterator, Optional, Union
 
 from . import _lancedb
 from .background_loop import LOOP
@@ -14,12 +13,66 @@ from .remote import ClientConfig, OAuthConfig
 from .remote.db import RemoteDBConnection
 
 
-@dataclass
-class ListDatabasesResponse:
-    """A page of database names and an optional continuation token."""
+class AsyncDatabaseNames(AsyncIterator[str]):
+    """An async iterator of database names with pagination state.
 
-    databases: list[str]
-    page_token: Optional[str] = None
+    Returned by [list_databases][lancedb.catalog.AsyncCatalog.list_databases].
+    """
+
+    def __init__(self, inner: _lancedb.DatabaseNames):
+        self._inner = inner
+
+    def __aiter__(self) -> "AsyncDatabaseNames":
+        return self
+
+    async def __anext__(self) -> str:
+        return await self._inner.__anext__()
+
+    def num_page_results(self) -> int:
+        """Return the number of names available without another REST request."""
+        return self._inner.num_page_results()
+
+    def page_token(self) -> Optional[str]:
+        """Return the token for the next REST request.
+
+        Before the first request, this is the supplied starting token (``None``
+        starts at the beginning). After the final page is fetched, it is ``None``,
+        even if names remain cached. Drain the cache before saving a token to
+        avoid skipping those names when resuming. A failed request terminates
+        iteration and leaves its token available for resuming a new iterator.
+        """
+        return self._inner.page_token()
+
+
+class DatabaseNames(Iterator[str]):
+    """A synchronous iterator of database names with pagination state.
+
+    Returned by [Catalog.list_databases][lancedb.catalog.Catalog.list_databases].
+    """
+
+    def __init__(self, inner: AsyncDatabaseNames):
+        self._inner = inner
+
+    def __iter__(self) -> "DatabaseNames":
+        return self
+
+    def __next__(self) -> str:
+        try:
+            return LOOP.run(self._inner.__anext__())
+        except StopAsyncIteration:
+            raise StopIteration from None
+
+    def num_page_results(self) -> int:
+        """Return the number of names available without another REST request."""
+        return self._inner.num_page_results()
+
+    def page_token(self) -> Optional[str]:
+        """Return the token for the next REST request.
+
+        See [page_token][lancedb.catalog.AsyncDatabaseNames.page_token]
+        for initial, cached, final-page, and error behavior.
+        """
+        return self._inner.page_token()
 
 
 class AsyncCatalog:
@@ -50,14 +103,25 @@ class AsyncCatalog:
         """Connect to an existing database by its logical name."""
         return AsyncConnection(await self._inner.connect_database(name))
 
-    async def list_databases(
-        self, *, limit: Optional[int] = None, page_token: Optional[str] = None
-    ) -> ListDatabasesResponse:
-        """List a page of databases. Pass the returned token for the next page."""
-        names, token = await self._inner.list_databases(
-            limit=limit, page_token=page_token
+    def list_databases(
+        self, *, page_token: Optional[str] = None, page_limit: Optional[int] = None
+    ) -> AsyncDatabaseNames:
+        """Iterate lazily over all database names using ``async for``.
+
+        ``page_token`` resumes from a saved token; ``None`` starts at the beginning.
+        ``page_limit`` sets the maximum names per REST response (not the total);
+        ``None`` uses the server default. Request errors are raised during iteration.
+
+        Examples
+        --------
+        ```python
+        async for name in catalog.list_databases():
+            print(name)
+        ```
+        """
+        return AsyncDatabaseNames(
+            self._inner.list_databases(page_token=page_token, page_limit=page_limit)
         )
-        return ListDatabasesResponse(names, token)
 
     async def drop_database(self, name: str, *, ignore_missing: bool = False) -> None:
         """Drop an empty database. A nonempty database is an error."""
@@ -73,7 +137,8 @@ class Catalog:
     ```python
     catalog = lancedb.connect_catalog("https://my-server.example", api_key="secret")
     db = catalog.create_database("analytics", exist_ok=True)
-    page = catalog.list_databases(limit=20)
+    for name in catalog.list_databases():
+        print(name)
     ```
     """
 
@@ -118,10 +183,17 @@ class Catalog:
         )
 
     def list_databases(
-        self, *, limit: Optional[int] = None, page_token: Optional[str] = None
-    ) -> ListDatabasesResponse:
-        """List a page of databases. Pass the returned token for the next page."""
-        return LOOP.run(self._inner.list_databases(limit=limit, page_token=page_token))
+        self, *, page_token: Optional[str] = None, page_limit: Optional[int] = None
+    ) -> DatabaseNames:
+        """Iterate lazily over database names, fetching more results as needed.
+
+        ``page_token`` resumes from a saved token; ``None`` starts at the beginning.
+        ``page_limit`` sets the maximum names per REST response (not the total);
+        ``None`` uses the server default. Request errors are raised during iteration.
+        """
+        return DatabaseNames(
+            self._inner.list_databases(page_token=page_token, page_limit=page_limit)
+        )
 
     def drop_database(self, name: str, *, ignore_missing: bool = False) -> None:
         """Drop an empty database. A nonempty database is an error."""

python/python/tests/test_catalog.py

@@ -55,7 +55,7 @@ def test_catalog_sync_scope_and_serialization(catalog_server):
             (200, {}),
             (200, {"tables": []}),
             (200, {"tables": []}),
-            (200, {"namespaces": ["team/search"], "page_token": "next"}),
+            (200, {"namespaces": ["team/search"]}),
             (204, None),
         ]
     )
@@ -85,13 +85,12 @@ def test_catalog_sync_scope_and_serialization(catalog_server):
         with pytest.raises(ValueError, match="sql_host_override must use"):
             connection.execute_query_async("SELECT 1")
     assert restored.table_names() == []
-    page = catalog.list_databases(limit=1, page_token="a/b")
-    assert page == lancedb.ListDatabasesResponse(["team/search"], "next")
+    assert list(catalog.list_databases()) == ["team/search"]
     catalog.drop_database("team/search", ignore_missing=True)
     assert requests[0][0] == "/v1/namespace/team%2Fsearch/create"
     assert requests[0][2] == {"mode": "ExistOk"}
     assert requests[1][0] == "/v1/namespace/team%2Fsearch/describe"
-    assert requests[4][0] == "/v1/namespace/$/list?limit=1&page_token=a%2Fb"
+    assert requests[4][0] == "/v1/namespace/$/list"
     assert requests[5][2] == {"mode": "Skip", "behavior": "Restrict"}
     for i, (_, headers, _) in enumerate(requests):
         headers = {key.lower(): value for key, value in headers.items()}
@@ -138,3 +137,129 @@ async def test_catalog_async_and_errors(catalog_server):
 def test_catalog_requires_remote_endpoint(endpoint):
     with pytest.raises(ValueError, match="endpoint"):
         lancedb.connect_catalog(endpoint)
+
+
+@pytest.mark.parametrize("asynchronous", [False, True])
+@pytest.mark.asyncio
+@pytest.mark.parametrize("page_token,page_limit", [(None, None), ("start/token", 2)])
+async def test_database_names_iteration(
+    catalog_server, asynchronous, page_token, page_limit
+):
+    endpoint, requests, responses = catalog_server
+    responses.extend(
+        [
+            (200, {"namespaces": ["first", "second"], "page_token": "a/b"}),
+            (200, {"namespaces": [], "page_token": "next"}),
+            (200, {"namespaces": ["last", "final"], "page_token": ""}),
+        ]
+    )
+    catalog = (
+        await lancedb.connect_catalog_async(endpoint)
+        if asynchronous
+        else lancedb.connect_catalog(endpoint)
+    )
+    names = catalog.list_databases(page_token=page_token, page_limit=page_limit)
+    assert requests == []
+    assert names.num_page_results() == 0
+    assert names.page_token() == page_token
+    for cached, expected in [(1, "first"), (0, "second")]:
+        assert (await names.__anext__() if asynchronous else next(names)) == expected
+        assert names.num_page_results() == cached
+        assert names.page_token() == "a/b"
+        assert len(requests) == 1
+    assert (await names.__anext__() if asynchronous else next(names)) == "last"
+    assert names.num_page_results() == 1
+    assert names.page_token() is None
+    assert (await names.__anext__() if asynchronous else next(names)) == "final"
+    assert names.num_page_results() == 0
+    assert names.page_token() is None
+    for _ in range(2):
+        with pytest.raises(StopAsyncIteration if asynchronous else StopIteration):
+            if asynchronous:
+                await names.__anext__()
+            else:
+                next(names)
+    from urllib.parse import parse_qs, urlsplit
+
+    for request, token in zip(requests, [page_token, "a/b", "next"]):
+        url = urlsplit(request[0])
+        assert url.path == "/v1/namespace/$/list"
+        expected = {}
+        if token is not None:
+            expected["page_token"] = [token]
+        if page_limit is not None:
+            expected["limit"] = [str(page_limit)]
+        assert parse_qs(url.query) == expected
+
+
+@pytest.mark.parametrize("asynchronous", [False, True])
+@pytest.mark.parametrize("fail", [False, True])
+@pytest.mark.asyncio
+async def test_database_names_empty_or_error(catalog_server, asynchronous, fail):
+    endpoint, requests, responses = catalog_server
+    responses.append(
+        (401, {"error": "unauthorized"}) if fail else (200, {"namespaces": []})
+    )
+    catalog = (
+        await lancedb.connect_catalog_async(endpoint)
+        if asynchronous
+        else lancedb.connect_catalog(endpoint)
+    )
+    names = catalog.list_databases()
+    if fail:
+        with pytest.raises(HttpError):
+            if asynchronous:
+                await names.__anext__()
+            else:
+                next(names)
+    assert ([name async for name in names] if asynchronous else list(names)) == []
+    assert len(requests) == 1
+
+
+@pytest.mark.parametrize("asynchronous", [False, True])
+@pytest.mark.asyncio
+async def test_database_names_resume(catalog_server, asynchronous):
+    endpoint, requests, responses = catalog_server
+    responses.extend(
+        [
+            (200, {"namespaces": ["first"], "page_token": "resume/token"}),
+            (200, {"namespaces": ["second"]}),
+        ]
+    )
+    catalog = (
+        await lancedb.connect_catalog_async(endpoint)
+        if asynchronous
+        else lancedb.connect_catalog(endpoint)
+    )
+    names = catalog.list_databases(page_limit=1)
+    assert (await names.__anext__() if asynchronous else next(names)) == "first"
+    assert names.num_page_results() == 0
+    resumed = catalog.list_databases(page_token=names.page_token(), page_limit=1)
+    assert resumed.page_token() == "resume/token"
+    assert len(requests) == 1
+    assert ([name async for name in resumed] if asynchronous else list(resumed)) == [
+        "second"
+    ]
+    assert resumed.page_token() is None
+    assert resumed.num_page_results() == 0
+    assert requests[1][0] == "/v1/namespace/$/list?limit=1&page_token=resume%2Ftoken"
+
+
+@pytest.mark.parametrize("asynchronous", [False, True])
+@pytest.mark.parametrize("page_limit", [0, 2**32 - 1])
+@pytest.mark.asyncio
+async def test_database_names_invalid_limit(catalog_server, asynchronous, page_limit):
+    endpoint, requests, _ = catalog_server
+    catalog = (
+        await lancedb.connect_catalog_async(endpoint)
+        if asynchronous
+        else lancedb.connect_catalog(endpoint)
+    )
+    names = catalog.list_databases(page_limit=page_limit)
+    with pytest.raises(ValueError, match="limit"):
+        if asynchronous:
+            await names.__anext__()
+        else:
+            next(names)
+    assert names.num_page_results() == 0
+    assert requests == []

python/src/catalog.rs

@@ -1,12 +1,13 @@
 // SPDX-License-Identifier: Apache-2.0
 // SPDX-FileCopyrightText: Copyright The LanceDB Authors
 
+use std::sync::{Arc, Mutex};
 use std::time::Duration;
 
-use lancedb::catalog::{
-    CatalogConnection, CreateDatabaseRequest, DropDatabaseRequest, ListDatabasesRequest,
-};
-use pyo3::exceptions::PyValueError;
+use futures::{StreamExt, future::poll_fn};
+
+use lancedb::catalog::{CatalogConnection, CreateDatabaseRequest, DropDatabaseRequest};
+use pyo3::exceptions::{PyStopAsyncIteration, PyValueError};
 use pyo3::{Bound, PyAny, PyRef, PyResult, Python, pyclass, pyfunction, pymethods};
 
 use crate::connection::{Connection, PyClientConfig};
@@ -67,19 +68,48 @@ impl Catalog {
         })
     }
 
-    #[pyo3(signature = (*, limit=None, page_token=None))]
-    fn list_databases<'py>(
-        self_: PyRef<'py, Self>,
-        limit: Option<u32>,
-        page_token: Option<String>,
-    ) -> PyResult<Bound<'py, PyAny>> {
+    #[pyo3(signature = (*, page_token=None, page_limit=None))]
+    fn list_databases(&self, page_token: Option<String>, page_limit: Option<u32>) -> DatabaseNames {
+        DatabaseNames {
+            inner: Arc::new(Mutex::new(
+                self.inner.list_databases(page_token, page_limit),
+            )),
+            next_lock: Arc::new(tokio::sync::Mutex::new(())),
+        }
+    }
+}
+
+#[pyclass]
+pub struct DatabaseNames {
+    inner: Arc<Mutex<lancedb::catalog::DatabaseNames>>,
+    next_lock: Arc<tokio::sync::Mutex<()>>,
+}
+
+#[pymethods]
+impl DatabaseNames {
+    fn num_page_results(&self, py: Python<'_>) -> usize {
+        py.detach(|| self.inner.lock().unwrap().num_page_results())
+    }
+
+    fn page_token(&self, py: Python<'_>) -> Option<String> {
+        py.detach(|| self.inner.lock().unwrap().page_token().map(str::to_owned))
+    }
+
+    fn __aiter__(self_: PyRef<'_, Self>) -> PyRef<'_, Self> {
+        self_
+    }
+
+    fn __anext__(self_: PyRef<'_, Self>) -> PyResult<Bound<'_, PyAny>> {
         let inner = self_.inner.clone();
+        let next_lock = self_.next_lock.clone();
         future_into_py(self_.py(), async move {
-            let mut request = ListDatabasesRequest::default();
-            request.limit = limit;
-            request.page_token = page_token;
-            let response = inner.list_databases(request).await.infer_error()?;
-            Ok((response.databases, response.page_token))
+            // Serialize advances while allowing synchronous state inspection
+            // between polls, including while a REST request is pending.
+            let _guard = next_lock.lock().await;
+            poll_fn(|cx| inner.lock().unwrap().poll_next_unpin(cx))
+                .await
+                .ok_or_else(|| PyStopAsyncIteration::new_err(""))?
+                .infer_error()
         })
     }
 }

python/src/lib.rs

@@ -64,6 +64,7 @@ pub fn _lancedb(py: Python, m: &Bound<'_, PyModule>) -> PyResult<()> {
     env_logger::init_from_env(env);
     m.add_class::<Connection>()?;
     m.add_class::<catalog::Catalog>()?;
+    m.add_class::<catalog::DatabaseNames>()?;
     m.add_function(wrap_pyfunction!(catalog::connect_catalog, m)?)?;
     m.add_class::<Session>()?;
     m.add_class::<Table>()?;