From de005f1d611a60d8ac38ae56e4e05db68f58919b Mon Sep 17 00:00:00 2001 From: Colin Patrick McCabe Date: Wed, 7 Oct 2026 10:24:20 -0700 Subject: [PATCH] refactor!: implement lazy iterators for several APIs (#4459) Implement lazy iterators for: - secrets - views - jobs - functions - materialized views --- docs/src/js/classes/Connection.md | 32 +- docs/src/js/classes/Listing.md | 91 ++++ docs/src/js/globals.md | 2 + docs/src/js/interfaces/ListingOptions.md | 29 + docs/src/python/python.md | 10 + nodejs/__test__/materialized_view.test.ts | 14 +- nodejs/__test__/remote.test.ts | 152 +++++- nodejs/lancedb/connection.ts | 50 +- nodejs/lancedb/index.ts | 2 + nodejs/lancedb/listing.ts | 73 +++ nodejs/src/connection.rs | 42 +- nodejs/src/lib.rs | 1 + nodejs/src/listing.rs | 83 +++ python/python/lancedb/_lancedb.pyi | 56 +- python/python/lancedb/db.py | 265 ++++++++-- python/python/lancedb/listing.py | 85 +++ python/python/lancedb/namespace.py | 40 +- python/python/lancedb/remote/db.py | 100 +++- .../tests/test_first_class_function_slice2.py | 14 +- python/python/tests/test_listing.py | 194 +++++++ .../python/tests/test_materialized_views.py | 36 +- python/python/tests/test_namespace.py | 25 + python/python/tests/test_remote_db.py | 4 +- python/src/connection.rs | 98 ++-- python/src/lib.rs | 5 + python/src/listing.rs | 131 +++++ rust/lancedb/src/connection.rs | 121 ++++- rust/lancedb/src/database.rs | 104 ++-- rust/lancedb/src/lib.rs | 1 + rust/lancedb/src/listing.rs | 261 +++++++++ rust/lancedb/src/materialized_view.rs | 23 +- rust/lancedb/src/remote/db.rs | 495 +++++++++++------- 32 files changed, 2185 insertions(+), 454 deletions(-) create mode 100644 docs/src/js/classes/Listing.md create mode 100644 docs/src/js/interfaces/ListingOptions.md create mode 100644 nodejs/lancedb/listing.ts create mode 100644 nodejs/src/listing.rs create mode 100644 python/python/lancedb/listing.py create mode 100644 python/python/tests/test_listing.py create mode 100644 python/src/listing.rs create mode 100644 rust/lancedb/src/listing.rs diff --git a/docs/src/js/classes/Connection.md b/docs/src/js/classes/Connection.md index 4bb769a0e..d5ab8c3d1 100644 --- a/docs/src/js/classes/Connection.md +++ b/docs/src/js/classes/Connection.md @@ -611,30 +611,42 @@ Return true if the connection has not been closed ### listJobs() ```ts -abstract listJobs(): Promise +abstract listJobs(options?): Listing ``` -List server-side jobs across the database's tables. +Iterate lazily over server-side jobs across the database's tables. +Use `for await...of`. Options control page size and the starting token. +See [Listing](Listing.md) for pagination state and error behavior. + +#### Parameters + +* **options?**: [`ListingOptions`](../interfaces/ListingOptions.md) #### Returns -`Promise`<[`JobInfo`](../interfaces/JobInfo.md)[]> +[`Listing`](Listing.md)<[`JobInfo`](../interfaces/JobInfo.md)> *** ### listMaterializedViews() ```ts -abstract listMaterializedViews(): Promise +abstract listMaterializedViews(options?): Listing ``` The names of the materialized views in this database. -Found by reading every table's schema, so this costs an open per table. +Iterate lazily with `for await...of`. Local listings inspect table schemas +one page at a time. Options control page size and the starting token. +See [Listing](Listing.md) for pagination state and error behavior. + +#### Parameters + +* **options?**: [`ListingOptions`](../interfaces/ListingOptions.md) #### Returns -`Promise`<`string`[]> +[`Listing`](Listing.md)<`string`> *** @@ -735,20 +747,24 @@ A page of table names and an ### listViews() ```ts -abstract listViews(namespacePath?): Promise +abstract listViews(namespacePath?, options?): Listing ``` The names of the views in one namespace. Names only; a definition comes from [describeView](Connection.md#describeview). +Iterate lazily with `for await...of`. See [Listing](Listing.md) for pagination +state and error behavior; options control page size and the starting token. #### Parameters * **namespacePath?**: `string`[] +* **options?**: [`ListingOptions`](../interfaces/ListingOptions.md) + #### Returns -`Promise`<`string`[]> +[`Listing`](Listing.md)<`string`> *** diff --git a/docs/src/js/classes/Listing.md b/docs/src/js/classes/Listing.md new file mode 100644 index 000000000..1a94df2dc --- /dev/null +++ b/docs/src/js/classes/Listing.md @@ -0,0 +1,91 @@ +[**@lancedb/lancedb**](../README.md) • **Docs** + +*** + +[@lancedb/lancedb](../globals.md) / Listing + +# Class: Listing<T> + +A lazy async iterator over database resources with pagination state. +Requests occur during iteration; request errors terminate the iterator. + +## Example + +```ts +for await (const name of connection.listViews([], { pageLimit: 20 })) { + console.log(name); +} +``` + +## Type Parameters + +• **T** + +## Implements + +- `AsyncIterableIterator`<`T`> + +## Methods + +### \[asyncIterator\]() + +```ts +asyncIterator: AsyncIterableIterator +``` + +#### Returns + +`AsyncIterableIterator`<`T`> + +#### Implementation of + +`AsyncIterableIterator.[asyncIterator]` + +*** + +### next() + +```ts +next(): Promise> +``` + +Fetch the next item, requesting another page only when needed. + +#### Returns + +`Promise`<`IteratorResult`<`T`, `any`>> + +#### Implementation of + +`AsyncIterableIterator.next` + +*** + +### numPageResults() + +```ts +numPageResults(): number +``` + +Number of cached items available without another request. + +#### Returns + +`number` + +*** + +### pageToken() + +```ts +pageToken(): undefined | string +``` + +Token for the next request. Initially this is the supplied starting token. +Returns undefined after fetching the final page, even with cached items. +Drain the cache before saving a token to avoid skipping items on resumption. +Failed requests terminate iteration and retain their token for a new iterator. + +#### Returns + +`undefined` \| `string` diff --git a/docs/src/js/globals.md b/docs/src/js/globals.md index 1a939ae28..9cf3136af 100644 --- a/docs/src/js/globals.md +++ b/docs/src/js/globals.md @@ -31,6 +31,7 @@ - [HeaderProvider](classes/HeaderProvider.md) - [Index](classes/Index.md) - [Job](classes/Job.md) +- [Listing](classes/Listing.md) - [MakeArrowTableOptions](classes/MakeArrowTableOptions.md) - [MatchQuery](classes/MatchQuery.md) - [MaterializedView](classes/MaterializedView.md) @@ -113,6 +114,7 @@ - [ListNamespacesResponse](interfaces/ListNamespacesResponse.md) - [ListTablesOptions](interfaces/ListTablesOptions.md) - [ListTablesResponse](interfaces/ListTablesResponse.md) +- [ListingOptions](interfaces/ListingOptions.md) - [LsmStats](interfaces/LsmStats.md) - [LsmWriteSpec](interfaces/LsmWriteSpec.md) - [MaterializedViewDefinition](interfaces/MaterializedViewDefinition.md) diff --git a/docs/src/js/interfaces/ListingOptions.md b/docs/src/js/interfaces/ListingOptions.md new file mode 100644 index 000000000..2070f8e3a --- /dev/null +++ b/docs/src/js/interfaces/ListingOptions.md @@ -0,0 +1,29 @@ +[**@lancedb/lancedb**](../README.md) • **Docs** + +*** + +[@lancedb/lancedb](../globals.md) / ListingOptions + +# Interface: ListingOptions + +Pagination controls for lazy resource listings. + +## Properties + +### pageLimit? + +```ts +optional pageLimit: number; +``` + +Maximum results per request, not a limit on the whole listing. + +*** + +### pageToken? + +```ts +optional pageToken: string; +``` + +Starting continuation token; omitted starts at the beginning. diff --git a/docs/src/python/python.md b/docs/src/python/python.md index 814bcb12a..863edbe4a 100644 --- a/docs/src/python/python.md +++ b/docs/src/python/python.md @@ -26,6 +26,14 @@ is also an [asynchronous API client](#connections-asynchronous). ::: lancedb.db.DBConnection +Resource listings (`list_secrets`, `list_views`, `list_jobs`, `list_functions`, and +`list_materialized_views`) return lazy iterators. Use `list(db.list_views())` to +collect synchronous results, or `[name async for name in db.list_views()]` with +an asynchronous connection. `page_limit` controls each request, and `page_token` +resumes a listing. Drain the iterator's cached results before saving its token. + +::: lancedb.listing.Listing + ::: lancedb.Session ## Catalogs (Synchronous) @@ -423,6 +431,8 @@ can be used to create, list, or open tables. ::: lancedb.db.AsyncConnection +::: lancedb.listing.AsyncListing + ## Namespaces (Asynchronous) ::: lancedb.connect_namespace_async diff --git a/nodejs/__test__/materialized_view.test.ts b/nodejs/__test__/materialized_view.test.ts index f0c97e13a..3e47a00e9 100644 --- a/nodejs/__test__/materialized_view.test.ts +++ b/nodejs/__test__/materialized_view.test.ts @@ -112,7 +112,7 @@ describe("materialized views", () => { await db.createMaterializedView("adults", "people", { where: "age >= 18", }); - expect(await db.listMaterializedViews()).toEqual(["adults"]); + expect(await collect(db.listMaterializedViews())).toEqual(["adults"]); await expect(db.openMaterializedView("people")).rejects.toThrow( "not a materialized view", ); @@ -121,7 +121,7 @@ describe("materialized views", () => { ); await db.dropMaterializedView("adults"); - expect(await db.listMaterializedViews()).toEqual([]); + expect(await collect(db.listMaterializedViews())).toEqual([]); }); it("returns a job when dropping a view asynchronously", async () => { @@ -130,7 +130,7 @@ describe("materialized views", () => { const job = await db.dropMaterializedViewAsync("adults"); expect(job.id).toBeNull(); await job.wait(); - expect(await db.listMaterializedViews()).toEqual([]); + expect(await collect(db.listMaterializedViews())).toEqual([]); }); it("rejects an invalid expression at create time", async () => { @@ -147,7 +147,7 @@ describe("materialized views", () => { db.createMaterializedView("bad", "people", { limit }), ).rejects.toThrow("non-negative integer"); } - expect(await db.listMaterializedViews()).toEqual([]); + expect(await collect(db.listMaterializedViews())).toEqual([]); const view = await db.createMaterializedView("copy", "people"); for (const sourceVersion of [-1, 1.5, Infinity, NaN]) { @@ -173,3 +173,9 @@ describe("materialized views", () => { expect(await view.table().countRows()).toBe(1); }); }); + +async function collect(items: AsyncIterable): Promise { + const result: T[] = []; + for await (const item of items) result.push(item); + return result; +} diff --git a/nodejs/__test__/remote.test.ts b/nodejs/__test__/remote.test.ts index c3e7968aa..3f2d4a436 100644 --- a/nodejs/__test__/remote.test.ts +++ b/nodejs/__test__/remote.test.ts @@ -258,7 +258,9 @@ describe("remote connection", () => { .end(JSON.stringify({ views: ["daily_sales"] })); }, async (db) => { - expect(await db.listMaterializedViews()).toEqual(["daily_sales"]); + expect(await collect(db.listMaterializedViews())).toEqual([ + "daily_sales", + ]); }, ); }); @@ -308,7 +310,7 @@ describe("remote connection", () => { .end(JSON.stringify({ views: ["adults"] })); }, async (db) => { - expect(await db.listViews()).toEqual(["adults"]); + expect(await collect(db.listViews())).toEqual(["adults"]); }, ); @@ -1377,7 +1379,7 @@ describe("remote connection jobs surface", () => { }); }, async (db) => { - const jobs = await db.listJobs(); + const jobs = await collect(db.listJobs()); expect(jobs.map((job) => job.jobId)).toEqual(["job-1", "job-2"]); expect(jobs[0].state).toEqual("running"); expect(jobs[1].state).toEqual("finished"); @@ -1473,3 +1475,147 @@ describe("remote connection jobs surface", () => { ); }); }); + +async function collect(items: AsyncIterable): Promise { + const result: T[] = []; + for await (const item of items) result.push(item); + return result; +} + +describe("lazy resource listings", () => { + for (const kind of ["views", "materializedViews", "jobs"] as const) { + const listing = ( + db: Connection, + options: { pageToken?: string; pageLimit?: number } = {}, + ) => { + if (kind === "jobs") return db.listJobs(options); + if (kind === "views") return db.listViews(["team"], options); + return db.listMaterializedViews(options); + }; + const page = (names: string[], token?: string) => { + const items = + kind === "jobs" + ? names.map((name) => ({ + // biome-ignore lint/style/useNamingConvention: server wire format + job_id: name, + table: "t", + // biome-ignore lint/style/useNamingConvention: server wire format + job_type: "create_index", + state: "in_progress", + // biome-ignore lint/style/useNamingConvention: server wire format + created_at_millis: 1, + })) + : names; + return { + [kind === "jobs" ? "jobs" : "views"]: items, + // biome-ignore lint/style/useNamingConvention: server wire format + page_token: token, + }; + }; + + it(`${kind} fetches lazily, skips empty pages, and resumes`, async () => { + const requests: { url: string; body: Record }[] = []; + const responses = [ + page(["a", "b"], "empty"), + page([], "last"), + page(["c", "d"], ""), + ]; + await withMockDatabase( + async (req, res) => { + const chunks: Buffer[] = []; + for await (const chunk of req) chunks.push(Buffer.from(chunk)); + const body = Buffer.concat(chunks).toString(); + requests.push({ + url: req.url ?? "", + body: body ? JSON.parse(body) : {}, + }); + res + .writeHead(200, { "content-type": "application/json" }) + .end(JSON.stringify(responses.shift())); + }, + async (db) => { + const items = listing(db, { pageToken: "start/token", pageLimit: 2 }); + expect(requests).toHaveLength(0); + expect(items.pageToken()).toBe("start/token"); + expect(items.numPageResults()).toBe(0); + expect((await items.next()).done).toBe(false); + expect(items.numPageResults()).toBe(1); + expect(items.pageToken()).toBe("empty"); + await items.next(); + expect(requests).toHaveLength(1); + const resumed = listing(db, { + pageToken: items.pageToken(), + pageLimit: 2, + }); + await resumed.next(); + expect(resumed.numPageResults()).toBe(1); + expect(resumed.pageToken()).toBeUndefined(); + expect(requests).toHaveLength(3); + await resumed.next(); + expect((await resumed.next()).done).toBe(true); + expect((await resumed.next()).done).toBe(true); + for (const [i, request] of requests.entries()) { + const token = ["start/token", "empty", "last"][i]; + const url = new URL(request.url, "http://localhost"); + if (kind === "jobs") { + expect(url.pathname).toBe("/v1/jobs/list"); + expect(request.body).toEqual({ + // biome-ignore lint/style/useNamingConvention: server wire format + page_token: token, + limit: 2, + }); + } else { + expect(url.pathname).toBe( + kind === "views" + ? "/v1/namespace/team/view/list" + : "/v1/namespace/$/materialized_view/list", + ); + expect(url.searchParams.get("page_token")).toBe(token); + expect(url.searchParams.get("limit")).toBe("2"); + } + } + }, + ); + }); + + it(`${kind} retains failed tokens and serializes concurrent advances`, async () => { + let requests = 0; + await withMockDatabase( + (_req, res) => { + requests += 1; + res + .writeHead(requests === 1 ? 401 : 200, { + "content-type": "application/json", + }) + .end(JSON.stringify(requests === 1 ? {} : page(["a", "b"]))); + }, + async (db) => { + const failed = listing(db, { pageToken: "retry" }); + await expect(failed.next()).rejects.toThrow(); + expect(failed.pageToken()).toBe("retry"); + expect((await failed.next()).done).toBe(true); + expect(requests).toBe(1); + const resumed = listing(db, { pageToken: failed.pageToken() }); + const results = await Promise.all([ + resumed.next(), + resumed.next(), + resumed.next(), + ]); + expect(results.filter((result) => result.done)).toHaveLength(1); + expect(requests).toBe(2); + for (const limit of [ + 0, + -1, + 1.5, + Number.NaN, + Number.POSITIVE_INFINITY, + 2147483648, + ]) { + expect(() => listing(db, { pageLimit: limit })).toThrow("limit"); + } + expect(requests).toBe(2); + }, + ); + }); + } +}); diff --git a/nodejs/lancedb/connection.ts b/nodejs/lancedb/connection.ts index 708f10fb4..5ba531eed 100644 --- a/nodejs/lancedb/connection.ts +++ b/nodejs/lancedb/connection.ts @@ -17,6 +17,7 @@ import { } from "./arrow"; import { EmbeddingFunctionConfig, getRegistry } from "./embedding/registry"; import { Job } from "./job"; +import { Listing, ListingOptions, validateListingOptions } from "./listing"; import { MaterializedView, MaterializedViewSelect, @@ -347,9 +348,11 @@ export abstract class Connection { /** * The names of the materialized views in this database. * - * Found by reading every table's schema, so this costs an open per table. + * Iterate lazily with `for await...of`. Local listings inspect table schemas + * one page at a time. Options control page size and the starting token. + * See {@link Listing} for pagination state and error behavior. */ - abstract listMaterializedViews(): Promise; + abstract listMaterializedViews(options?: ListingOptions): Listing; /** * Drop the materialized view named `name`. @@ -421,8 +424,13 @@ export abstract class Connection { * The names of the views in one namespace. * * Names only; a definition comes from {@link describeView}. + * Iterate lazily with `for await...of`. See {@link Listing} for pagination + * state and error behavior; options control page size and the starting token. */ - abstract listViews(namespacePath?: string[]): Promise; + abstract listViews( + namespacePath?: string[], + options?: ListingOptions, + ): Listing; abstract openTable( name: string, @@ -637,8 +645,12 @@ export abstract class Connection { */ abstract openJob(jobId: string): Promise; - /** List server-side jobs across the database's tables. */ - abstract listJobs(): Promise; + /** + * Iterate lazily over server-side jobs across the database's tables. + * Use `for await...of`. Options control page size and the starting token. + * See {@link Listing} for pagination state and error behavior. + */ + abstract listJobs(options?: ListingOptions): Listing; /** * Request cancellation of a server-side job by id. @@ -741,8 +753,11 @@ export class LocalConnection extends Connection { return new MaterializedView(new LocalTable(innerTable)); } - async listMaterializedViews(): Promise { - return await this.inner.listMaterializedViews(); + listMaterializedViews(options: ListingOptions = {}): Listing { + validateListingOptions(options); + return new Listing( + this.inner.listMaterializedViews(options.pageToken, options.pageLimit), + ); } async dropMaterializedView( @@ -788,8 +803,18 @@ export class LocalConnection extends Connection { return new Job(await this.inner.dropViewAsync(name, namespacePath ?? [])); } - async listViews(namespacePath?: string[]): Promise { - return this.inner.listViews(namespacePath ?? []); + listViews( + namespacePath?: string[], + options: ListingOptions = {}, + ): Listing { + validateListingOptions(options); + return new Listing( + this.inner.listViews( + namespacePath ?? [], + options.pageToken, + options.pageLimit, + ), + ); } async listTables( @@ -1057,8 +1082,11 @@ export class LocalConnection extends Connection { return new Job(await this.inner.openJob(jobId)); } - async listJobs(): Promise { - return this.inner.listJobs(); + listJobs(options: ListingOptions = {}): Listing { + validateListingOptions(options); + return new Listing( + this.inner.listJobs(options.pageToken, options.pageLimit), + ); } async cancelJob(jobId: string): Promise { diff --git a/nodejs/lancedb/index.ts b/nodejs/lancedb/index.ts index 694db9fa7..be46e23dc 100644 --- a/nodejs/lancedb/index.ts +++ b/nodejs/lancedb/index.ts @@ -649,3 +649,5 @@ export { DatabaseNames, connectCatalog, } from "./catalog"; + +export { Listing, ListingOptions } from "./listing"; diff --git a/nodejs/lancedb/listing.ts b/nodejs/lancedb/listing.ts new file mode 100644 index 000000000..87333a28f --- /dev/null +++ b/nodejs/lancedb/listing.ts @@ -0,0 +1,73 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The LanceDB Authors + +/** Pagination controls for lazy resource listings. */ +export interface ListingOptions { + /** Starting continuation token; omitted starts at the beginning. */ + pageToken?: string; + /** Maximum results per request, not a limit on the whole listing. */ + pageLimit?: number; +} + +/** @internal */ +export function validateListingOptions(options: ListingOptions): void { + if ( + options.pageLimit !== undefined && + (!Number.isInteger(options.pageLimit) || + options.pageLimit <= 0 || + options.pageLimit > 2147483647) + ) { + throw new Error( + "Listing page limit must be an integer between 1 and 2147483647", + ); + } +} + +/** + * A lazy async iterator over database resources with pagination state. + * Requests occur during iteration; request errors terminate the iterator. + * + * @example + * ```ts + * for await (const name of connection.listViews([], { pageLimit: 20 })) { + * console.log(name); + * } + * ``` + */ +export class Listing implements AsyncIterableIterator { + /** @hidden */ + constructor( + private readonly inner: { + next(): Promise; + numPageResults(): number; + pageToken(): string | null | undefined; + }, + ) {} + + [Symbol.asyncIterator](): AsyncIterableIterator { + return this; + } + + /** Fetch the next item, 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 cached items available without another request. */ + numPageResults(): number { + return this.inner.numPageResults(); + } + + /** + * Token for the next request. Initially this is the supplied starting token. + * Returns undefined after fetching the final page, even with cached items. + * Drain the cache before saving a token to avoid skipping items on resumption. + * Failed requests terminate iteration and retain their token for a new iterator. + */ + pageToken(): string | undefined { + return this.inner.pageToken() ?? undefined; + } +} diff --git a/nodejs/src/connection.rs b/nodejs/src/connection.rs index 4a820af5e..58ec9859a 100644 --- a/nodejs/src/connection.rs +++ b/nodejs/src/connection.rs @@ -371,13 +371,15 @@ impl Connection { } #[napi(catch_unwind)] - pub async fn list_materialized_views(&self) -> napi::Result> { - let views = self - .get_inner()? - .list_materialized_views() - .await - .default_error()?; - Ok(views) + pub fn list_materialized_views( + &self, + page_token: Option, + page_limit: Option, + ) -> napi::Result { + Ok(crate::listing::NameListing::new( + self.get_inner()? + .list_materialized_views(crate::listing::options(page_token, page_limit)), + )) } /// Drop a materialized view. @@ -460,12 +462,18 @@ impl Connection { /// The names of the views in one namespace. #[napi(catch_unwind)] - pub async fn list_views( + pub fn list_views( &self, namespace_path: Option>, - ) -> napi::Result> { - let ns = namespace_path.unwrap_or_default(); - self.get_inner()?.list_views(&ns).await.default_error() + page_token: Option, + page_limit: Option, + ) -> napi::Result { + Ok(crate::listing::NameListing::new( + self.get_inner()?.list_views( + &namespace_path.unwrap_or_default(), + crate::listing::options(page_token, page_limit), + ), + )) } /// Start dropping a materialized view and return its cleanup job. @@ -588,9 +596,15 @@ impl Connection { /// List server-side jobs across the database's tables. #[napi(catch_unwind)] - pub async fn list_jobs(&self) -> napi::Result> { - let jobs = self.get_inner()?.list_jobs().await.default_error()?; - Ok(jobs.into_iter().map(Into::into).collect()) + pub fn list_jobs( + &self, + page_token: Option, + page_limit: Option, + ) -> napi::Result { + Ok(crate::listing::JobListing::new( + self.get_inner()? + .list_jobs(crate::listing::options(page_token, page_limit)), + )) } /// Request cancellation of a server-side job by id. Returns true if the diff --git a/nodejs/src/lib.rs b/nodejs/src/lib.rs index 1bb5f761d..f47bc113e 100644 --- a/nodejs/src/lib.rs +++ b/nodejs/src/lib.rs @@ -18,6 +18,7 @@ mod header; mod index; mod iterator; mod job; +mod listing; pub mod merge; pub mod otel; pub mod permutation; diff --git a/nodejs/src/listing.rs b/nodejs/src/listing.rs new file mode 100644 index 000000000..ed2bce9df --- /dev/null +++ b/nodejs/src/listing.rs @@ -0,0 +1,83 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The LanceDB Authors + +use crate::error::NapiErrorExt; +use futures::{StreamExt, future::poll_fn}; +use lancedb::listing::{Listing, ListingOptions}; +use napi_derive::napi; +use std::sync::Mutex; + +pub fn options(page_token: Option, page_limit: Option) -> ListingOptions { + let mut options = ListingOptions::default(); + options.page_token = page_token; + options.page_limit = page_limit; + options +} + +#[napi] +pub struct NameListing { + inner: Mutex>, + next_lock: futures::lock::Mutex<()>, +} +impl NameListing { + pub fn new(inner: Listing) -> Self { + Self { + inner: Mutex::new(inner), + next_lock: futures::lock::Mutex::new(()), + } + } +} +#[napi] +impl NameListing { + #[napi] + pub fn num_page_results(&self) -> u32 { + self.inner.lock().unwrap().num_page_results() as u32 + } + #[napi] + pub fn page_token(&self) -> Option { + self.inner.lock().unwrap().page_token().map(str::to_owned) + } + #[napi] + pub async fn next(&self) -> napi::Result> { + let _guard = self.next_lock.lock().await; + let value = poll_fn(|cx| self.inner.lock().unwrap().poll_next_unpin(cx)) + .await + .transpose() + .default_error()?; + Ok(value) + } +} + +#[napi] +pub struct JobListing { + inner: Mutex>, + next_lock: futures::lock::Mutex<()>, +} +impl JobListing { + pub fn new(inner: Listing) -> Self { + Self { + inner: Mutex::new(inner), + next_lock: futures::lock::Mutex::new(()), + } + } +} +#[napi] +impl JobListing { + #[napi] + pub fn num_page_results(&self) -> u32 { + self.inner.lock().unwrap().num_page_results() as u32 + } + #[napi] + pub fn page_token(&self) -> Option { + self.inner.lock().unwrap().page_token().map(str::to_owned) + } + #[napi] + pub async fn next(&self) -> napi::Result> { + let _guard = self.next_lock.lock().await; + let value = poll_fn(|cx| self.inner.lock().unwrap().poll_next_unpin(cx)) + .await + .transpose() + .default_error()?; + Ok(value.map(|value| value.into())) + } +} diff --git a/python/python/lancedb/_lancedb.pyi b/python/python/lancedb/_lancedb.pyi index 8ff53bcbe..f7477e891 100644 --- a/python/python/lancedb/_lancedb.pyi +++ b/python/python/lancedb/_lancedb.pyi @@ -155,9 +155,13 @@ class Connection(object): async def get_function( self, name: str, version: str, namespace_path: Optional[List[str]] = None ) -> str: ... - async def list_functions( - self, namespace_path: Optional[List[str]] = None - ) -> List[str]: ... + def list_functions( + self, + namespace_path: Optional[List[str]] = None, + *, + page_token: Optional[str] = None, + page_limit: Optional[int] = None, + ) -> FunctionListing: ... async def drop_function( self, name: str, version: str, namespace_path: Optional[List[str]] = None ) -> bool: ... @@ -170,9 +174,13 @@ class Connection(object): async def alter_secret( self, name: str, value: str, namespace_path: Optional[List[str]] = None ) -> None: ... - async def list_secrets( - self, namespace_path: Optional[List[str]] = None - ) -> List[str]: ... + def list_secrets( + self, + namespace_path: Optional[List[str]] = None, + *, + page_token: Optional[str] = None, + page_limit: Optional[int] = None, + ) -> NameListing: ... async def drop_secret( self, name: str, namespace_path: Optional[List[str]] = None ) -> None: ... @@ -191,10 +199,16 @@ class Connection(object): async def drop_view_async( self, name: str, namespace_path: Optional[List[str]] = None ) -> Job: ... - async def list_views( - self, namespace_path: Optional[List[str]] = None - ) -> List[str]: ... - async def list_jobs(self) -> List[JobInfo]: ... + def list_views( + self, + namespace_path: Optional[List[str]] = None, + *, + page_token: Optional[str] = None, + page_limit: Optional[int] = None, + ) -> NameListing: ... + def list_jobs( + self, *, page_token: Optional[str] = None, page_limit: Optional[int] = None + ) -> JobListing: ... async def cancel_job(self, job_id: str) -> bool: ... async def pause_job(self, job_id: str) -> str: ... async def resume_job(self, job_id: str) -> str: ... @@ -265,7 +279,9 @@ class Connection(object): limit: Optional[int] = None, with_no_data: bool = False, ) -> Job: ... - async def list_materialized_views(self) -> List[str]: ... + def list_materialized_views( + self, *, page_token: Optional[str] = None, page_limit: Optional[int] = None + ) -> NameListing: ... async def drop_materialized_view( self, name: str, namespace_path: Optional[List[str]] = None ) -> None: ... @@ -961,3 +977,21 @@ async def connect_catalog( read_consistency_interval: Optional[float] = None, oauth_config: Optional[Any] = None, ) -> Catalog: ... + +class NameListing: + def num_page_results(self) -> int: ... + def page_token(self) -> Optional[str]: ... + def __aiter__(self) -> NameListing: ... + async def __anext__(self) -> str: ... + +class FunctionListing: + def num_page_results(self) -> int: ... + def page_token(self) -> Optional[str]: ... + def __aiter__(self) -> FunctionListing: ... + async def __anext__(self) -> str: ... + +class JobListing: + def num_page_results(self) -> int: ... + def page_token(self) -> Optional[str]: ... + def __aiter__(self) -> JobListing: ... + async def __anext__(self) -> JobInfo: ... diff --git a/python/python/lancedb/db.py b/python/python/lancedb/db.py index 3a9815cb4..f890c5a14 100644 --- a/python/python/lancedb/db.py +++ b/python/python/lancedb/db.py @@ -31,6 +31,7 @@ if sys.version_info >= (3, 12): else: from overrides import EnforceOverrides, override # type: ignore +from lancedb.listing import AsyncListing, Listing from lancedb.embeddings.registry import EmbeddingFunctionRegistry from lancedb.common import data_to_reader, sanitize_uri, validate_schema @@ -668,11 +669,16 @@ class DBConnection(EnforceOverrides): "materialized views are not supported on this connection type" ) - def list_materialized_views(self) -> List[str]: + def list_materialized_views( + self, *, page_token: Optional[str] = None, page_limit: Optional[int] = None + ) -> Listing[str]: """The names of the materialized views in this database. - Found by reading every table's schema, so this costs an open per - table. + Local listings inspect table schemas one page at a time. + + Returns a lazy iterator. ``page_limit`` limits each request, not the total; + ``page_token`` resumes from a saved token. Requests and errors occur during + iteration. See [pagination state][lancedb.listing.AsyncListing]. """ raise NotImplementedError( "materialized views are not supported on this connection type" @@ -885,8 +891,12 @@ class DBConnection(EnforceOverrides): ) def list_functions( - self, *, namespace_path: Optional[List[str]] = None - ) -> List[FunctionVersion]: + self, + *, + namespace_path: Optional[List[str]] = None, + page_token: Optional[str] = None, + page_limit: Optional[int] = None, + ) -> Listing[FunctionVersion]: """List every published immutable Function version in a namespace. Functions in child namespaces are not included. Results are ordered by @@ -900,6 +910,10 @@ class DBConnection(EnforceOverrides): ```python [(function.name, function.version) for function in db.list_functions()] ``` + + Returns a lazy iterator. ``page_limit`` limits each request, not the total; + ``page_token`` resumes from a saved token. Requests and errors occur during + iteration. See [pagination state][lancedb.listing.AsyncListing]. """ raise NotImplementedError( "Function catalog operations are not supported for this connection type" @@ -971,11 +985,21 @@ class DBConnection(EnforceOverrides): "Secret operations are not supported for this connection type" ) - def list_secrets(self, *, namespace_path: Optional[List[str]] = None) -> List[str]: + def list_secrets( + self, + *, + namespace_path: Optional[List[str]] = None, + page_token: Optional[str] = None, + page_limit: Optional[int] = None, + ) -> Listing[str]: """The names of every Secret in this database. Names only. No method returns a stored credential, by construction rather than by policy. Local connections raise ``NotImplementedError``. + + Returns a lazy iterator. ``page_limit`` limits each request, not the total; + ``page_token`` resumes from a saved token. Requests and errors occur during + iteration. See [pagination state][lancedb.listing.AsyncListing]. """ raise NotImplementedError( "Secret operations are not supported for this connection type" @@ -1073,12 +1097,22 @@ class DBConnection(EnforceOverrides): "View operations are not supported for this connection type" ) - def list_views(self, *, namespace_path: Optional[List[str]] = None) -> List[str]: + def list_views( + self, + *, + namespace_path: Optional[List[str]] = None, + page_token: Optional[str] = None, + page_limit: Optional[int] = None, + ) -> Listing[str]: """The names of the views in one namespace. Names only; a definition comes from [describe_view][lancedb.db.DBConnection.describe_view]. Local connections raise ``NotImplementedError``. + + Returns a lazy iterator. ``page_limit`` limits each request, not the total; + ``page_token`` resumes from a saved token. Requests and errors occur during + iteration. See [pagination state][lancedb.listing.AsyncListing]. """ raise NotImplementedError( "View operations are not supported for this connection type" @@ -1097,8 +1131,15 @@ class DBConnection(EnforceOverrides): """ raise NotImplementedError("open_job is not supported for this connection type") - def list_jobs(self) -> List[JobInfo]: - """List server-side jobs across the database's tables.""" + def list_jobs( + self, *, page_token: Optional[str] = None, page_limit: Optional[int] = None + ) -> Listing[JobInfo]: + """List server-side jobs across the database's tables. + + Returns a lazy iterator. ``page_limit`` limits each request, not the total; + ``page_token`` resumes from a saved token. Requests and errors occur during + iteration. See [pagination state][lancedb.listing.AsyncListing]. + """ raise NotImplementedError("list_jobs is not supported for this connection type") def cancel_job(self, job_id: str) -> bool: @@ -1699,9 +1740,20 @@ class LanceDBConnection(DBConnection): return view @override - def list_materialized_views(self) -> List[str]: - """The names of the materialized views in this database.""" - return LOOP.run(self._conn.list_materialized_views()) + def list_materialized_views( + self, *, page_token: Optional[str] = None, page_limit: Optional[int] = None + ) -> Listing[str]: + """The names of the materialized views in this database. + + Returns a lazy iterator. ``page_limit`` limits each request, not the total; + ``page_token`` resumes from a saved token. Requests and errors occur during + iteration. See [pagination state][lancedb.listing.AsyncListing]. + """ + return Listing( + self._conn.list_materialized_views( + page_token=page_token, page_limit=page_limit + ) + ) @override def drop_materialized_view( @@ -1896,9 +1948,25 @@ class LanceDBConnection(DBConnection): @override def list_functions( - self, *, namespace_path: Optional[List[str]] = None - ) -> List[FunctionVersion]: - return LOOP.run(self._conn.list_functions(namespace_path=namespace_path)) + self, + *, + namespace_path: Optional[List[str]] = None, + page_token: Optional[str] = None, + page_limit: Optional[int] = None, + ) -> Listing[FunctionVersion]: + """List functions. + + Returns a lazy iterator. ``page_limit`` limits each request, not the total; + ``page_token`` resumes from a saved token. Requests and errors occur during + iteration. See [pagination state][lancedb.listing.AsyncListing]. + """ + return Listing( + self._conn.list_functions( + namespace_path=namespace_path, + page_token=page_token, + page_limit=page_limit, + ) + ) @override def drop_function( @@ -1942,8 +2010,26 @@ class LanceDBConnection(DBConnection): LOOP.run(self._conn.alter_secret(name, value, namespace_path=namespace_path)) @override - def list_secrets(self, *, namespace_path: Optional[List[str]] = None) -> List[str]: - return LOOP.run(self._conn.list_secrets(namespace_path=namespace_path)) + def list_secrets( + self, + *, + namespace_path: Optional[List[str]] = None, + page_token: Optional[str] = None, + page_limit: Optional[int] = None, + ) -> Listing[str]: + """List secrets. + + Returns a lazy iterator. ``page_limit`` limits each request, not the total; + ``page_token`` resumes from a saved token. Requests and errors occur during + iteration. See [pagination state][lancedb.listing.AsyncListing]. + """ + return Listing( + self._conn.list_secrets( + namespace_path=namespace_path, + page_token=page_token, + page_limit=page_limit, + ) + ) @override def drop_secret( @@ -1986,13 +2072,40 @@ class LanceDBConnection(DBConnection): ) @override - def list_views(self, *, namespace_path: Optional[List[str]] = None) -> List[str]: - return LOOP.run(self._conn.list_views(namespace_path=namespace_path)) + def list_views( + self, + *, + namespace_path: Optional[List[str]] = None, + page_token: Optional[str] = None, + page_limit: Optional[int] = None, + ) -> Listing[str]: + """List views. + + Returns a lazy iterator. ``page_limit`` limits each request, not the total; + ``page_token`` resumes from a saved token. Requests and errors occur during + iteration. See [pagination state][lancedb.listing.AsyncListing]. + """ + return Listing( + self._conn.list_views( + namespace_path=namespace_path, + page_token=page_token, + page_limit=page_limit, + ) + ) @override - def list_jobs(self) -> List[JobInfo]: - """List server-side jobs across the database's tables.""" - return LOOP.run(self._conn.list_jobs()) + def list_jobs( + self, *, page_token: Optional[str] = None, page_limit: Optional[int] = None + ) -> Listing[JobInfo]: + """List server-side jobs across the database's tables. + + Returns a lazy iterator. ``page_limit`` limits each request, not the total; + ``page_token`` resumes from a saved token. Requests and errors occur during + iteration. See [pagination state][lancedb.listing.AsyncListing]. + """ + return Listing( + self._conn.list_jobs(page_token=page_token, page_limit=page_limit) + ) @override def cancel_job(self, job_id: str) -> bool: @@ -2680,13 +2793,22 @@ class AsyncConnection(object): await view.definition() return view - async def list_materialized_views(self) -> List[str]: + def list_materialized_views( + self, *, page_token: Optional[str] = None, page_limit: Optional[int] = None + ) -> AsyncListing[str]: """The names of the materialized views in this database. - Found by reading every table's schema, so this costs an open per - table. + Local listings inspect table schemas one page at a time. + + Returns a lazy iterator. ``page_limit`` limits each request, not the total; + ``page_token`` resumes from a saved token. Requests and errors occur during + iteration. See [pagination state][lancedb.listing.AsyncListing]. """ - return await self._inner.list_materialized_views() + return AsyncListing( + self._inner.list_materialized_views( + page_token=page_token, page_limit=page_limit + ) + ) async def drop_materialized_view( self, @@ -2908,19 +3030,31 @@ class AsyncConnection(object): await self._inner.get_function(name, version, namespace_path) ) - async def list_functions( - self, *, namespace_path: Optional[List[str]] = None - ) -> List[FunctionVersion]: + def list_functions( + self, + *, + namespace_path: Optional[List[str]] = None, + page_token: Optional[str] = None, + page_limit: Optional[int] = None, + ) -> AsyncListing[FunctionVersion]: """List every published immutable Function version in a namespace. Functions in child namespaces are not included. Results are ordered by Function name then version. Local connections raise ``NotImplementedError``. + + Returns a lazy iterator. ``page_limit`` limits each request, not the total; + ``page_token`` resumes from a saved token. Requests and errors occur during + iteration. See [pagination state][lancedb.listing.AsyncListing]. """ - return [ - FunctionVersion.from_json(value) - for value in await self._inner.list_functions(namespace_path) - ] + return AsyncListing( + self._inner.list_functions( + namespace_path=namespace_path, + page_token=page_token, + page_limit=page_limit, + ), + FunctionVersion.from_json, + ) async def drop_function( self, @@ -2978,12 +3112,25 @@ class AsyncConnection(object): list(validate_namespace_path(namespace_path)), ) - async def list_secrets( - self, *, namespace_path: Optional[List[str]] = None - ) -> List[str]: - """The names of every Secret in this database. Names only.""" - return await self._inner.list_secrets( - list(validate_namespace_path(namespace_path)) + def list_secrets( + self, + *, + namespace_path: Optional[List[str]] = None, + page_token: Optional[str] = None, + page_limit: Optional[int] = None, + ) -> AsyncListing[str]: + """The names of every Secret in this database. Names only. + + Returns a lazy iterator. ``page_limit`` limits each request, not the total; + ``page_token`` resumes from a saved token. Requests and errors occur during + iteration. See [pagination state][lancedb.listing.AsyncListing]. + """ + return AsyncListing( + self._inner.list_secrets( + namespace_path=namespace_path, + page_token=page_token, + page_limit=page_limit, + ) ) async def drop_secret( @@ -3055,15 +3202,39 @@ class AsyncConnection(object): await self._inner.drop_view_async(name, namespace_path=namespace_path) ) - async def list_views( - self, *, namespace_path: Optional[List[str]] = None - ) -> List[str]: - """The names of the views in one namespace.""" - return await self._inner.list_views(list(namespace_path or [])) + def list_views( + self, + *, + namespace_path: Optional[List[str]] = None, + page_token: Optional[str] = None, + page_limit: Optional[int] = None, + ) -> AsyncListing[str]: + """The names of the views in one namespace. - async def list_jobs(self) -> List[JobInfo]: - """List server-side jobs across the database's tables.""" - return await self._inner.list_jobs() + Returns a lazy iterator. ``page_limit`` limits each request, not the total; + ``page_token`` resumes from a saved token. Requests and errors occur during + iteration. See [pagination state][lancedb.listing.AsyncListing]. + """ + return AsyncListing( + self._inner.list_views( + namespace_path=namespace_path, + page_token=page_token, + page_limit=page_limit, + ) + ) + + def list_jobs( + self, *, page_token: Optional[str] = None, page_limit: Optional[int] = None + ) -> AsyncListing[JobInfo]: + """List server-side jobs across the database's tables. + + Returns a lazy iterator. ``page_limit`` limits each request, not the total; + ``page_token`` resumes from a saved token. Requests and errors occur during + iteration. See [pagination state][lancedb.listing.AsyncListing]. + """ + return AsyncListing( + self._inner.list_jobs(page_token=page_token, page_limit=page_limit) + ) async def cancel_job(self, job_id: str) -> bool: """Request cancellation of a server-side job by id. diff --git a/python/python/lancedb/listing.py b/python/python/lancedb/listing.py new file mode 100644 index 000000000..0620aa7d8 --- /dev/null +++ b/python/python/lancedb/listing.py @@ -0,0 +1,85 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright The LanceDB Authors + +"""Lazy, resumable resource listings.""" + +from __future__ import annotations + +from typing import Any, AsyncIterator, Callable, Generic, Iterator, Optional, TypeVar + +from .background_loop import LOOP + +T = TypeVar("T") + + +class AsyncListing(AsyncIterator[T], Generic[T]): + """A lazy async iterator with inspectable pagination state. + + Returned by connection resource-listing methods. For example:: + + async for name in connection.list_views(page_limit=20): + print(name) + + Requests happen during iteration; request errors terminate the iterator. + """ + + def __init__(self, inner: Any, convert: Callable[[Any], T] = lambda value: value): + self._inner = inner + self._convert = convert + + def __aiter__(self) -> AsyncListing[T]: + return self + + async def __anext__(self) -> T: + return self._convert(await self._inner.__anext__()) + + def num_page_results(self) -> int: + """Return the number of cached items available without another request.""" + return self._inner.num_page_results() + + def page_token(self) -> Optional[str]: + """Return the token for the next request. + + Initially this is the supplied starting token (``None`` starts at the + beginning). After fetching the final page it is ``None``, even if items + remain cached. Drain the cache before saving a token to avoid skipping + items when resuming. A failed request terminates iteration and retains + its token for retrying with a new iterator. + """ + return self._inner.page_token() + + +class Listing(Iterator[T], Generic[T]): + """A lazy synchronous iterator with inspectable pagination state. + + For example:: + + for name in connection.list_views(page_limit=20): + print(name) + + Use ``list(connection.list_views())`` to collect all results. + """ + + def __init__(self, inner: AsyncListing[T]): + self._inner = inner + + def __iter__(self) -> Listing[T]: + return self + + def __next__(self) -> T: + try: + return LOOP.run(self._inner.__anext__()) + except StopAsyncIteration: + raise StopIteration from None + + def num_page_results(self) -> int: + """Return the number of cached items available without another request.""" + return self._inner.num_page_results() + + def page_token(self) -> Optional[str]: + """Return the next request's token. + + See [AsyncListing.page_token][lancedb.listing.AsyncListing.page_token] + for initial, cached, final-page, and error behavior. + """ + return self._inner.page_token() diff --git a/python/python/lancedb/namespace.py b/python/python/lancedb/namespace.py index 16be63e2d..a3571f64b 100644 --- a/python/python/lancedb/namespace.py +++ b/python/python/lancedb/namespace.py @@ -11,7 +11,7 @@ through a namespace abstraction. from __future__ import annotations import sys -from typing import TYPE_CHECKING, Any, Dict, Iterable, List, Optional, Union +from typing import Any, Dict, Iterable, List, Optional, Union from uuid import UUID if sys.version_info >= (3, 12): @@ -19,8 +19,7 @@ if sys.version_info >= (3, 12): else: from overrides import override -if TYPE_CHECKING: - from lancedb.query import Query +from lancedb.query import Query from datetime import timedelta import pyarrow as pa @@ -52,6 +51,7 @@ from lancedb.background_loop import LOOP from lancedb.arrow import AsyncRecordBatchReader from lancedb.db import AsyncConnection, DBConnection from lancedb.job import AsyncJob, Job +from lancedb.listing import AsyncListing, Listing from lancedb.sql import AsyncQuery as AsyncSqlQuery from lancedb.sql import QueryDescription from lance_namespace import ( @@ -698,9 +698,20 @@ class LanceNamespaceDBConnection(DBConnection): return view @override - def list_materialized_views(self) -> List[str]: - """The names of the materialized views in the root namespace.""" - return LOOP.run(self._inner.list_materialized_views()) + def list_materialized_views( + self, *, page_token: Optional[str] = None, page_limit: Optional[int] = None + ) -> Listing[str]: + """The names of the materialized views in the root namespace. + + Returns a lazy iterator. ``page_limit`` limits each request, not the total; + ``page_token`` resumes from a saved token. Requests and errors occur during + iteration. See [pagination state][lancedb.listing.AsyncListing]. + """ + return Listing( + self._inner.list_materialized_views( + page_token=page_token, page_limit=page_limit + ) + ) @override def drop_materialized_view( @@ -1297,9 +1308,20 @@ class AsyncLanceNamespaceDBConnection: await view.definition() return view - async def list_materialized_views(self) -> List[str]: - """The names of the materialized views in the root namespace.""" - return await self._inner.list_materialized_views() + def list_materialized_views( + self, *, page_token: Optional[str] = None, page_limit: Optional[int] = None + ) -> AsyncListing[str]: + """The names of the materialized views in the root namespace. + + Returns a lazy iterator. ``page_limit`` limits each request, not the total; + ``page_token`` resumes from a saved token. Requests and errors occur during + iteration. See [pagination state][lancedb.listing.AsyncListing]. + """ + return AsyncListing( + self._inner.list_materialized_views( + page_token=page_token, page_limit=page_limit + ) + ) async def drop_materialized_view( self, name: str, namespace_path: Optional[List[str]] = None diff --git a/python/python/lancedb/remote/db.py b/python/python/lancedb/remote/db.py index 135bca8fd..2688c3efc 100644 --- a/python/python/lancedb/remote/db.py +++ b/python/python/lancedb/remote/db.py @@ -2,6 +2,8 @@ # SPDX-FileCopyrightText: Copyright The LanceDB Authors +from __future__ import annotations + from dataclasses import replace from datetime import timedelta import json @@ -34,6 +36,7 @@ from lancedb.remote import ClientConfig, RetryConfig, TimeoutConfig, TlsConfig import pyarrow as pa from ..common import DATA +from lancedb.listing import Listing from ..db import DBConnection, LOOP from ..functions import FunctionVersion, UdfDefinition from ..job import AsyncJob, Job @@ -769,8 +772,20 @@ class RemoteDBConnection(DBConnection): return view @override - def list_materialized_views(self) -> List[str]: - return LOOP.run(self._conn.list_materialized_views()) + def list_materialized_views( + self, *, page_token: Optional[str] = None, page_limit: Optional[int] = None + ) -> Listing[str]: + """List materialized views. + + Returns a lazy iterator. ``page_limit`` limits each request, not the total; + ``page_token`` resumes from a saved token. Requests and errors occur during + iteration. See [pagination state][lancedb.listing.AsyncListing]. + """ + return Listing( + self._conn.list_materialized_views( + page_token=page_token, page_limit=page_limit + ) + ) @override def drop_materialized_view( @@ -885,9 +900,25 @@ class RemoteDBConnection(DBConnection): @override def list_functions( - self, *, namespace_path: Optional[List[str]] = None - ) -> List[FunctionVersion]: - return LOOP.run(self._conn.list_functions(namespace_path=namespace_path)) + self, + *, + namespace_path: Optional[List[str]] = None, + page_token: Optional[str] = None, + page_limit: Optional[int] = None, + ) -> Listing[FunctionVersion]: + """List functions. + + Returns a lazy iterator. ``page_limit`` limits each request, not the total; + ``page_token`` resumes from a saved token. Requests and errors occur during + iteration. See [pagination state][lancedb.listing.AsyncListing]. + """ + return Listing( + self._conn.list_functions( + namespace_path=namespace_path, + page_token=page_token, + page_limit=page_limit, + ) + ) @override def drop_function( @@ -937,8 +968,26 @@ class RemoteDBConnection(DBConnection): return LOOP.run(self._conn.describe_secret(name, namespace_path=namespace_path)) @override - def list_secrets(self, *, namespace_path: Optional[List[str]] = None) -> List[str]: - return LOOP.run(self._conn.list_secrets(namespace_path=namespace_path)) + def list_secrets( + self, + *, + namespace_path: Optional[List[str]] = None, + page_token: Optional[str] = None, + page_limit: Optional[int] = None, + ) -> Listing[str]: + """List secrets. + + Returns a lazy iterator. ``page_limit`` limits each request, not the total; + ``page_token`` resumes from a saved token. Requests and errors occur during + iteration. See [pagination state][lancedb.listing.AsyncListing]. + """ + return Listing( + self._conn.list_secrets( + namespace_path=namespace_path, + page_token=page_token, + page_limit=page_limit, + ) + ) @override def drop_secret( @@ -974,13 +1023,40 @@ class RemoteDBConnection(DBConnection): return Job(job) @override - def list_views(self, *, namespace_path: Optional[List[str]] = None) -> List[str]: - return LOOP.run(self._conn.list_views(namespace_path=namespace_path)) + def list_views( + self, + *, + namespace_path: Optional[List[str]] = None, + page_token: Optional[str] = None, + page_limit: Optional[int] = None, + ) -> Listing[str]: + """List views. + + Returns a lazy iterator. ``page_limit`` limits each request, not the total; + ``page_token`` resumes from a saved token. Requests and errors occur during + iteration. See [pagination state][lancedb.listing.AsyncListing]. + """ + return Listing( + self._conn.list_views( + namespace_path=namespace_path, + page_token=page_token, + page_limit=page_limit, + ) + ) @override - def list_jobs(self) -> List["JobInfo"]: - """List server-side jobs across the database's tables.""" - return LOOP.run(self._conn.list_jobs()) + def list_jobs( + self, *, page_token: Optional[str] = None, page_limit: Optional[int] = None + ) -> Listing[JobInfo]: + """List server-side jobs across the database's tables. + + Returns a lazy iterator. ``page_limit`` limits each request, not the total; + ``page_token`` resumes from a saved token. Requests and errors occur during + iteration. See [pagination state][lancedb.listing.AsyncListing]. + """ + return Listing( + self._conn.list_jobs(page_token=page_token, page_limit=page_limit) + ) @override def cancel_job(self, job_id: str) -> bool: diff --git a/python/python/tests/test_first_class_function_slice2.py b/python/python/tests/test_first_class_function_slice2.py index 4610170f6..c18893f1b 100644 --- a/python/python/tests/test_first_class_function_slice2.py +++ b/python/python/tests/test_first_class_function_slice2.py @@ -1499,7 +1499,7 @@ def test_local_function_catalog_operations_are_not_supported(tmp_path): with pytest.raises(NotImplementedError, match=message): db.get_function("normalize_score", version=FUNCTION_VERSION) with pytest.raises(NotImplementedError, match=message): - db.list_functions() + list(db.list_functions()) with pytest.raises(NotImplementedError, match=message): db.drop_function("normalize_score", version=FUNCTION_VERSION) @@ -1720,7 +1720,7 @@ def test_remote_secret_verbs_round_trip(): ) assert db.create_secret("openai-prod", "sk-live-0001") is None assert db.alter_secret("openai-prod", "sk-live-0002") is None - assert db.list_secrets() == ["openai-prod", "hf-prod"] + assert list(db.list_secrets()) == ["openai-prod", "hf-prod"] assert db.drop_secret("openai-prod") is None routes = [path for path, _ in state["requests"]] @@ -1781,7 +1781,7 @@ def test_remote_list_functions_paginates_and_returns_typed_versions(): ) created = db.create_function(normalize_score) state["requests"].clear() - functions = db.list_functions() + functions = list(db.list_functions()) assert functions == [created] assert state["requests"] == [ @@ -1805,7 +1805,7 @@ async def test_async_remote_list_functions_returns_typed_versions(): registration = await db.create_function_async(normalize_score) created = await registration.wait() state["requests"].clear() - functions = await db.list_functions() + functions = [item async for item in db.list_functions()] assert functions == [created] assert [path for path, _ in state["requests"]] == [ @@ -1878,7 +1878,7 @@ def test_remote_function_verbs_address_a_namespace_in_the_path(): assert reopened(value=lancedb.col("score")).function.namespace_path == tuple( namespace ) - db.list_functions(namespace_path=namespace) + list(db.list_functions(namespace_path=namespace)) assert ( db.drop_function( "normalize_score", version=FUNCTION_VERSION, namespace_path=namespace @@ -1918,7 +1918,7 @@ async def test_async_remote_function_verbs_address_a_namespace_in_the_path(): await db.get_function( "normalize_score", version=FUNCTION_VERSION, namespace_path=namespace ) - await db.list_functions(namespace_path=namespace) + [item async for item in db.list_functions(namespace_path=namespace)] assert ( await db.drop_function( "normalize_score", version=FUNCTION_VERSION, namespace_path=namespace @@ -1947,6 +1947,6 @@ def test_a_function_namespace_path_must_be_a_list_of_segments(): client_config={"retry_config": {"retries": 0}}, ) with pytest.raises(TypeError): - db.list_functions(namespace_path="analytics") + list(db.list_functions(namespace_path="analytics")) assert state["requests"] == [] diff --git a/python/python/tests/test_listing.py b/python/python/tests/test_listing.py new file mode 100644 index 000000000..4a2745eb8 --- /dev/null +++ b/python/python/tests/test_listing.py @@ -0,0 +1,194 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright The LanceDB Authors + +import asyncio +from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer +import json +from pathlib import Path +from threading import Thread +from urllib.parse import parse_qs, urlsplit + +import lancedb +from lancedb.remote.errors import HttpError +import pytest + + +@pytest.fixture +def listing_server(): + requests, responses = [], [] + + class Handler(BaseHTTPRequestHandler): + def handle_request(self): + body = self.rfile.read(int(self.headers.get("Content-Length", "0"))) + requests.append((self.command, self.path, json.loads(body) if body else {})) + status, response = responses.pop(0) + self.send_response(status) + self.send_header("Content-Type", "application/json") + self.end_headers() + self.wfile.write(json.dumps(response).encode()) + + do_GET = handle_request + do_POST = handle_request + + with ThreadingHTTPServer(("127.0.0.1", 0), Handler) as server: + thread = Thread(target=server.serve_forever, daemon=True) + thread.start() + try: + yield f"http://127.0.0.1:{server.server_port}", requests, responses + finally: + server.shutdown() + thread.join() + + +def response_page(kind, names, token): + if kind == "secrets": + items = [{"name": name} for name in names] + elif kind == "jobs": + items = [ + { + "job_id": name, + "table": "t", + "job_type": "create_index", + "state": "in_progress", + "created_at_millis": 1, + } + for name in names + ] + elif kind == "functions": + fixture = ( + Path(__file__).resolve().parents[3] + / "rust/lancedb/tests/fixtures/first_class_functions/v1" + / "remote_function_version.canonical.json" + ) + definition = json.loads(fixture.read_text()) + items = [ + {"name": "embed", "version": "1", "definition": definition} for _ in names + ] + else: + items = names + key = "views" if kind == "materialized_views" else kind + return {key: items, "page_token": token} + + +async def connection(endpoint, asynchronous): + kwargs = dict( + api_key="fake", + host_override=endpoint, + client_config={"retry_config": {"retries": 0}}, + ) + if asynchronous: + return await lancedb.connect_async("db://dev", **kwargs) + return lancedb.connect("db://dev", **kwargs) + + +async def advance(items, asynchronous): + return await items.__anext__() if asynchronous else next(items) + + +KINDS = ["secrets", "views", "jobs", "functions", "materialized_views"] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("kind", KINDS) +@pytest.mark.parametrize("asynchronous", [False, True]) +async def test_listing_lazy_pages_and_resume(listing_server, kind, asynchronous): + endpoint, requests, responses = listing_server + responses.extend( + (200, response_page(kind, names, token)) + for names, token in [(["a", "b"], "empty"), ([], "last"), (["c", "d"], "")] + ) + db = await connection(endpoint, asynchronous) + method = getattr(db, "list_" + kind) + kwargs = ( + {} if kind in ("jobs", "materialized_views") else {"namespace_path": ["team"]} + ) + items = method(page_token="start/token", page_limit=2, **kwargs) + assert requests == [] + assert items.page_token() == "start/token" + assert items.num_page_results() == 0 + first = await advance(items, asynchronous) + if kind == "functions": + assert first.name == "embed" + assert first.version == "1" + elif kind == "jobs": + assert first.job_id == "a" + assert first.state == "running" + else: + assert first == "a" + assert items.num_page_results() == 1 + assert items.page_token() == "empty" + await advance(items, asynchronous) + assert items.num_page_results() == 0 + assert len(requests) == 1 + resumed = method(page_token=items.page_token(), page_limit=2, **kwargs) + await advance(resumed, asynchronous) + assert len(requests) == 3 + assert resumed.page_token() is None + assert resumed.num_page_results() == 1 + remaining = [item async for item in resumed] if asynchronous else list(resumed) + assert len(remaining) == 1 + assert resumed.num_page_results() == 0 + assert ([item async for item in resumed] if asynchronous else list(resumed)) == [] + for (verb, url, body), token in zip(requests, ["start/token", "empty", "last"]): + path = urlsplit(url) + if kind == "jobs": + assert verb == "POST" + assert path.path == "/v1/jobs/list" + assert body == {"page_token": token, "limit": 2} + else: + namespace = "$" if kind == "materialized_views" else "team" + resource = kind.removesuffix("s") + assert path.path == f"/v1/namespace/{namespace}/{resource}/list" + expected = {"page_token": [token], "limit": ["2"]} + if kind == "functions": + expected["include_definition"] = ["true"] + assert parse_qs(path.query) == expected + assert responses == [] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("kind", KINDS) +@pytest.mark.parametrize("asynchronous", [False, True]) +async def test_listing_error_retains_token(listing_server, kind, asynchronous): + endpoint, requests, responses = listing_server + responses.extend([(401, {}), (200, response_page(kind, ["a"], None))]) + db = await connection(endpoint, asynchronous) + method = getattr(db, "list_" + kind) + items = method(page_token="retry") + with pytest.raises(HttpError): + await advance(items, asynchronous) + assert items.page_token() == "retry" + assert items.num_page_results() == 0 + assert ([item async for item in items] if asynchronous else list(items)) == [] + assert len(requests) == 1 + resumed = method(page_token=items.page_token()) + assert len([item async for item in resumed] if asynchronous else list(resumed)) == 1 + + +@pytest.mark.asyncio +@pytest.mark.parametrize("kind", KINDS) +async def test_listing_serializes_concurrent_advances(listing_server, kind): + endpoint, requests, responses = listing_server + responses.append((200, response_page(kind, ["a", "b"], None))) + db = await connection(endpoint, True) + items = getattr(db, "list_" + kind)() + results = await asyncio.gather( + items.__anext__(), items.__anext__(), items.__anext__(), return_exceptions=True + ) + assert ( + len([result for result in results if isinstance(result, StopAsyncIteration)]) + == 1 + ) + assert len(requests) == 1 + + +@pytest.mark.asyncio +@pytest.mark.parametrize("kind", KINDS) +@pytest.mark.parametrize("page_limit", [0, 2**32 - 1]) +async def test_listing_invalid_limit_does_not_fetch(listing_server, kind, page_limit): + endpoint, requests, _ = listing_server + db = await connection(endpoint, True) + items = getattr(db, "list_" + kind)(page_limit=page_limit) + with pytest.raises(ValueError, match="limit"): + await items.__anext__() + assert requests == [] diff --git a/python/python/tests/test_materialized_views.py b/python/python/tests/test_materialized_views.py index bef173323..1edd10670 100644 --- a/python/python/tests/test_materialized_views.py +++ b/python/python/tests/test_materialized_views.py @@ -93,7 +93,7 @@ def test_remote_list_uses_namespace_route(): host_override=host, client_config={"retry_config": {"retries": 0}}, ) - assert db.list_materialized_views() == ["daily_sales"] + assert list(db.list_materialized_views()) == ["daily_sales"] assert requests == ["/v1/namespace/$/materialized_view/list"] @@ -190,7 +190,7 @@ def test_create_and_refresh_jobs(tmp_path): drop_job = db.drop_materialized_view_async("adults") assert drop_job.id is None assert drop_job.wait() is None - assert "adults" not in db.list_materialized_views() + assert "adults" not in list(db.list_materialized_views()) def test_definition_round_trips(tmp_path): @@ -250,14 +250,14 @@ def test_list_and_not_a_view(tmp_path): db = make_db(tmp_path) db.create_materialized_view("adults", "people", where="age >= 18") - assert db.list_materialized_views() == ["adults"] + assert list(db.list_materialized_views()) == ["adults"] with pytest.raises(ValueError, match="not a materialized view"): db.open_materialized_view("people") with pytest.raises(ValueError, match="not a materialized view"): db.drop_materialized_view("people") db.drop_materialized_view("adults") - assert db.list_materialized_views() == [] + assert list(db.list_materialized_views()) == [] def test_invalid_expression_fails_at_create(tmp_path): @@ -285,7 +285,7 @@ async def test_async_create_refresh_and_open(tmp_path): reopened = await db.open_materialized_view("shouts") definition = await reopened.definition() assert definition.query == "SELECT upper(name) AS shout FROM people" - assert await db.list_materialized_views() == ["shouts"] + assert [item async for item in db.list_materialized_views()] == ["shouts"] @pytest.mark.asyncio @@ -309,7 +309,7 @@ async def test_async_create_and_refresh_jobs(tmp_path): drop_job = await db.drop_materialized_view_async("adults") assert drop_job.id is None assert await drop_job.wait() is None - assert "adults" not in await db.list_materialized_views() + assert "adults" not in [item async for item in db.list_materialized_views()] @pytest.mark.asyncio @@ -399,7 +399,7 @@ def test_namespace_connection_materialized_views(tmp_path): view = db.create_materialized_view("adults", "people", where="age >= 18") view.refresh() assert view.table.count_rows() == 1 - assert db.list_materialized_views() == ["adults"] + assert list(db.list_materialized_views()) == ["adults"] reopened = db.open_materialized_view("adults") assert reopened.definition.query.startswith("SELECT name, age FROM ") @@ -415,7 +415,7 @@ def test_namespace_connection_materialized_views(tmp_path): assert db.drop_materialized_view_async("job_view").wait() is None db.drop_materialized_view("adults") - assert db.list_materialized_views() == [] + assert list(db.list_materialized_views()) == [] @pytest.mark.asyncio @@ -434,7 +434,7 @@ async def test_async_namespace_connection_materialized_views(tmp_path): view = await db.create_materialized_view("adults", "people", where="age >= 18") await view.refresh() assert await view.table.count_rows() == 1 - assert await db.list_materialized_views() == ["adults"] + assert [item async for item in db.list_materialized_views()] == ["adults"] reopened = await db.open_materialized_view("adults") assert (await reopened.definition()).query.startswith("SELECT name, age FROM ") @@ -460,7 +460,7 @@ async def test_async_namespace_connection_materialized_views(tmp_path): drop_job = await db.drop_materialized_view_async("job_view") assert await drop_job.wait() is None await db.drop_materialized_view("adults") - assert await db.list_materialized_views() == [] + assert [item async for item in db.list_materialized_views()] == [] def test_stored_queries_and_legacy_layouts_are_read(): @@ -509,3 +509,19 @@ def test_stored_queries_and_legacy_layouts_are_read(): ): with pytest.raises(NotImplementedError, match="cannot refresh"): read(newer) + + +def test_local_materialized_view_listing_pages_and_resume(tmp_path): + db = make_db(tmp_path) + db.create_materialized_view("a_view", "people", with_no_data=True) + db.create_materialized_view("z_view", "people", with_no_data=True) + db.create_table("ordinary", [{"id": 1}]) + names = db.list_materialized_views(page_limit=1) + assert names.num_page_results() == 0 + first = next(names) + assert names.num_page_results() == 0 + assert names.page_token() is not None + resumed = db.list_materialized_views(page_token=names.page_token(), page_limit=1) + assert sorted([first, *resumed]) == ["a_view", "z_view"] + assert resumed.page_token() is None + assert list(resumed) == [] diff --git a/python/python/tests/test_namespace.py b/python/python/tests/test_namespace.py index 7f095b71d..d9af24a90 100644 --- a/python/python/tests/test_namespace.py +++ b/python/python/tests/test_namespace.py @@ -10,6 +10,7 @@ import pytest import pyarrow as pa import lancedb from lance_namespace.errors import NamespaceNotEmptyError, TableNotFoundError +from lancedb.listing import AsyncListing, Listing from lancedb.namespace import _MAX_QUERY_K from lancedb.table import AsyncTable, LanceTable @@ -104,6 +105,18 @@ class TestNamespaceConnection: assert isinstance(db, lancedb.LanceNamespaceDBConnection) assert len(list(db.table_names())) == 0 + def test_list_materialized_views(self): + db = lancedb.connect_namespace("dir", {"root": self.temp_dir}) + db.create_table("source", [{"id": 1}]) + db.create_materialized_view("view", "source", with_no_data=True) + + names = db.list_materialized_views(page_limit=1) + assert isinstance(names, Listing) + assert names.num_page_results() == 0 + assert list(names) == ["view"] + assert names.num_page_results() == 0 + assert names.page_token() is None + def test_sync_builtin_namespace_uses_rust_without_python_client(self, monkeypatch): """Built-in sync namespace connections should not construct or call the Python namespace client for normal namespace/table management.""" @@ -605,6 +618,18 @@ class TestAsyncNamespaceConnection: table_names = await db.table_names() assert len(list(table_names)) == 0 + async def test_list_materialized_views(self): + db = lancedb.connect_namespace_async("dir", {"root": self.temp_dir}) + await db.create_table("source", [{"id": 1}]) + await db.create_materialized_view("view", "source", with_no_data=True) + + names = db.list_materialized_views(page_limit=1) + assert isinstance(names, AsyncListing) + assert names.num_page_results() == 0 + assert [name async for name in names] == ["view"] + assert names.num_page_results() == 0 + assert names.page_token() is None + async def test_async_builtin_namespace_uses_rust_without_python_client( self, monkeypatch ): diff --git a/python/python/tests/test_remote_db.py b/python/python/tests/test_remote_db.py index effe6f1a7..b72965faa 100644 --- a/python/python/tests/test_remote_db.py +++ b/python/python/tests/test_remote_db.py @@ -3125,7 +3125,7 @@ def test_remote_connection_jobs_surface(): request.end_headers() with mock_lancedb_connection(handler) as db: - jobs = db.list_jobs() + jobs = list(db.list_jobs()) assert [job.job_id for job in jobs] == ["job-1", "job-2"] assert jobs[0].state == "running" assert jobs[0].table == "t1" @@ -3285,7 +3285,7 @@ def test_view_crud_addresses_its_own_routes(): described = db.describe_view("adults", namespace_path=["analytics"]) assert described.schema == view.schema - assert db.list_views(namespace_path=["analytics"]) == ["adults"] + assert list(db.list_views(namespace_path=["analytics"])) == ["adults"] db.drop_view("adults", namespace_path=["analytics"]) assert paths == [ diff --git a/python/src/connection.rs b/python/src/connection.rs index 737ef05e2..71ba70e0e 100644 --- a/python/src/connection.rs +++ b/python/src/connection.rs @@ -464,12 +464,16 @@ impl Connection { }) } - pub fn list_materialized_views(self_: PyRef<'_, Self>) -> PyResult> { - let inner = self_.get_inner()?.clone(); - future_into_py(self_.py(), async move { - let views = inner.list_materialized_views().await.infer_error()?; - Ok(views) - }) + #[pyo3(signature = (*, page_token=None, page_limit=None))] + pub fn list_materialized_views( + &self, + page_token: Option, + page_limit: Option, + ) -> PyResult { + Ok(crate::listing::NameListing::new( + self.get_inner()? + .list_materialized_views(crate::listing::options(page_token, page_limit)), + )) } #[pyo3(signature = (name, namespace_path=None))] @@ -776,22 +780,19 @@ impl Connection { }) } - #[pyo3(signature = (namespace_path=None))] + #[pyo3(signature = (namespace_path=None, *, page_token=None, page_limit=None))] pub fn list_functions( - self_: PyRef<'_, Self>, + &self, namespace_path: Option>, - ) -> PyResult> { - let inner = self_.get_inner()?.clone(); - let namespace_path = namespace_path.unwrap_or_default(); - future_into_py(self_.py(), async move { - inner - .list_functions(&namespace_path) - .await - .infer_error()? - .into_iter() - .map(|function| function.to_canonical_json().infer_error()) - .collect::>>() - }) + page_token: Option, + page_limit: Option, + ) -> PyResult { + Ok(crate::listing::FunctionListing::new( + self.get_inner()?.list_functions( + &namespace_path.unwrap_or_default(), + crate::listing::options(page_token, page_limit), + ), + )) } #[pyo3(signature = (name, version, namespace_path=None))] @@ -863,16 +864,19 @@ impl Connection { }) } - #[pyo3(signature = (namespace_path=None))] + #[pyo3(signature = (namespace_path=None, *, page_token=None, page_limit=None))] pub fn list_secrets( - self_: PyRef<'_, Self>, + &self, namespace_path: Option>, - ) -> PyResult> { - let inner = self_.get_inner()?.clone(); - let namespace_path = namespace_path.unwrap_or_default(); - future_into_py(self_.py(), async move { - inner.list_secrets(&namespace_path).await.infer_error() - }) + page_token: Option, + page_limit: Option, + ) -> PyResult { + Ok(crate::listing::NameListing::new( + self.get_inner()?.list_secrets( + &namespace_path.unwrap_or_default(), + crate::listing::options(page_token, page_limit), + ), + )) } #[pyo3(signature = (name, namespace_path=None))] @@ -973,27 +977,31 @@ impl Connection { }) } - #[pyo3(signature = (namespace_path=None))] + #[pyo3(signature = (namespace_path=None, *, page_token=None, page_limit=None))] pub fn list_views( - self_: PyRef<'_, Self>, + &self, namespace_path: Option>, - ) -> PyResult> { - let inner = self_.get_inner()?.clone(); - let namespace_path = namespace_path.unwrap_or_default(); - future_into_py(self_.py(), async move { - inner.list_views(&namespace_path).await.infer_error() - }) + page_token: Option, + page_limit: Option, + ) -> PyResult { + Ok(crate::listing::NameListing::new( + self.get_inner()?.list_views( + &namespace_path.unwrap_or_default(), + crate::listing::options(page_token, page_limit), + ), + )) } - pub fn list_jobs(self_: PyRef<'_, Self>) -> PyResult> { - let inner = self_.get_inner()?.clone(); - future_into_py(self_.py(), async move { - let jobs = inner.list_jobs().await.infer_error()?; - Ok(jobs - .into_iter() - .map(crate::job::JobInfo::from) - .collect::>()) - }) + #[pyo3(signature = (*, page_token=None, page_limit=None))] + pub fn list_jobs( + &self, + page_token: Option, + page_limit: Option, + ) -> PyResult { + Ok(crate::listing::JobListing::new( + self.get_inner()? + .list_jobs(crate::listing::options(page_token, page_limit)), + )) } pub fn cancel_job(self_: PyRef<'_, Self>, job_id: String) -> PyResult> { diff --git a/python/src/lib.rs b/python/src/lib.rs index 57d30d183..9f88030ca 100644 --- a/python/src/lib.rs +++ b/python/src/lib.rs @@ -29,6 +29,7 @@ pub mod expr; pub mod header; pub mod index; pub mod job; +mod listing; pub mod namespace; pub mod oauth; pub mod otel; @@ -65,6 +66,10 @@ pub fn _lancedb(py: Python, m: &Bound<'_, PyModule>) -> PyResult<()> { m.add_class::()?; m.add_class::()?; m.add_class::()?; + 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/python/src/listing.rs b/python/src/listing.rs new file mode 100644 index 000000000..2d1a3916e --- /dev/null +++ b/python/src/listing.rs @@ -0,0 +1,131 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The LanceDB Authors + +use crate::error::PythonErrorExt; +use crate::runtime::future_into_py; +use futures::{StreamExt, future::poll_fn}; +use lancedb::listing::{Listing, ListingOptions}; +use pyo3::exceptions::PyStopAsyncIteration; +use pyo3::{Bound, PyAny, PyRef, PyResult, Python, pyclass, pymethods}; +use std::sync::{Arc, Mutex}; + +pub fn options(page_token: Option, page_limit: Option) -> ListingOptions { + let mut options = ListingOptions::default(); + options.page_token = page_token; + options.page_limit = page_limit; + options +} + +#[pyclass] +pub struct NameListing { + inner: Arc>>, + next_lock: Arc>, +} +impl NameListing { + pub fn new(inner: Listing) -> Self { + Self { + inner: Arc::new(Mutex::new(inner)), + next_lock: Arc::new(tokio::sync::Mutex::new(())), + } + } +} +#[pymethods] +impl NameListing { + 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 _guard = next_lock.lock().await; + let value = poll_fn(|cx| inner.lock().unwrap().poll_next_unpin(cx)) + .await + .ok_or_else(|| PyStopAsyncIteration::new_err(""))? + .infer_error()?; + Ok(value) + }) + } +} + +#[pyclass] +pub struct FunctionListing { + inner: Arc>>, + next_lock: Arc>, +} +impl FunctionListing { + pub fn new(inner: Listing) -> Self { + Self { + inner: Arc::new(Mutex::new(inner)), + next_lock: Arc::new(tokio::sync::Mutex::new(())), + } + } +} +#[pymethods] +impl FunctionListing { + 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 _guard = next_lock.lock().await; + let value = poll_fn(|cx| inner.lock().unwrap().poll_next_unpin(cx)) + .await + .ok_or_else(|| PyStopAsyncIteration::new_err(""))? + .infer_error()?; + value.to_canonical_json().infer_error() + }) + } +} + +#[pyclass] +pub struct JobListing { + inner: Arc>>, + next_lock: Arc>, +} +impl JobListing { + pub fn new(inner: Listing) -> Self { + Self { + inner: Arc::new(Mutex::new(inner)), + next_lock: Arc::new(tokio::sync::Mutex::new(())), + } + } +} +#[pymethods] +impl JobListing { + 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 _guard = next_lock.lock().await; + let value = poll_fn(|cx| inner.lock().unwrap().poll_next_unpin(cx)) + .await + .ok_or_else(|| PyStopAsyncIteration::new_err(""))? + .infer_error()?; + Ok(crate::job::JobInfo::from(value)) + }) + } +} diff --git a/rust/lancedb/src/connection.rs b/rust/lancedb/src/connection.rs index 1d0a4a0ab..b299cf4ff 100644 --- a/rust/lancedb/src/connection.rs +++ b/rust/lancedb/src/connection.rs @@ -28,6 +28,7 @@ use crate::database::{ }; use crate::embeddings::{EmbeddingRegistry, MemoryRegistry}; use crate::error::{Error, Result}; +use crate::listing::{Listing, ListingOptions}; #[cfg(feature = "remote")] use crate::remote::{ client::ClientConfig, @@ -625,8 +626,7 @@ impl Connection { /// List every published immutable Function version in `namespace_path` of /// the remote catalog. Functions in child namespaces are not included. /// - /// Results are ordered by Function name then version. The client walks all - /// server pages before returning. Local databases return + /// Results are ordered by Function name then version. Local databases return /// [`Error::NotSupported`]. /// /// # Example @@ -635,17 +635,30 @@ impl Connection { /// # async fn list_functions( /// # connection: &lancedb::Connection, /// # ) -> Result<(), Box> { - /// for function in connection.list_functions(&[]).await? { + /// use futures::TryStreamExt; + /// let mut functions = connection.list_functions(&[], Default::default()); + /// while let Some(function) = functions.try_next().await? { /// println!("{} {}", function.name(), function.version()); /// } /// # Ok(()) /// # } /// ``` - pub async fn list_functions( + /// + /// Results are fetched lazily, one page at a time. See [`Listing`] + /// for cached-result and continuation-token semantics. Errors terminate iteration. + /// + pub fn list_functions( &self, namespace_path: &[String], - ) -> Result> { - self.internal.list_functions(namespace_path).await + options: ListingOptions, + ) -> Listing { + let database = self.database().clone(); + let namespace_path = namespace_path.to_vec(); + Listing::new(options, move |options| { + let database = database.clone(); + let namespace_path = namespace_path.clone(); + async move { database.list_functions(&namespace_path, options).await } + }) } /// Remove the current Function name binding, retaining the object history. @@ -723,11 +736,37 @@ impl Connection { /// Names only. No path in this API returns a stored credential, by /// construction rather than by policy. Local databases return /// [`Error::NotSupported`]. - pub async fn list_secrets(&self, namespace_path: &[String]) -> Result> { - for segment in namespace_path { - validate_secret_component("Secret namespace path segment", segment)?; - } - self.internal.list_secrets(namespace_path).await + /// + /// Results are fetched lazily, one page at a time. See [`Listing`] + /// for cached-result and continuation-token semantics. Errors terminate iteration. + /// + /// ``` + /// # async fn example(connection: &lancedb::Connection) -> lancedb::Result<()> { + /// use futures::TryStreamExt; + /// let mut items = connection.list_secrets(&[], Default::default()); + /// while let Some(item) = items.try_next().await? { + /// println!("{item:?}"); + /// } + /// # Ok(()) + /// # } + /// ``` + pub fn list_secrets( + &self, + namespace_path: &[String], + options: ListingOptions, + ) -> Listing { + let database = self.database().clone(); + let namespace_path = namespace_path.to_vec(); + Listing::new(options, move |options| { + let database = database.clone(); + let namespace_path = namespace_path.clone(); + async move { + for segment in &namespace_path { + validate_secret_component("Secret namespace path segment", segment)?; + } + database.list_secrets(&namespace_path, options).await + } + }) } /// Drop a Secret. @@ -795,7 +834,8 @@ impl Connection { /// let described = connection.describe_view("recent_orders", &namespace).await?; /// println!("{} in {:?}", described.query, described.default_namespace_path); /// - /// let names = connection.list_views(&namespace).await?; + /// use futures::TryStreamExt; + /// let names: Vec<_> = connection.list_views(&namespace, Default::default()).try_collect().await?; /// assert!(names.iter().any(|name| name == "recent_orders")); /// /// // Dropping the view leaves `orders` untouched. @@ -871,11 +911,36 @@ impl Connection { /// The names of the views in one namespace. /// /// Names only; a definition is query metadata and comes from - /// [`Self::describe_view`]. The client walks all server pages before - /// returning. Local databases return [`Error::NotSupported`]. - pub async fn list_views(&self, namespace_path: &[String]) -> Result> { - validate_namespace(namespace_path)?; - self.internal.list_views(namespace_path).await + /// [`Self::describe_view`]. Local databases return [`Error::NotSupported`]. + /// + /// Results are fetched lazily, one page at a time. See [`Listing`] + /// for cached-result and continuation-token semantics. Errors terminate iteration. + /// + /// ``` + /// # async fn example(connection: &lancedb::Connection) -> lancedb::Result<()> { + /// use futures::TryStreamExt; + /// let mut items = connection.list_views(&[], Default::default()); + /// while let Some(item) = items.try_next().await? { + /// println!("{item:?}"); + /// } + /// # Ok(()) + /// # } + /// ``` + pub fn list_views( + &self, + namespace_path: &[String], + options: ListingOptions, + ) -> Listing { + let database = self.database().clone(); + let namespace_path = namespace_path.to_vec(); + Listing::new(options, move |options| { + let database = database.clone(); + let namespace_path = namespace_path.clone(); + async move { + validate_namespace(&namespace_path)?; + database.list_views(&namespace_path, options).await + } + }) } /// Rename a table in the database. @@ -934,8 +999,26 @@ impl Connection { } /// List server-side jobs across the database's tables. - pub async fn list_jobs(&self) -> Result> { - self.internal.list_jobs().await + /// + /// Results are fetched lazily, one page at a time. See [`Listing`] + /// for cached-result and continuation-token semantics. Errors terminate iteration. + /// + /// ``` + /// # async fn example(connection: &lancedb::Connection) -> lancedb::Result<()> { + /// use futures::TryStreamExt; + /// let mut items = connection.list_jobs(Default::default()); + /// while let Some(item) = items.try_next().await? { + /// println!("{item:?}"); + /// } + /// # Ok(()) + /// # } + /// ``` + pub fn list_jobs(&self, options: ListingOptions) -> Listing { + let database = self.database().clone(); + Listing::new(options, move |options| { + let database = database.clone(); + async move { database.list_jobs(options).await } + }) } /// Request cancellation of a server-side job by id. Returns true if the diff --git a/rust/lancedb/src/database.rs b/rust/lancedb/src/database.rs index 2518edf13..ad4e8118d 100644 --- a/rust/lancedb/src/database.rs +++ b/rust/lancedb/src/database.rs @@ -14,7 +14,7 @@ //! * Tables may be managed by a database system (e.g. Postgres) //! * A custom table implementation (e.g. remote table, etc.) may be used -use std::collections::{HashMap, HashSet}; +use std::collections::HashMap; use std::sync::Arc; use std::time::Duration; @@ -29,6 +29,7 @@ use lance_namespace::models::{ use crate::data::scannable::Scannable; use crate::error::Result; use crate::job::Job; +use crate::listing::{ListingOptions, ListingPage}; use crate::materialized_view::CreateMaterializedViewRequest; use crate::secrets::SecretInfo; use crate::table::{BaseTable, WriteOptions}; @@ -362,52 +363,49 @@ pub trait Database: ) -> Result { job_op_not_supported("remote materialized-view drop") } - /// List materialized-view names in a namespace. + /// Fetch one page of materialized-view names in a namespace. + /// Local backends filter a page of table names, so an empty result may + /// still have a continuation token. #[doc(hidden)] - async fn list_materialized_views(&self, namespace_path: &[String]) -> Result> { + async fn list_materialized_views( + &self, + namespace_path: &[String], + options: ListingOptions, + ) -> Result> { + let response = self + .list_tables(ListTablesRequest { + id: Some(namespace_path.to_vec()), + page_token: options.page_token, + limit: options.page_limit.map(|limit| limit as i32), + ..Default::default() + }) + .await?; let mut names = Vec::new(); - let mut page_token = None; - let mut seen_page_tokens = HashSet::new(); - loop { - let response = self - .list_tables(ListTablesRequest { - id: Some(namespace_path.to_vec()), - page_token: page_token.clone(), - ..Default::default() + for name in response.tables { + let Ok(table) = self + .open_table(OpenTableRequest { + name: name.clone(), + namespace_path: namespace_path.to_vec(), + index_cache_size: None, + lance_read_params: None, + location: None, + namespace_client: None, + managed_versioning: None, }) - .await?; - for name in response.tables { - let Ok(table) = self - .open_table(OpenTableRequest { - name: name.clone(), - namespace_path: namespace_path.to_vec(), - index_cache_size: None, - lance_read_params: None, - location: None, - namespace_client: None, - managed_versioning: None, - }) - .await - else { - continue; - }; - let schema = table.schema().await?; - if crate::materialized_view::read_definition(schema.metadata())?.is_some() { - names.push(name); - } - } - let Some(next_page_token) = response.page_token.filter(|token| !token.is_empty()) + .await else { - break; + continue; }; - if !seen_page_tokens.insert(next_page_token.clone()) { - return Err(crate::Error::Runtime { - message: "materialized-view listing repeated a page token".into(), - }); + let schema = table.schema().await?; + if crate::materialized_view::read_definition(schema.metadata())?.is_some() { + names.push(name); } - page_token = Some(next_page_token); } - Ok(names) + + Ok(ListingPage { + items: names, + page_token: response.page_token, + }) } /// Look up one exact immutable Function version. async fn get_function( @@ -418,12 +416,12 @@ pub trait Database: ) -> Result { function_catalog_not_supported() } - /// List every published immutable Function version in one namespace of the - /// remote catalog. + /// Fetch one page of published immutable Function versions in one namespace. async fn list_functions( &self, _namespace_path: &[String], - ) -> Result> { + _options: ListingOptions, + ) -> Result> { function_catalog_not_supported() } /// Remove the current Function name binding, retaining the object history. @@ -469,11 +467,15 @@ pub trait Database: ) -> Result<()> { secret_catalog_not_supported() } - /// The names of every Secret in this database. + /// Fetch one page of Secret names in this database. /// /// Names only. No API path returns a stored credential, by construction /// rather than by policy. - async fn list_secrets(&self, _namespace_path: &[String]) -> Result> { + async fn list_secrets( + &self, + _namespace_path: &[String], + _options: ListingOptions, + ) -> Result> { secret_catalog_not_supported() } /// Drop a Secret. Functions bound to it fail at their next job, which is @@ -516,8 +518,12 @@ pub trait Database: async fn drop_view_async(&self, _name: &str, _namespace_path: &[String]) -> Result { view_ops_not_supported() } - /// The names of the views in one namespace. - async fn list_views(&self, _namespace_path: &[String]) -> Result> { + /// Fetch one page of view names in one namespace. + async fn list_views( + &self, + _namespace_path: &[String], + _options: ListingOptions, + ) -> Result> { view_ops_not_supported() } /// Open a job by id, returning a handle with its record already @@ -526,8 +532,8 @@ pub trait Database: async fn open_job(&self, _job_id: &str) -> Result { job_op_not_supported("open_job") } - /// List server-side jobs across the database's tables. - async fn list_jobs(&self) -> Result> { + /// Fetch one page of server-side jobs across the database's tables. + async fn list_jobs(&self, _options: ListingOptions) -> Result> { job_op_not_supported("list_jobs") } /// Request cancellation of a job by id. Returns true if the server diff --git a/rust/lancedb/src/lib.rs b/rust/lancedb/src/lib.rs index 530328cbb..e215f3c6b 100644 --- a/rust/lancedb/src/lib.rs +++ b/rust/lancedb/src/lib.rs @@ -187,6 +187,7 @@ pub mod index; pub mod io; pub mod ipc; pub mod job; +pub mod listing; pub mod materialized_view; #[cfg(feature = "metrics-otel")] pub mod metrics_otel; diff --git a/rust/lancedb/src/listing.rs b/rust/lancedb/src/listing.rs new file mode 100644 index 000000000..710202d9d --- /dev/null +++ b/rust/lancedb/src/listing.rs @@ -0,0 +1,261 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The LanceDB Authors + +//! Lazy, resumable listings of database resources. + +use std::collections::HashSet; +use std::pin::Pin; +use std::task::{Context, Poll}; + +use futures::future::BoxFuture; +use futures::{Future, FutureExt, Stream}; + +use crate::{Error, Result}; + +/// Pagination controls for a resource listing. +#[derive(Clone, Debug, Default)] +#[non_exhaustive] +pub struct ListingOptions { + /// Starting continuation token; `None` starts at the beginning. + pub page_token: Option, + /// Maximum results per request, not a limit on the entire listing. + /// `None` uses the backend default. Must be between 1 and `i32::MAX`. + pub page_limit: Option, +} + +impl ListingOptions { + /// Resume at a saved continuation token. + pub fn page_token(mut self, token: impl Into) -> Self { + self.page_token = Some(token.into()); + self + } + + /// Set the maximum number of results requested per page. + pub fn page_limit(mut self, limit: u32) -> Self { + self.page_limit = Some(limit); + self + } +} + +/// One backend page of a resource listing. +#[derive(Debug)] +pub struct ListingPage { + /// Items returned on this page, possibly empty even when more pages exist. + pub items: Vec, + /// Token for the next page; `None` or an empty string ends the listing. + pub page_token: Option, +} + +type FetchPage = + Box BoxFuture<'static, Result>> + Send + Sync>; + +/// A lazy stream of resources with inspectable pagination state. +/// +/// No requests are made until the stream is polled. Empty pages with continuation +/// tokens are skipped. Errors terminate the stream and retain the failed request's +/// token, allowing a new stream to retry it. Repeated tokens are rejected. +/// +/// A token advances when its page is fetched. Drain [`Self::num_page_results`] +/// before saving [`Self::page_token`] to avoid skipping cached items on resumption. +/// After the final page is fetched the token is `None`, even with cached items. +/// Use [`futures::TryStreamExt::try_collect`] to collect a listing into a vector. +pub struct Listing { + fetch: FetchPage, + items: std::vec::IntoIter, + options: ListingOptions, + done: bool, + seen_tokens: HashSet, + pending: Option>>>, +} + +impl Listing { + pub(crate) fn new(options: ListingOptions, fetch: F) -> Self + where + F: Fn(ListingOptions) -> Fut + Send + Sync + 'static, + Fut: Future>> + Send + 'static, + { + let seen_tokens = options + .page_token + .iter() + .filter(|t| !t.is_empty()) + .cloned() + .collect(); + Self { + fetch: Box::new(move |options| fetch(options).boxed()), + items: Vec::new().into_iter(), + options, + done: false, + seen_tokens, + pending: None, + } + } + + /// Number of items available without another backend request. + pub fn num_page_results(&self) -> usize { + self.items.len() + } + + /// Token for the next request. Initially the supplied token; `None` after + /// fetching the final page. Drain cached items before saving this token. + pub fn page_token(&self) -> Option<&str> { + self.options.page_token.as_deref() + } +} + +impl Stream for Listing { + type Item = Result; + + fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { + let this = self.get_mut(); + loop { + if let Some(item) = this.items.next() { + return Poll::Ready(Some(Ok(item))); + } + if this.done { + return Poll::Ready(None); + } + if this + .options + .page_limit + .is_some_and(|limit| limit == 0 || limit > i32::MAX as u32) + { + this.done = true; + return Poll::Ready(Some(Err(Error::InvalidInput { + message: "Listing page limit must be between 1 and 2147483647".into(), + }))); + } + let pending = this + .pending + .get_or_insert_with(|| (this.fetch)(this.options.clone())); + let response = futures::ready!(pending.as_mut().poll(cx)); + this.pending = None; + match response { + Ok(response) => { + let token = response.page_token.filter(|token| !token.is_empty()); + if token + .as_ref() + .is_some_and(|token| !this.seen_tokens.insert(token.clone())) + { + this.done = true; + return Poll::Ready(Some(Err(Error::Runtime { + message: "Listing response repeated a page_token".into(), + }))); + } + this.done = token.is_none(); + this.options.page_token = token; + this.items = response.items.into_iter(); + } + Err(error) => { + this.done = true; + return Poll::Ready(Some(Err(error))); + } + } + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + use futures::{StreamExt, TryStreamExt}; + use std::sync::{Arc, Mutex}; + + #[tokio::test] + async fn lazy_pages_empty_pages_and_resume() { + let requests = Arc::new(Mutex::new(Vec::new())); + let fetch = { + let requests = requests.clone(); + move |options: ListingOptions| { + requests.lock().unwrap().push(options.clone()); + async move { + let (items, token) = match options.page_token.as_deref() { + Some("start") => (vec![1, 2], Some("empty")), + Some("empty") => (vec![], Some("last")), + Some("last") => (vec![3, 4], Some("")), + _ => panic!("unexpected request"), + }; + Ok(ListingPage { + items, + page_token: token.map(str::to_owned), + }) + } + } + }; + let mut listing = Listing::new( + ListingOptions::default().page_token("start").page_limit(2), + fetch.clone(), + ); + assert!(requests.lock().unwrap().is_empty()); + assert_eq!(listing.page_token(), Some("start")); + assert_eq!(listing.num_page_results(), 0); + assert_eq!(listing.try_next().await.unwrap(), Some(1)); + assert_eq!(listing.num_page_results(), 1); + assert_eq!(listing.page_token(), Some("empty")); + assert_eq!(listing.try_next().await.unwrap(), Some(2)); + assert_eq!(requests.lock().unwrap().len(), 1); + let mut resumed = Listing::new( + ListingOptions::default() + .page_token(listing.page_token().unwrap()) + .page_limit(2), + fetch, + ); + assert_eq!(resumed.try_next().await.unwrap(), Some(3)); + assert_eq!(resumed.page_token(), None); + assert_eq!(resumed.num_page_results(), 1); + assert_eq!(resumed.try_next().await.unwrap(), Some(4)); + assert_eq!(resumed.try_next().await.unwrap(), None); + assert_eq!(resumed.try_next().await.unwrap(), None); + let requests = requests.lock().unwrap(); + assert_eq!(requests.len(), 3); + assert!(requests.iter().all(|request| request.page_limit == Some(2))); + } + + #[tokio::test] + async fn failure_retains_token_and_terminates() { + let mut listing = + Listing::::new(ListingOptions::default().page_token("retry"), |_| async { + Err(Error::Runtime { + message: "failed".into(), + }) + }); + assert!(listing.next().await.unwrap().is_err()); + assert_eq!(listing.page_token(), Some("retry")); + assert_eq!(listing.num_page_results(), 0); + assert!(listing.next().await.is_none()); + } + + #[tokio::test] + async fn invalid_limits_do_not_fetch() { + for limit in [0, u32::MAX] { + let mut listing = + Listing::::new(ListingOptions::default().page_limit(limit), |_| async { + panic!("invalid limits must not reach the backend") + }); + assert!(matches!( + listing.next().await, + Some(Err(Error::InvalidInput { .. })) + )); + assert!(listing.next().await.is_none()); + } + } + + #[tokio::test] + async fn jobs_can_exceed_the_old_page_cap() { + let listing = Listing::new(ListingOptions::default(), |options| async move { + let n: usize = options + .page_token + .as_deref() + .unwrap_or("0") + .parse() + .unwrap(); + Ok(ListingPage { + items: vec![n], + page_token: (n < 110).then(|| (n + 1).to_string()), + }) + }); + assert_eq!( + listing.try_collect::>().await.unwrap(), + (0..=110).collect::>() + ); + } +} diff --git a/rust/lancedb/src/materialized_view.rs b/rust/lancedb/src/materialized_view.rs index 70b39506b..82e6aa802 100644 --- a/rust/lancedb/src/materialized_view.rs +++ b/rust/lancedb/src/materialized_view.rs @@ -27,6 +27,7 @@ use crate::database::{CreateTableRequest, Database, OpenTableRequest}; use crate::embeddings::EmbeddingDefinition; use crate::function::FunctionBinding; use crate::job::Job; +use crate::listing::{Listing, ListingOptions}; use crate::table::Table; use crate::table::computed_columns::{ FUNCTION_BINDINGS_META_KEY, computed_column_from_field, computed_columns, @@ -1805,8 +1806,26 @@ impl Connection { } /// The names of materialized views in the root namespace. - pub async fn list_materialized_views(&self) -> Result> { - self.database().list_materialized_views(&[]).await + /// + /// Results are fetched lazily, one page at a time. See [`Listing`] + /// for cached-result and continuation-token semantics. Errors terminate iteration. + /// + /// ``` + /// # async fn example(connection: &lancedb::Connection) -> lancedb::Result<()> { + /// use futures::TryStreamExt; + /// let mut items = connection.list_materialized_views(Default::default()); + /// while let Some(item) = items.try_next().await? { + /// println!("{item:?}"); + /// } + /// # Ok(()) + /// # } + /// ``` + pub fn list_materialized_views(&self, options: ListingOptions) -> Listing { + let database = self.database().clone(); + Listing::new(options, move |options| { + let database = database.clone(); + async move { database.list_materialized_views(&[], options).await } + }) } /// Drop a materialized view. diff --git a/rust/lancedb/src/remote/db.rs b/rust/lancedb/src/remote/db.rs index 1453d8d73..c44f7558e 100644 --- a/rust/lancedb/src/remote/db.rs +++ b/rust/lancedb/src/remote/db.rs @@ -1,7 +1,7 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright The LanceDB Authors -use std::collections::{HashMap, HashSet}; +use std::collections::HashMap; use std::sync::Arc; use async_trait::async_trait; @@ -29,6 +29,7 @@ use crate::function::{ PythonRuntimeSpec, }; use crate::job::Job; +use crate::listing::{ListingOptions, ListingPage}; use crate::materialized_view::CreateMaterializedViewRequest; use crate::remote::job::{PauseJobResponse, RemoteJob, ResumeJobResponse, job_state_to_client}; use crate::remote::util::stream_as_body; @@ -843,10 +844,6 @@ struct RemoteListViewsResponse { page_token: Option, } -/// Bound on `list_jobs` page walking; a warning is logged when the listing -/// is truncated at this many pages. -const MAX_LIST_JOBS_PAGES: usize = 100; - #[async_trait] impl Database for RemoteDatabase { fn uri(&self) -> &str { @@ -919,45 +916,28 @@ impl Database for RemoteDatabase { Ok(job) } - async fn list_materialized_views(&self, namespace_path: &[String]) -> Result> { - #[derive(serde::Deserialize)] - struct ListMaterializedViewsResponse { - #[serde(default)] - views: Vec, - #[serde(default)] - page_token: Option, - } - + async fn list_materialized_views( + &self, + namespace_path: &[String], + options: ListingOptions, + ) -> Result> { let namespace_id = build_namespace_identifier(namespace_path)?; let path = format!("/v1/namespace/{namespace_id}/materialized_view/list"); - let mut views = Vec::new(); - let mut page_token: Option = None; - let mut seen_page_tokens = HashSet::new(); - loop { - let mut req = self.client.get(&path); - if let Some(token) = &page_token { - req = req.query(&[("page_token", token)]); - } - let (request_id, response) = self.client.send(req).await?; - let response = self.client.check_response(&request_id, response).await?; - let status = response.status(); - let response: ListMaterializedViewsResponse = - response.json().await.err_to_http(request_id.clone())?; - views.extend(response.views); - let Some(next_page_token) = response.page_token.filter(|token| !token.is_empty()) - else { - break; - }; - if !seen_page_tokens.insert(next_page_token.clone()) { - return Err(Error::Http { - source: "Materialized-view listing response repeated a page_token".into(), - request_id, - status_code: Some(status), - }); - } - page_token = Some(next_page_token); + let mut req = self.client.get(&path); + if let Some(token) = options.page_token { + req = req.query(&[("page_token", token)]); } - Ok(views) + if let Some(limit) = options.page_limit { + req = req.query(&[("limit", limit)]); + } + let (request_id, response) = self.client.send(req).await?; + let response = self.client.check_response(&request_id, response).await?; + let response: RemoteListViewsResponse = response.json().await.err_to_http(request_id)?; + let items = response.views; + Ok(ListingPage { + items, + page_token: response.page_token, + }) } async fn create_function_async( @@ -1008,45 +988,34 @@ impl Database for RemoteDatabase { response.json().await.err_to_http(request_id) } - async fn list_functions(&self, namespace_path: &[String]) -> Result> { + async fn list_functions( + &self, + namespace_path: &[String], + options: ListingOptions, + ) -> Result> { let namespace_id = build_namespace_identifier(namespace_path)?; let path = format!("/v1/namespace/{namespace_id}/function/list"); - let mut functions = Vec::new(); - let mut page_token: Option = None; - let mut seen_page_tokens = HashSet::new(); - loop { - let mut req = self - .client - .get(&path) - .query(&[("include_definition", true)]); - if let Some(token) = &page_token { - req = req.query(&[("page_token", token)]); - } - let (request_id, response) = self.client.send(req).await?; - let response = self.client.check_response(&request_id, response).await?; - let status = response.status(); - let response: RemoteListFunctionsResponse = - response.json().await.err_to_http(request_id.clone())?; - functions.extend( - response - .functions - .into_iter() - .map(|listed| listed.definition), - ); - let Some(next_page_token) = response.page_token.filter(|token| !token.is_empty()) - else { - break; - }; - if !seen_page_tokens.insert(next_page_token.clone()) { - return Err(Error::Http { - source: "Function listing response repeated a page_token".into(), - request_id, - status_code: Some(status), - }); - } - page_token = Some(next_page_token); + let mut req = self.client.get(&path); + if let Some(token) = options.page_token { + req = req.query(&[("page_token", token)]); } - Ok(functions) + if let Some(limit) = options.page_limit { + req = req.query(&[("limit", limit)]); + } + req = req.query(&[("include_definition", true)]); + let (request_id, response) = self.client.send(req).await?; + let response = self.client.check_response(&request_id, response).await?; + let response: RemoteListFunctionsResponse = + response.json().await.err_to_http(request_id)?; + let items = response + .functions + .into_iter() + .map(|listed| listed.definition) + .collect(); + Ok(ListingPage { + items, + page_token: response.page_token, + }) } async fn drop_function( @@ -1129,37 +1098,32 @@ impl Database for RemoteDatabase { .await } - async fn list_secrets(&self, namespace_path: &[String]) -> Result> { + async fn list_secrets( + &self, + namespace_path: &[String], + options: ListingOptions, + ) -> Result> { let namespace_id = build_namespace_identifier(namespace_path)?; let path = format!("/v1/namespace/{namespace_id}/secret/list"); - let mut names = Vec::new(); - let mut page_token: Option = None; - let mut seen_page_tokens = HashSet::new(); - loop { - let mut req = self.client.get(&path); - if let Some(token) = &page_token { - req = req.query(&[("page_token", token)]); - } - let (request_id, response) = self.client.send(req).await?; - let response = self.client.check_response(&request_id, response).await?; - let status = response.status(); - let response: RemoteListSecretsResponse = - response.json().await.err_to_http(request_id.clone())?; - names.extend(response.secrets.into_iter().map(|secret| secret.name)); - let Some(next_page_token) = response.page_token.filter(|token| !token.is_empty()) - else { - break; - }; - if !seen_page_tokens.insert(next_page_token.clone()) { - return Err(Error::Http { - source: "Secret listing response repeated a page_token".into(), - request_id, - status_code: Some(status), - }); - } - page_token = Some(next_page_token); + let mut req = self.client.get(&path); + if let Some(token) = options.page_token { + req = req.query(&[("page_token", token)]); } - Ok(names) + if let Some(limit) = options.page_limit { + req = req.query(&[("limit", limit)]); + } + let (request_id, response) = self.client.send(req).await?; + let response = self.client.check_response(&request_id, response).await?; + let response: RemoteListSecretsResponse = response.json().await.err_to_http(request_id)?; + let items = response + .secrets + .into_iter() + .map(|secret| secret.name) + .collect(); + Ok(ListingPage { + items, + page_token: response.page_token, + }) } async fn drop_secret(&self, name: &str, namespace_path: &[String]) -> Result<()> { @@ -1248,37 +1212,28 @@ impl Database for RemoteDatabase { } } - async fn list_views(&self, namespace_path: &[String]) -> Result> { + async fn list_views( + &self, + namespace_path: &[String], + options: ListingOptions, + ) -> Result> { let namespace_id = build_namespace_identifier(namespace_path)?; let path = format!("/v1/namespace/{namespace_id}/view/list"); - let mut views = Vec::new(); - let mut page_token: Option = None; - let mut seen_page_tokens = HashSet::new(); - loop { - let mut req = self.client.get(&path); - if let Some(token) = &page_token { - req = req.query(&[("page_token", token)]); - } - let (request_id, response) = self.client.send(req).await?; - let response = self.client.check_response(&request_id, response).await?; - let status = response.status(); - let response: RemoteListViewsResponse = - response.json().await.err_to_http(request_id.clone())?; - views.extend(response.views); - let Some(next_page_token) = response.page_token.filter(|token| !token.is_empty()) - else { - break; - }; - if !seen_page_tokens.insert(next_page_token.clone()) { - return Err(Error::Http { - source: "View listing response repeated a page_token".into(), - request_id, - status_code: Some(status), - }); - } - page_token = Some(next_page_token); + let mut req = self.client.get(&path); + if let Some(token) = options.page_token { + req = req.query(&[("page_token", token)]); } - Ok(views) + if let Some(limit) = options.page_limit { + req = req.query(&[("limit", limit)]); + } + let (request_id, response) = self.client.send(req).await?; + let response = self.client.check_response(&request_id, response).await?; + let response: RemoteListViewsResponse = response.json().await.err_to_http(request_id)?; + let items = response.views; + Ok(ListingPage { + items, + page_token: response.page_token, + }) } async fn open_job(&self, job_id: &str) -> Result { @@ -1295,47 +1250,33 @@ impl Database for RemoteDatabase { } } - async fn list_jobs(&self) -> Result> { - let mut out = Vec::new(); - let mut page_token: Option = None; - let mut seen_page_tokens = HashSet::new(); - for page in 0..MAX_LIST_JOBS_PAGES { - let mut body = serde_json::json!({}); - if let Some(token) = &page_token { - body["page_token"] = serde_json::Value::String(token.clone()); - } - let req = self.client.post("/v1/jobs/list").json(&body); - let (request_id, rsp) = self.client.send(req).await?; - let rsp = self.client.check_response(&request_id, rsp).await?; - let status = rsp.status(); - let body: RemoteListJobsResponse = rsp.json().await.err_to_http(request_id.clone())?; - out.extend(body.jobs.into_iter().map(|row| JobInfo { + async fn list_jobs(&self, options: ListingOptions) -> Result> { + let mut body = serde_json::json!({}); + if let Some(token) = options.page_token { + body["page_token"] = token.into(); + } + if let Some(limit) = options.page_limit { + body["limit"] = limit.into(); + } + let req = self.client.post("/v1/jobs/list").json(&body); + let (request_id, rsp) = self.client.send(req).await?; + let rsp = self.client.check_response(&request_id, rsp).await?; + let response: RemoteListJobsResponse = rsp.json().await.err_to_http(request_id)?; + let items = response + .jobs + .into_iter() + .map(|row| JobInfo { job_id: row.job_id, table: row.table, job_type: row.job_type, state: job_state_to_client(&row.state), created_at_millis: row.created_at_millis, - })); - let Some(next_page_token) = body.page_token.filter(|token| !token.is_empty()) else { - break; - }; - if !seen_page_tokens.insert(next_page_token.clone()) { - return Err(Error::Http { - source: "Job listing response repeated a page_token".into(), - request_id, - status_code: Some(status), - }); - } - page_token = Some(next_page_token); - if page + 1 == MAX_LIST_JOBS_PAGES { - log::warn!( - "list_jobs truncated after {} pages ({} jobs)", - MAX_LIST_JOBS_PAGES, - out.len() - ); - } - } - Ok(out) + }) + .collect(); + Ok(ListingPage { + items, + page_token: response.page_token, + }) } async fn cancel_job(&self, job_id: &str) -> Result { @@ -1905,6 +1846,8 @@ impl From for RemoteOptions { #[cfg(test)] mod tests { use super::{NamespaceHeaderProviderContext, build_cache_key}; + use crate::listing::{Listing, ListingOptions}; + use futures::TryStreamExt; use std::collections::HashMap; use std::sync::atomic::{AtomicUsize, Ordering}; use std::sync::{Arc, OnceLock}; @@ -2123,9 +2066,19 @@ mod tests { } }); assert_eq!( - db.list_materialized_views(&["analytics".into()]) - .await - .unwrap(), + Listing::new(Default::default(), { + let db = Arc::new(db); + move |options| { + let db = db.clone(); + async move { + db.list_materialized_views(&["analytics".into()], options) + .await + } + } + }) + .try_collect::>() + .await + .unwrap(), ["adults"] ); } @@ -3534,6 +3487,110 @@ mod tests { } } + #[tokio::test] + async fn test_resource_listings_are_lazy_and_forward_pagination() { + async fn check(mut items: Listing, count: Arc) { + assert_eq!(count.load(Ordering::SeqCst), 0); + assert_eq!(items.num_page_results(), 0); + assert_eq!(items.page_token(), Some("start/token")); + assert!(items.try_next().await.unwrap().is_some()); + assert_eq!(items.num_page_results(), 1); + assert_eq!(items.page_token(), Some("empty")); + assert!(items.try_next().await.unwrap().is_some()); + assert_eq!(count.load(Ordering::SeqCst), 1); + assert!(items.try_next().await.unwrap().is_some()); + assert_eq!(count.load(Ordering::SeqCst), 3); + assert_eq!(items.num_page_results(), 1); + assert_eq!(items.page_token(), None); + assert!(items.try_next().await.unwrap().is_some()); + assert!(items.try_next().await.unwrap().is_none()); + assert!(items.try_next().await.unwrap().is_none()); + } + for resource in ["secret", "view", "function", "materialized_view", "jobs"] { + let count = Arc::new(AtomicUsize::new(0)); + let conn = Connection::new_with_handler({ + let count = count.clone(); + move |request| { + let page = count.fetch_add(1, Ordering::SeqCst); + assert!(page < 3); + let token = ["start/token", "empty", "last"][page]; + if resource == "jobs" { + assert_eq!(request.url().path(), "/v1/jobs/list"); + let body: serde_json::Value = + serde_json::from_slice(request.body().unwrap().as_bytes().unwrap()) + .unwrap(); + assert_eq!(body, serde_json::json!({"page_token": token, "limit": 2})); + } else { + let ns = if resource == "materialized_view" { + "$" + } else { + "team" + }; + assert_eq!( + request.url().path(), + format!("/v1/namespace/{ns}/{resource}/list") + ); + let query = request.url().query_pairs().collect::>(); + assert_eq!(query.get("page_token").unwrap(), token); + assert_eq!(query.get("limit").unwrap(), "2"); + if resource == "function" { + assert_eq!(query.get("include_definition").unwrap(), "true"); + } + } + let (key, item) = match resource { + "secret" => ("secrets", serde_json::json!({"name": "a"})), + "jobs" => ( + "jobs", + serde_json::json!({"job_id": "a", "table": "t", "job_type": "create_index", "state": "in_progress", "created_at_millis": 1}), + ), + "function" => ( + "functions", + serde_json::json!({"name": "embed", "version": "1", "definition": serde_json::from_str::(include_str!("../../tests/fixtures/first_class_functions/v1/remote_function_version.canonical.json")).unwrap()}), + ), + _ => ("views", serde_json::json!("a")), + }; + let items = if page == 1 { + vec![] + } else { + vec![item.clone(), item] + }; + http::Response::builder().status(200).body(serde_json::json!({key: items, "page_token": (["empty", "last", ""][page])}).to_string()).unwrap() + } + }); + let options = ListingOptions::default() + .page_token("start/token") + .page_limit(2); + let namespace = vec!["team".to_owned()]; + match resource { + "secret" => check(conn.list_secrets(&namespace, options), count).await, + "view" => check(conn.list_views(&namespace, options), count).await, + "function" => check(conn.list_functions(&namespace, options), count).await, + "materialized_view" => check(conn.list_materialized_views(options), count).await, + _ => check(conn.list_jobs(options), count).await, + } + } + } + + #[tokio::test] + async fn test_list_jobs_does_not_truncate_after_100_pages() { + let conn = Connection::new_with_handler(|request| { + let body: serde_json::Value = + serde_json::from_slice(request.body().unwrap().as_bytes().unwrap()).unwrap(); + let page: u32 = body["page_token"].as_str().unwrap_or("0").parse().unwrap(); + http::Response::builder().status(200).body(serde_json::json!({ + "jobs": [{"job_id": page.to_string(), "table": "t", "job_type": "create_index", "state": "in_progress", "created_at_millis": 1}], + "page_token": (page < 110).then(|| (page + 1).to_string()), + }).to_string()).unwrap() + }); + let jobs = conn + .list_jobs(Default::default()) + .try_collect::>() + .await + .unwrap(); + assert_eq!(jobs.len(), 111); + assert_eq!(jobs.last().unwrap().job_id, "110"); + } + #[tokio::test] async fn test_list_jobs_paginates() { let page = Arc::new(AtomicUsize::new(0)); @@ -3563,7 +3620,11 @@ mod tests { } } }); - let jobs = conn.list_jobs().await.unwrap(); + let jobs = conn + .list_jobs(Default::default()) + .try_collect::>() + .await + .unwrap(); assert_eq!(jobs.len(), 3); assert_eq!(jobs[0].job_id, "job-1"); assert_eq!(jobs[0].table, "t1"); @@ -3592,17 +3653,14 @@ mod tests { .unwrap() }); - let error = conn.list_jobs().await.unwrap_err(); + let mut items = conn.list_jobs(Default::default()); + let error = (&mut items).try_collect::>().await.unwrap_err(); assert!( - matches!( - &error, - Error::Http { - status_code: Some(http::StatusCode::OK), - .. - } - ), + matches!(&error, Error::Runtime { message } if message.contains("repeated a page_token")), "got {error:?}" ); + assert!(items.page_token().is_some()); + assert!(items.try_next().await.unwrap().is_none()); assert_eq!(requests.load(Ordering::SeqCst), 2); } @@ -3870,7 +3928,8 @@ mod tests { ); let by_segment = conn - .list_secrets(&[component.to_string()]) + .list_secrets(&[component.to_string()], Default::default()) + .try_collect::>() .await .expect_err("a dot-only namespace segment must be refused"); assert!( @@ -4069,7 +4128,10 @@ mod tests { http::Response::builder().status(200).body(body).unwrap() }); assert_eq!( - conn.list_secrets(&[]).await.unwrap(), + conn.list_secrets(&[], Default::default()) + .try_collect::>() + .await + .unwrap(), vec!["openai-prod".to_string(), "hf-prod".to_string()] ); } @@ -4084,7 +4146,11 @@ mod tests { .body(r#"{"secrets":[{"name":"openai-prod"}],"page_token":"same"}"#) .unwrap() }); - let error = conn.list_secrets(&[]).await.unwrap_err(); + let error = conn + .list_secrets(&[], Default::default()) + .try_collect::>() + .await + .unwrap_err(); assert!( error.to_string().contains("repeated a page_token"), "{error}" @@ -4147,9 +4213,13 @@ mod tests { .body(r#"{"secrets":[]}"#) .unwrap() }); - conn.list_secrets(&["prod".to_string(), "vision".to_string()]) - .await - .unwrap(); + conn.list_secrets( + &["prod".to_string(), "vision".to_string()], + Default::default(), + ) + .try_collect::>() + .await + .unwrap(); } /// A view description carries the schema in the namespace spec's JSON @@ -4314,7 +4384,10 @@ mod tests { http::Response::builder().status(200).body(body).unwrap() }); assert_eq!( - conn.list_views(&["analytics".into()]).await.unwrap(), + conn.list_views(&["analytics".into()], Default::default()) + .try_collect::>() + .await + .unwrap(), vec!["adults".to_string()] ); } @@ -4329,7 +4402,11 @@ mod tests { .body(r#"{"views":["adults"],"page_token":"same"}"#) .unwrap() }); - let error = conn.list_views(&[]).await.unwrap_err(); + let error = conn + .list_views(&[], Default::default()) + .try_collect::>() + .await + .unwrap_err(); assert!( error.to_string().contains("repeated a page_token"), "{error}" @@ -4451,7 +4528,11 @@ mod tests { } } }); - let functions = conn.list_functions(&[]).await.unwrap(); + let functions = conn + .list_functions(&[], Default::default()) + .try_collect::>() + .await + .unwrap(); assert_eq!(functions.len(), 1); assert_eq!(functions[0].name(), "embed"); assert_eq!(functions[0].version(), "1"); @@ -4474,7 +4555,11 @@ mod tests { .unwrap() }); - let functions = conn.list_functions(&[]).await.unwrap(); + let functions = conn + .list_functions(&[], Default::default()) + .try_collect::>() + .await + .unwrap(); assert!(functions.is_empty()); assert_eq!(requests.load(Ordering::SeqCst), 1); } @@ -4515,17 +4600,14 @@ mod tests { .unwrap() }); - let error = conn.list_functions(&[]).await.unwrap_err(); + let mut items = conn.list_functions(&[], Default::default()); + let error = (&mut items).try_collect::>().await.unwrap_err(); assert!( - matches!( - &error, - Error::Http { - status_code: Some(http::StatusCode::OK), - .. - } - ), + matches!(&error, Error::Runtime { message } if message.contains("repeated a page_token")), "got {error:?}" ); + assert!(items.page_token().is_some()); + assert!(items.try_next().await.unwrap().is_none()); assert_eq!(requests.load(Ordering::SeqCst), 3); } @@ -4603,7 +4685,13 @@ mod tests { .unwrap(); assert_eq!(job.id(), Some("job-function-1")); conn.get_function("embed", "1", &namespace).await.unwrap(); - assert!(conn.list_functions(&namespace).await.unwrap().is_empty()); + assert!( + conn.list_functions(&namespace, Default::default()) + .try_collect::>() + .await + .unwrap() + .is_empty() + ); assert!(conn.drop_function("embed", "1", &namespace).await.unwrap()); } @@ -4616,7 +4704,12 @@ mod tests { }); let namespace = ["analytics".to_string(), String::new()]; assert!(conn.get_function("embed", "1", &namespace).await.is_err()); - assert!(conn.list_functions(&namespace).await.is_err()); + assert!( + conn.list_functions(&namespace, Default::default()) + .try_collect::>() + .await + .is_err() + ); assert!(conn.drop_function("embed", "1", &namespace).await.is_err()); }