diff --git a/docs/src/js/classes/Catalog.md b/docs/src/js/classes/Catalog.md index deac5e3ce..8351380e2 100644 --- a/docs/src/js/classes/Catalog.md +++ b/docs/src/js/classes/Catalog.md @@ -89,19 +89,30 @@ Drop an empty database. The server rejects nonempty databases. ### listDatabases() ```ts -listDatabases(options): Promise +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`<[`ListDatabasesResponse`](../interfaces/ListDatabasesResponse.md)> +[`DatabaseNames`](DatabaseNames.md) + +#### Example + +```ts +for await (const name of catalog.listDatabases({ pageLimit: 20 })) { + console.log(name); +} +``` diff --git a/docs/src/js/classes/DatabaseNames.md b/docs/src/js/classes/DatabaseNames.md new file mode 100644 index 000000000..6b78f93b8 --- /dev/null +++ b/docs/src/js/classes/DatabaseNames.md @@ -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`<`string`> + +## Methods + +### \[asyncIterator\]() + +```ts +asyncIterator: AsyncIterableIterator +``` + +#### Returns + +`AsyncIterableIterator`<`string`> + +#### Implementation of + +`AsyncIterableIterator.[asyncIterator]` + +*** + +### next() + +```ts +next(): Promise> +``` + +Fetch the next name, requesting another page only when needed. + +#### Returns + +`Promise`<`IteratorResult`<`string`, `any`>> + +#### 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` diff --git a/docs/src/js/functions/connectCatalog.md b/docs/src/js/functions/connectCatalog.md index 2bdb3182f..ca51c0ed4 100644 --- a/docs/src/js/functions/connectCatalog.md +++ b/docs/src/js/functions/connectCatalog.md @@ -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); +} ``` diff --git a/docs/src/js/globals.md b/docs/src/js/globals.md index 0c48ad4ec..1a939ae28 100644 --- a/docs/src/js/globals.md +++ b/docs/src/js/globals.md @@ -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) diff --git a/docs/src/js/interfaces/ListDatabasesResponse.md b/docs/src/js/interfaces/ListDatabasesResponse.md deleted file mode 100644 index a1a247746..000000000 --- a/docs/src/js/interfaces/ListDatabasesResponse.md +++ /dev/null @@ -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; -``` diff --git a/docs/src/python/python.md b/docs/src/python/python.md index 66fda4378..f946fbb13 100644 --- a/docs/src/python/python.md +++ b/docs/src/python/python.md @@ -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. diff --git a/nodejs/__test__/catalog.test.ts b/nodejs/__test__/catalog.test.ts index f9158c94f..945956b5b 100644 --- a/nodejs/__test__/catalog.test.ts +++ b/nodejs/__test__/catalog.test.ts @@ -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", diff --git a/nodejs/lancedb/catalog.ts b/nodejs/lancedb/catalog.ts index 5cfe2c5ea..e344d8187 100644 --- a/nodejs/lancedb/catalog.ts +++ b/nodejs/lancedb/catalog.ts @@ -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 | Promise>); } -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 { + /** @hidden */ + constructor(private readonly inner: NativeDatabaseNames) {} + + [Symbol.asyncIterator](): AsyncIterableIterator { + return this; + } + + /** Fetch the next name, requesting another page only when needed. */ + async next(): Promise> { + 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 { + /** + * 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( diff --git a/nodejs/lancedb/index.ts b/nodejs/lancedb/index.ts index 72b8c4a56..694db9fa7 100644 --- a/nodejs/lancedb/index.ts +++ b/nodejs/lancedb/index.ts @@ -646,6 +646,6 @@ export async function connectNamespace( export { Catalog, CatalogOptions, - ListDatabasesResponse, + DatabaseNames, connectCatalog, } from "./catalog"; diff --git a/nodejs/src/catalog.rs b/nodejs/src/catalog.rs index c2288bc4a..82edda5da 100644 --- a/nodejs/src/catalog.rs +++ b/nodejs/src/catalog.rs @@ -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, } -#[napi(object)] -pub struct ListDatabasesResponse { - pub databases: Vec, - pub page_token: Option, +/// A lazy iterator over database names with inspectable pagination state. +#[napi] +pub struct DatabaseNames { + inner: Mutex, + 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 { + 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> { + // 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, page_token: Option, - ) -> Result { - 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, + ) -> DatabaseNames { + DatabaseNames { + inner: Mutex::new(self.inner.list_databases(page_token, page_limit)), + next_lock: futures::lock::Mutex::new(()), + } } } diff --git a/python/python/lancedb/__init__.py b/python/python/lancedb/__init__.py index d3a678359..45a69a92e 100644 --- a/python/python/lancedb/__init__.py +++ b/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", diff --git a/python/python/lancedb/_lancedb.pyi b/python/python/lancedb/_lancedb.pyi index 176e57ac3..bbd6a3c5d 100644 --- a/python/python/lancedb/_lancedb.pyi +++ b/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, diff --git a/python/python/lancedb/catalog.py b/python/python/lancedb/catalog.py index 4fd941ddd..70aab1bbb 100644 --- a/python/python/lancedb/catalog.py +++ b/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.""" diff --git a/python/python/tests/test_catalog.py b/python/python/tests/test_catalog.py index ca5d68fdf..e7b7b1ec2 100644 --- a/python/python/tests/test_catalog.py +++ b/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 == [] diff --git a/python/src/catalog.rs b/python/src/catalog.rs index f61b409ff..e7cc56c49 100644 --- a/python/src/catalog.rs +++ b/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, - page_token: Option, - ) -> PyResult> { + #[pyo3(signature = (*, page_token=None, page_limit=None))] + fn list_databases(&self, page_token: Option, page_limit: Option) -> 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>, + next_lock: Arc>, +} + +#[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 { + 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> { 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() }) } } diff --git a/python/src/lib.rs b/python/src/lib.rs index ab04663c0..57d30d183 100644 --- a/python/src/lib.rs +++ b/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::()?; m.add_class::()?; + m.add_class::()?; m.add_function(wrap_pyfunction!(catalog::connect_catalog, m)?)?; m.add_class::()?; m.add_class::()?; diff --git a/rust/lancedb/src/catalog.rs b/rust/lancedb/src/catalog.rs index 5eaa1d49e..576eb8e88 100644 --- a/rust/lancedb/src/catalog.rs +++ b/rust/lancedb/src/catalog.rs @@ -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, } +/// 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, + names: std::vec::IntoIter, + page_token: Option, + page_limit: Option, + done: bool, + pending: Option>>, +} + +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; + + fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { + 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 { - self.catalog.list_databases(request).await + page_token: Option, + page_limit: Option, + ) -> 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. diff --git a/rust/lancedb/src/remote/catalog.rs b/rust/lancedb/src/remote/catalog.rs index 1c0129889..72b77a87e 100644 --- a/rust/lancedb/src/remote/catalog.rs +++ b/rust/lancedb/src/remote/catalog.rs @@ -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 [