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.
This commit is contained in:
Colin Patrick McCabe
2026-10-05 19:41:53 -07:00
committed by GitHub
parent f4544628ba
commit d864e7198f
18 changed files with 842 additions and 115 deletions
+15 -4
View File
@@ -89,19 +89,30 @@ Drop an empty database. The server rejects nonempty databases.
### listDatabases()
```ts
listDatabases(options): Promise<ListDatabasesResponse>
listDatabases(options): DatabaseNames
```
List one page of databases; pass pageToken from a response for the next page.
Iterate lazily over all database names using `for await...of`.
pageToken resumes from a saved token; omitted starts at the beginning.
pageLimit limits each REST response, not the total; omitted uses the server default.
Request errors are raised during iteration and terminate the iterator.
#### Parameters
* **options** = `{}`
* **options.limit?**: `number`
* **options.pageLimit?**: `number`
* **options.pageToken?**: `string`
#### Returns
`Promise`&lt;[`ListDatabasesResponse`](../interfaces/ListDatabasesResponse.md)&gt;
[`DatabaseNames`](DatabaseNames.md)
#### Example
```ts
for await (const name of catalog.listDatabases({ pageLimit: 20 })) {
console.log(name);
}
```
+79
View File
@@ -0,0 +1,79 @@
[**@lancedb/lancedb**](../README.md) • **Docs**
***
[@lancedb/lancedb](../globals.md) / DatabaseNames
# Class: DatabaseNames
A lazy async iterator of database names with pagination state.
Returned by [Catalog.listDatabases](Catalog.md#listdatabases).
## Implements
- `AsyncIterableIterator`&lt;`string`&gt;
## Methods
### \[asyncIterator\]()
```ts
asyncIterator: AsyncIterableIterator<string>
```
#### Returns
`AsyncIterableIterator`&lt;`string`&gt;
#### Implementation of
`AsyncIterableIterator.[asyncIterator]`
***
### next()
```ts
next(): Promise<IteratorResult<string, any>>
```
Fetch the next name, requesting another page only when needed.
#### Returns
`Promise`&lt;`IteratorResult`&lt;`string`, `any`&gt;&gt;
#### Implementation of
`AsyncIterableIterator.next`
***
### numPageResults()
```ts
numPageResults(): number
```
Number of names available without another REST request.
#### Returns
`number`
***
### pageToken()
```ts
pageToken(): undefined | string
```
Token for the next REST request. Initially this is the supplied starting token.
Returns undefined after the final page is fetched, even if names remain cached.
Drain the cache before saving a token to avoid skipping names when resuming.
A failed request terminates iteration and retains its token for a new iterator.
#### Returns
`undefined` \| `string`
+3 -1
View File
@@ -28,5 +28,7 @@ headers; opened database connections inherit authentication and client options.
```ts
const catalog = await connectCatalog("https://my-server.example", { apiKey: "secret" });
const db = await catalog.createDatabase("analytics", { existOk: true });
const page = await catalog.listDatabases({ limit: 20 });
for await (const name of catalog.listDatabases({ pageLimit: 20 })) {
console.log(name);
}
```
+1 -1
View File
@@ -27,6 +27,7 @@
- [Branches](classes/Branches.md)
- [Catalog](classes/Catalog.md)
- [Connection](classes/Connection.md)
- [DatabaseNames](classes/DatabaseNames.md)
- [HeaderProvider](classes/HeaderProvider.md)
- [Index](classes/Index.md)
- [Job](classes/Job.md)
@@ -108,7 +109,6 @@
- [JobEventsOptions](interfaces/JobEventsOptions.md)
- [JobFailureInfo](interfaces/JobFailureInfo.md)
- [JobInfo](interfaces/JobInfo.md)
- [ListDatabasesResponse](interfaces/ListDatabasesResponse.md)
- [ListNamespacesOptions](interfaces/ListNamespacesOptions.md)
- [ListNamespacesResponse](interfaces/ListNamespacesResponse.md)
- [ListTablesOptions](interfaces/ListTablesOptions.md)
@@ -1,23 +0,0 @@
[**@lancedb/lancedb**](../README.md) • **Docs**
***
[@lancedb/lancedb](../globals.md) / ListDatabasesResponse
# Interface: ListDatabasesResponse
## Properties
### databases
```ts
databases: string[];
```
***
### pageToken?
```ts
optional pageToken: string;
```
+4 -1
View File
@@ -37,7 +37,8 @@ are ordinary connections. Dropping a database requires it to be empty.
::: lancedb.catalog.Catalog
::: lancedb.catalog.ListDatabasesResponse
::: lancedb.catalog.DatabaseNames
## Remote SQL
@@ -440,6 +441,8 @@ latency on your workload before opting in.
::: lancedb.catalog.AsyncCatalog
::: lancedb.catalog.AsyncDatabaseNames
Connections represent a connection to a LanceDb database and
can be used to create, list, or open tables.
+136 -5
View File
@@ -80,9 +80,12 @@ describe("remote catalog", () => {
const second = await catalog.connectDatabase("other");
expect(await second.tableNames()).toEqual([]);
expect(await first.tableNames()).toEqual([]);
expect(
await catalog.listDatabases({ limit: 1, pageToken: "a/b" }),
).toEqual({ databases: ["team/search"], pageToken: "next" });
const names = catalog.listDatabases({ pageLimit: 1, pageToken: "a/b" });
expect(await names.next()).toEqual({
done: false,
value: "team/search",
});
expect(names.pageToken()).toBe("next");
await catalog.dropDatabase("team/search", { ignoreMissing: true });
expect(requests[0].url).toBe("/v1/namespace/team%2Fsearch/create");
expect(requests[0].body).toEqual({ mode: "ExistOk" });
@@ -129,11 +132,139 @@ describe("remote catalog", () => {
);
});
it("lazily traverses pages, including empty pages, and exposes cached state", async () => {
await withCatalog(
[
// biome-ignore lint/style/useNamingConvention: server wire format
[200, { namespaces: ["a", "b"], page_token: "empty" }],
// biome-ignore lint/style/useNamingConvention: server wire format
[200, { namespaces: [], page_token: "last" }],
// biome-ignore lint/style/useNamingConvention: server wire format
[200, { namespaces: ["c", "d"], page_token: "" }],
],
async (catalog, requests) => {
const names = catalog.listDatabases({ pageLimit: 2 });
expect(names[Symbol.asyncIterator]()).toBe(names);
expect(requests).toHaveLength(0);
expect(names.numPageResults()).toBe(0);
expect(names.pageToken()).toBeUndefined();
expect(await names.next()).toEqual({ done: false, value: "a" });
expect(names.numPageResults()).toBe(1);
expect(names.pageToken()).toBe("empty");
expect(await names.next()).toEqual({ done: false, value: "b" });
expect(requests).toHaveLength(1);
expect(await names.next()).toEqual({ done: false, value: "c" });
expect(requests).toHaveLength(3);
expect(names.pageToken()).toBeUndefined();
expect(names.numPageResults()).toBe(1);
const remaining = [];
for await (const name of names) remaining.push(name);
expect(remaining).toEqual(["d"]);
expect(names.numPageResults()).toBe(0);
expect(await names.next()).toEqual({ done: true, value: undefined });
expect(requests.map((request) => request.url)).toEqual([
"/v1/namespace/$/list?limit=2",
"/v1/namespace/$/list?limit=2&page_token=empty",
"/v1/namespace/$/list?limit=2&page_token=last",
]);
},
);
});
it("resumes from a saved token after draining the cache", async () => {
await withCatalog(
[
// biome-ignore lint/style/useNamingConvention: server wire format
[200, { namespaces: ["a", "b"], page_token: "next/token" }],
[200, { namespaces: ["c"] }],
],
async (catalog, requests) => {
const names = catalog.listDatabases({
pageToken: "start/token",
pageLimit: 2,
});
expect(names.pageToken()).toBe("start/token");
expect(await names.next()).toEqual({ done: false, value: "a" });
while (names.numPageResults() > 0) await names.next();
const resumed = catalog.listDatabases({
pageToken: names.pageToken(),
pageLimit: 2,
});
const results = [];
for await (const name of resumed) results.push(name);
expect(results).toEqual(["c"]);
expect(requests.map((request) => request.url)).toEqual([
"/v1/namespace/$/list?limit=2&page_token=start%2Ftoken",
"/v1/namespace/$/list?limit=2&page_token=next%2Ftoken",
]);
},
);
});
it("terminates on errors and retains the failed request token", async () => {
await withCatalog(
[
// biome-ignore lint/style/useNamingConvention: server wire format
[200, { namespaces: ["a"], page_token: "retry" }],
[400, {}],
[200, { namespaces: ["b"] }],
],
async (catalog, requests) => {
const names = catalog.listDatabases();
await names.next();
await expect(names.next()).rejects.toThrow();
expect(names.pageToken()).toBe("retry");
expect(names.numPageResults()).toBe(0);
expect(await names.next()).toEqual({ done: true, value: undefined });
expect(requests).toHaveLength(2);
const resumed = catalog.listDatabases({ pageToken: names.pageToken() });
expect(await resumed.next()).toEqual({ done: false, value: "b" });
expect(await resumed.next()).toEqual({ done: true, value: undefined });
},
);
});
it("handles an empty listing and concurrent advances", async () => {
await withCatalog(
[
[200, { namespaces: [] }],
[200, { namespaces: ["a", "b"] }],
],
async (catalog, requests) => {
const empty = catalog.listDatabases();
expect(await empty.next()).toEqual({ done: true, value: undefined });
expect(await empty.next()).toEqual({ done: true, value: undefined });
const names = catalog.listDatabases();
const pending = [names.next(), names.next(), names.next()];
expect(names.numPageResults()).toBe(0);
expect(names.pageToken()).toBeUndefined();
const results = await Promise.all(pending);
expect(
results
.filter((result) => !result.done)
.map((result) => result.value)
.sort(),
).toEqual(["a", "b"]);
expect(results.filter((result) => result.done)).toHaveLength(1);
expect(requests).toHaveLength(2);
},
);
});
it("validates endpoints and pagination", async () => {
await expect(connectCatalog("/tmp/catalog")).rejects.toThrow();
const catalog = await connectCatalog("http://127.0.0.1:1");
for (const limit of [0, -1, 1.5, 2147483648]) {
await expect(catalog.listDatabases({ limit })).rejects.toThrow("limit");
for (const limit of [
0,
-1,
1.5,
2147483648,
Number.NaN,
Number.POSITIVE_INFINITY,
]) {
expect(() => catalog.listDatabases({ pageLimit: limit })).toThrow(
"limit",
);
}
await expect(catalog.connectDatabase("a$b")).rejects.toThrow(
"Invalid database name",
+62 -12
View File
@@ -5,9 +5,9 @@ import { Connection, LocalConnection } from "./connection";
import { HeaderProvider } from "./header";
import {
JsHeaderProvider,
ListDatabasesResponse,
Catalog as NativeCatalog,
CatalogOptions as NativeCatalogOptions,
DatabaseNames as NativeDatabaseNames,
} from "./native.js";
import { OAuthConfig } from "./oauth";
@@ -21,7 +21,41 @@ export interface CatalogOptions
| (() => Record<string, string> | Promise<Record<string, string>>);
}
export type { ListDatabasesResponse } from "./native.js";
/**
* A lazy async iterator of database names with pagination state.
* Returned by {@link Catalog.listDatabases}.
*/
export class DatabaseNames implements AsyncIterableIterator<string> {
/** @hidden */
constructor(private readonly inner: NativeDatabaseNames) {}
[Symbol.asyncIterator](): AsyncIterableIterator<string> {
return this;
}
/** Fetch the next name, requesting another page only when needed. */
async next(): Promise<IteratorResult<string>> {
const value = await this.inner.next();
return value == null
? { done: true, value: undefined }
: { done: false, value };
}
/** Number of names available without another REST request. */
numPageResults(): number {
return this.inner.numPageResults();
}
/**
* Token for the next REST request. Initially this is the supplied starting token.
* Returns undefined after the final page is fetched, even if names remain cached.
* Drain the cache before saving a token to avoid skipping names when resuming.
* A failed request terminates iteration and retains its token for a new iterator.
*/
pageToken(): string | undefined {
return this.inner.pageToken() ?? undefined;
}
}
/** A remote catalog manages databases through the server's root namespace. */
export class Catalog {
@@ -56,21 +90,35 @@ export class Catalog {
await this.inner.dropDatabase(name, options.ignoreMissing);
}
/** List one page of databases; pass pageToken from a response for the next page. */
async listDatabases(
options: { limit?: number; pageToken?: string } = {},
): Promise<ListDatabasesResponse> {
/**
* Iterate lazily over all database names using `for await...of`.
* pageToken resumes from a saved token; omitted starts at the beginning.
* pageLimit limits each REST response, not the total; omitted uses the server default.
* Request errors are raised during iteration and terminate the iterator.
*
* @example
* ```ts
* for await (const name of catalog.listDatabases({ pageLimit: 20 })) {
* console.log(name);
* }
* ```
*/
listDatabases(
options: { pageLimit?: number; pageToken?: string } = {},
): DatabaseNames {
if (
options.limit !== undefined &&
(!Number.isInteger(options.limit) ||
options.limit <= 0 ||
options.limit > 2147483647)
options.pageLimit !== undefined &&
(!Number.isInteger(options.pageLimit) ||
options.pageLimit <= 0 ||
options.pageLimit > 2147483647)
) {
throw new Error(
"Database list limit must be an integer between 1 and 2147483647",
);
}
return this.inner.listDatabases(options.limit, options.pageToken);
return new DatabaseNames(
this.inner.listDatabases(options.pageToken, options.pageLimit),
);
}
}
@@ -82,7 +130,9 @@ export class Catalog {
* ```ts
* const catalog = await connectCatalog("https://my-server.example", { apiKey: "secret" });
* const db = await catalog.createDatabase("analytics", { existOk: true });
* const page = await catalog.listDatabases({ limit: 20 });
* for await (const name of catalog.listDatabases({ pageLimit: 20 })) {
* console.log(name);
* }
* ```
*/
export async function connectCatalog(
+1 -1
View File
@@ -646,6 +646,6 @@ export async function connectNamespace(
export {
Catalog,
CatalogOptions,
ListDatabasesResponse,
DatabaseNames,
connectCatalog,
} from "./catalog";
+42 -19
View File
@@ -1,12 +1,12 @@
// SPDX-License-Identifier: Apache-2.0
// SPDX-FileCopyrightText: Copyright The LanceDB Authors
use std::sync::Arc;
use std::sync::{Arc, Mutex};
use std::time::Duration;
use lancedb::catalog::{
CatalogConnection, CreateDatabaseRequest, DropDatabaseRequest, ListDatabasesRequest,
};
use futures::{StreamExt, future::poll_fn};
use lancedb::catalog::{CatalogConnection, CreateDatabaseRequest, DropDatabaseRequest};
use napi::bindgen_prelude::*;
use napi_derive::napi;
@@ -25,10 +25,37 @@ pub struct CatalogOptions {
pub oauth_config: Option<OAuthConfig>,
}
#[napi(object)]
pub struct ListDatabasesResponse {
pub databases: Vec<String>,
pub page_token: Option<String>,
/// A lazy iterator over database names with inspectable pagination state.
#[napi]
pub struct DatabaseNames {
inner: Mutex<lancedb::catalog::DatabaseNames>,
next_lock: futures::lock::Mutex<()>,
}
#[napi]
impl DatabaseNames {
/// Number of names cached without another REST request.
#[napi]
pub fn num_page_results(&self) -> u32 {
self.inner.lock().unwrap().num_page_results() as u32
}
/// Continuation token for the next REST request.
#[napi]
pub fn page_token(&self) -> Option<String> {
self.inner.lock().unwrap().page_token().map(str::to_owned)
}
/// Fetch the next name, returning None when exhausted.
#[napi]
pub async fn next(&self) -> Result<Option<String>> {
// Serialize advances but allow synchronous inspection between polls.
let _guard = self.next_lock.lock().await;
poll_fn(|cx| self.inner.lock().unwrap().poll_next_unpin(cx))
.await
.transpose()
.default_error()
}
}
#[napi]
@@ -106,18 +133,14 @@ impl Catalog {
.default_error()
}
#[napi]
pub async fn list_databases(
pub fn list_databases(
&self,
limit: Option<u32>,
page_token: Option<String>,
) -> Result<ListDatabasesResponse> {
let mut request = ListDatabasesRequest::default();
request.limit = limit;
request.page_token = page_token;
let response = self.inner.list_databases(request).await.default_error()?;
Ok(ListDatabasesResponse {
databases: response.databases,
page_token: response.page_token,
})
page_limit: Option<u32>,
) -> DatabaseNames {
DatabaseNames {
inner: Mutex::new(self.inner.list_databases(page_token, page_limit)),
next_lock: futures::lock::Mutex::new(()),
}
}
}
-2
View File
@@ -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",
+9 -3
View File
@@ -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,
+91 -19
View File
@@ -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."""
+129 -4
View File
@@ -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 == []
+45 -15
View File
@@ -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()
})
}
}
+1
View File
@@ -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>()?;
+100 -5
View File
@@ -15,7 +15,12 @@
//! ```
use std::fmt;
use std::pin::Pin;
use std::sync::Arc;
use std::task::{Context, Poll};
use futures::future::BoxFuture;
use futures::{FutureExt, Stream};
use crate::Result;
use crate::connection::Connection;
@@ -118,6 +123,73 @@ pub struct ListDatabasesResponse {
pub page_token: Option<String>,
}
/// A lazy stream of database names with inspectable pagination state.
///
/// The continuation token advances when a page is fetched, so resuming from it
/// skips any names still cached in this stream. Drain the cache before saving a
/// token if those names must be included. Errors terminate the stream; the token
/// for the failed request remains available for resuming with a new stream.
pub struct DatabaseNames {
catalog: Arc<dyn Catalog>,
names: std::vec::IntoIter<String>,
page_token: Option<String>,
page_limit: Option<u32>,
done: bool,
pending: Option<BoxFuture<'static, Result<ListDatabasesResponse>>>,
}
impl DatabaseNames {
/// Number of names available without another REST request.
pub fn num_page_results(&self) -> usize {
self.names.len()
}
/// Token for the next REST request, or `None` after the last page is fetched.
///
/// Before the first request, returns the supplied starting token (`None`
/// means start from the beginning). A final page can still have cached names.
pub fn page_token(&self) -> Option<&str> {
self.page_token.as_deref()
}
}
impl Stream for DatabaseNames {
type Item = Result<String>;
fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
let this = self.get_mut();
loop {
if let Some(name) = this.names.next() {
return Poll::Ready(Some(Ok(name)));
}
if this.done {
return Poll::Ready(None);
}
let pending = this.pending.get_or_insert_with(|| {
let catalog = this.catalog.clone();
let request = ListDatabasesRequest {
page_token: this.page_token.clone(),
limit: this.page_limit,
};
async move { catalog.list_databases(request).await }.boxed()
});
let response = futures::ready!(pending.as_mut().poll(cx));
this.pending = None;
match response {
Ok(response) => {
this.page_token = response.page_token.filter(|token| !token.is_empty());
this.done = this.page_token.is_none();
this.names = response.databases.into_iter();
}
Err(error) => {
this.done = true;
return Poll::Ready(Some(Err(error)));
}
}
}
}
}
/// A backend that manages databases. Implementations own database lifecycle semantics.
#[async_trait::async_trait]
pub trait Catalog: Send + Sync + std::fmt::Debug + 'static {
@@ -188,12 +260,35 @@ impl CatalogConnection {
self.catalog.drop_database(request.into()).await
}
/// List a page of databases. Pass [`ListDatabasesRequest::default`] for the first page.
pub async fn list_databases(
/// Iterate lazily over all database names, fetching more results as needed.
///
/// `page_token` starts at a saved continuation token; `None` starts at the
/// beginning. `page_limit` limits each REST response (not the whole stream);
/// `None` uses the server default. Errors terminate the listing.
///
/// ```
/// # async fn example(catalog: &lancedb::catalog::CatalogConnection) -> lancedb::Result<()> {
/// use futures::TryStreamExt;
/// let mut names = catalog.list_databases(None, None);
/// while let Some(name) = names.try_next().await? {
/// println!("{name}");
/// }
/// # Ok(())
/// # }
/// ```
pub fn list_databases(
&self,
request: ListDatabasesRequest,
) -> Result<ListDatabasesResponse> {
self.catalog.list_databases(request).await
page_token: Option<String>,
page_limit: Option<u32>,
) -> DatabaseNames {
DatabaseNames {
catalog: self.catalog.clone(),
names: Vec::new().into_iter(),
page_token,
page_limit,
done: false,
pending: None,
}
}
/// Connect to a database without creating it if it is missing.
+124
View File
@@ -416,6 +416,7 @@ mod tests {
RemoteCatalog::try_new(&endpoint, options).unwrap(),
));
let page = catalog
.catalog()
.list_databases(ListDatabasesRequest::default().limit(1).page_token("a/b"))
.await
.unwrap();
@@ -443,6 +444,7 @@ mod tests {
.unwrap();
assert!(
catalog
.catalog()
.list_databases(ListDatabasesRequest::default())
.await
.unwrap()
@@ -538,6 +540,128 @@ mod tests {
);
}
#[tokio::test]
async fn database_names_follow_pages_lazily() {
use futures::TryStreamExt;
let (endpoint, task) = server(vec![
(
200,
json!({"namespaces": ["first", "second"], "page_token": "a/b"}),
),
(200, json!({"namespaces": [], "page_token": "next"})),
(200, json!({"namespaces": ["last"], "page_token": ""})),
])
.await;
let catalog = CatalogConnection::new(Arc::new(
RemoteCatalog::try_new(endpoint, RemoteCatalogOptions::default()).unwrap(),
));
let mut names = catalog.list_databases(None, None);
drop(catalog);
assert_eq!(names.num_page_results(), 0);
assert_eq!(names.page_token(), None);
assert_eq!(names.try_next().await.unwrap().as_deref(), Some("first"));
assert_eq!(names.num_page_results(), 1);
assert_eq!(names.page_token(), Some("a/b"));
assert_eq!(names.try_next().await.unwrap().as_deref(), Some("second"));
assert_eq!(names.num_page_results(), 0);
assert_eq!(names.page_token(), Some("a/b"));
assert!(!task.is_finished());
assert_eq!(names.try_next().await.unwrap().as_deref(), Some("last"));
assert_eq!(names.try_next().await.unwrap(), None);
assert_eq!(names.try_next().await.unwrap(), None);
let requests = task.await.unwrap();
assert_eq!(requests[0].line, "GET /v1/namespace/$/list HTTP/1.1");
assert_eq!(
requests[1].line,
"GET /v1/namespace/$/list?page_token=a%2Fb HTTP/1.1"
);
assert_eq!(
requests[2].line,
"GET /v1/namespace/$/list?page_token=next HTTP/1.1"
);
}
#[tokio::test]
async fn database_names_stop_after_error() {
use futures::StreamExt;
let (endpoint, task) = server(vec![
(200, json!({"namespaces": ["first"], "page_token": "next"})),
(401, json!({"error": "unauthorized"})),
])
.await;
let catalog = CatalogConnection::new(Arc::new(
RemoteCatalog::try_new(endpoint, RemoteCatalogOptions::default()).unwrap(),
));
let mut names = catalog.list_databases(None, None);
assert_eq!(names.next().await.unwrap().unwrap(), "first");
assert!(names.next().await.unwrap().is_err());
assert_eq!(names.num_page_results(), 0);
assert_eq!(names.page_token(), Some("next"));
assert!(names.next().await.is_none());
task.await.unwrap();
}
#[tokio::test]
async fn database_names_can_resume_with_page_limit() {
use futures::TryStreamExt;
let (endpoint, task) = server(vec![
(
200,
json!({"namespaces": ["first", "second"], "page_token": "next/token"}),
),
(200, json!({"namespaces": ["third", "fourth"]})),
])
.await;
let catalog = CatalogConnection::new(Arc::new(
RemoteCatalog::try_new(endpoint, RemoteCatalogOptions::default()).unwrap(),
));
let mut names = catalog.list_databases(Some("start/token".into()), Some(2));
assert_eq!(names.page_token(), Some("start/token"));
assert_eq!(names.num_page_results(), 0);
assert_eq!(names.try_next().await.unwrap().as_deref(), Some("first"));
assert_eq!(names.page_token(), Some("next/token"));
assert_eq!(names.num_page_results(), 1);
assert_eq!(names.try_next().await.unwrap().as_deref(), Some("second"));
assert_eq!(names.num_page_results(), 0);
let token = names.page_token().map(str::to_owned);
drop(names);
let mut resumed = catalog.list_databases(token, Some(2));
assert_eq!(resumed.try_next().await.unwrap().as_deref(), Some("third"));
assert_eq!(resumed.page_token(), None);
assert_eq!(resumed.num_page_results(), 1);
assert_eq!(resumed.try_next().await.unwrap().as_deref(), Some("fourth"));
assert_eq!(resumed.num_page_results(), 0);
assert_eq!(resumed.try_next().await.unwrap(), None);
let requests = task.await.unwrap();
assert_eq!(
requests[0].line,
"GET /v1/namespace/$/list?limit=2&page_token=start%2Ftoken HTTP/1.1"
);
assert_eq!(
requests[1].line,
"GET /v1/namespace/$/list?limit=2&page_token=next%2Ftoken HTTP/1.1"
);
}
#[tokio::test]
async fn database_names_reject_invalid_page_limit() {
use futures::TryStreamExt;
let catalog = CatalogConnection::new(Arc::new(
RemoteCatalog::try_new("http://127.0.0.1:1", RemoteCatalogOptions::default()).unwrap(),
));
for limit in [0, u32::MAX] {
let mut names = catalog.list_databases(None, Some(limit));
assert!(matches!(
names.try_next().await,
Err(Error::InvalidInput { .. })
));
assert_eq!(names.num_page_results(), 0);
assert_eq!(names.try_next().await.unwrap(), None);
}
}
#[tokio::test]
async fn catalog_validates_before_sending_requests() {
for endpoint in [