From 631ea41a635890218bf7ee1ab80df6caccc7606c Mon Sep 17 00:00:00 2001 From: Colin Patrick McCabe Date: Thu, 22 Oct 2026 15:41:30 -0700 Subject: [PATCH] feat!: return lazy iterators for table and namespace listings (#4599) Make table and namespace listings consistent with other resource listings across Rust, Python, and TypeScript. Callers can iterate over individual names across all pages, with configurable page sizes and continuation tokens. In order to avoid breaking existing Python users, this change adds some compatibility methods to the returned Python iterators. Code such as: ``` for t in db.list_tables().tables: print(t) ``` will continue to work as expected because the iterator has a `tables` field and a `page_token` field, which function as expected to reveal the first page of results. In other words, it is never necessary to use the iterator __as an iterator__. It may be used as if it were an old-style response object. (As the iterator progresses, the `tables` list may switch to the second, third, etc. page of results.) Similar compatibility code allows the async python API to keep working as well. --- docs/src/js/classes/Connection.md | 80 ++------ docs/src/js/globals.md | 4 - .../js/interfaces/ListNamespacesOptions.md | 27 --- .../js/interfaces/ListNamespacesResponse.md | 23 --- docs/src/js/interfaces/ListTablesOptions.md | 35 ---- docs/src/js/interfaces/ListTablesResponse.md | 23 --- docs/src/js/interfaces/TableNamesOptions.md | 2 +- docs/src/python/python.md | 23 ++- docs/web-tests/ts/connection.test.ts | 16 +- docs/web-tests/ts/quickstart.test.ts | 7 +- nodejs/__test__/connection.test.ts | 156 ++++++++-------- nodejs/lancedb/connection.ts | 133 ++++---------- nodejs/lancedb/index.ts | 4 - nodejs/src/connection.rs | 101 ++--------- python/python/lancedb/_lancedb.pyi | 11 +- python/python/lancedb/db.py | 70 ++++--- python/python/lancedb/listing.py | 171 ++++++++++++++++++ python/python/lancedb/namespace.py | 39 ++-- python/python/lancedb/remote/db.py | 16 +- python/python/tests/test_embeddings.py | 4 + python/python/tests/test_listing.py | 137 ++++++++++++++ python/python/tests/test_ollama_embeddings.py | 2 +- python/src/connection.rs | 56 ++---- python/src/listing.rs | 30 ++- rust/lancedb/src/connection.rs | 161 +++++++++-------- rust/lancedb/src/listing.rs | 130 +++++++++---- rust/lancedb/src/remote/db.rs | 66 ++----- 27 files changed, 814 insertions(+), 713 deletions(-) delete mode 100644 docs/src/js/interfaces/ListNamespacesOptions.md delete mode 100644 docs/src/js/interfaces/ListNamespacesResponse.md delete mode 100644 docs/src/js/interfaces/ListTablesOptions.md delete mode 100644 docs/src/js/interfaces/ListTablesResponse.md diff --git a/docs/src/js/classes/Connection.md b/docs/src/js/classes/Connection.md index d5ab8c3d1..c110eb38c 100644 --- a/docs/src/js/classes/Connection.md +++ b/docs/src/js/classes/Connection.md @@ -653,94 +653,48 @@ See [Listing](Listing.md) for pagination state and error behavior. ### listNamespaces() ```ts -abstract listNamespaces(namespacePath?, options?): Promise +abstract listNamespaces(namespacePath?, options?): Listing ``` -List the immediate child namespaces under the given parent. - -Results may be paginated. To retrieve subsequent pages, pass the -`pageToken` returned by a previous call. +Lazily list immediate child namespaces, 100 per page by default. #### Parameters * **namespacePath?**: `string`[] - The parent namespace path. Defaults - to the root namespace if omitted. -* **options?**: `Partial`<[`ListNamespacesOptions`](../interfaces/ListNamespacesOptions.md)> - Pagination options - (`pageToken`, `limit`). +* **options?**: [`ListingOptions`](../interfaces/ListingOptions.md) #### Returns -`Promise`<[`ListNamespacesResponse`](../interfaces/ListNamespacesResponse.md)> - -Child namespace names and - an optional token for fetching the next page. +[`Listing`](Listing.md)<`string`> *** ### listTables() -#### listTables(options) - ```ts -abstract listTables(options?): Promise +abstract listTables(namespacePath?, options?): Listing ``` -List a page of the tables in this database. +Lazily list tables in a namespace, requesting 100 per page by default. -To retrieve the tables after the page, pass the `pageToken` the response -carries back in. A page can be shorter than `limit` without being the last -one, so walk until a response carries no page token: - -```ts -const names = []; -let pageToken = undefined; -do { - const page = await conn.listTables({ pageToken, limit: 100 }); - names.push(...page.tables); - pageToken = page.pageToken; -} while (pageToken); -``` - -##### Parameters - -* **options?**: `Partial`<[`ListTablesOptions`](../interfaces/ListTablesOptions.md)> - Pagination options - (`pageToken`, `limit`). - -##### Returns - -`Promise`<[`ListTablesResponse`](../interfaces/ListTablesResponse.md)> - -A page of table names and an - optional token for the tables after it. - -#### listTables(namespacePath, options) - -```ts -abstract listTables(namespacePath?, options?): Promise -``` - -List a page of the tables in this database. - -##### Parameters +#### Parameters * **namespacePath?**: `string`[] - The namespace path to list tables from - (defaults to root namespace) -* **options?**: `Partial`<[`ListTablesOptions`](../interfaces/ListTablesOptions.md)> - Pagination options - (`pageToken`, `limit`). +* **options?**: [`ListingOptions`](../interfaces/ListingOptions.md) -##### Returns +#### Returns -`Promise`<[`ListTablesResponse`](../interfaces/ListTablesResponse.md)> +[`Listing`](Listing.md)<`string`> -A page of table names and an - optional token for the tables after it. +#### Example + +```ts +for await (const name of conn.listTables([], { pageLimit: 20 })) { + console.log(name); +} +``` *** diff --git a/docs/src/js/globals.md b/docs/src/js/globals.md index 646341190..e3830e9ad 100644 --- a/docs/src/js/globals.md +++ b/docs/src/js/globals.md @@ -108,10 +108,6 @@ - [JobEventsOptions](interfaces/JobEventsOptions.md) - [JobFailureInfo](interfaces/JobFailureInfo.md) - [JobInfo](interfaces/JobInfo.md) -- [ListNamespacesOptions](interfaces/ListNamespacesOptions.md) -- [ListNamespacesResponse](interfaces/ListNamespacesResponse.md) -- [ListTablesOptions](interfaces/ListTablesOptions.md) -- [ListTablesResponse](interfaces/ListTablesResponse.md) - [ListingOptions](interfaces/ListingOptions.md) - [MaterializedViewDefinition](interfaces/MaterializedViewDefinition.md) - [MemtableStats](interfaces/MemtableStats.md) diff --git a/docs/src/js/interfaces/ListNamespacesOptions.md b/docs/src/js/interfaces/ListNamespacesOptions.md deleted file mode 100644 index 79680417f..000000000 --- a/docs/src/js/interfaces/ListNamespacesOptions.md +++ /dev/null @@ -1,27 +0,0 @@ -[**@lancedb/lancedb**](../README.md) • **Docs** - -*** - -[@lancedb/lancedb](../globals.md) / ListNamespacesOptions - -# Interface: ListNamespacesOptions - -## Properties - -### limit? - -```ts -optional limit: number; -``` - -An optional limit to the number of results to return. - -*** - -### pageToken? - -```ts -optional pageToken: string; -``` - -Token from a previous response for pagination. diff --git a/docs/src/js/interfaces/ListNamespacesResponse.md b/docs/src/js/interfaces/ListNamespacesResponse.md deleted file mode 100644 index a99ab480f..000000000 --- a/docs/src/js/interfaces/ListNamespacesResponse.md +++ /dev/null @@ -1,23 +0,0 @@ -[**@lancedb/lancedb**](../README.md) • **Docs** - -*** - -[@lancedb/lancedb](../globals.md) / ListNamespacesResponse - -# Interface: ListNamespacesResponse - -## Properties - -### namespaces - -```ts -namespaces: string[]; -``` - -*** - -### pageToken? - -```ts -optional pageToken: string; -``` diff --git a/docs/src/js/interfaces/ListTablesOptions.md b/docs/src/js/interfaces/ListTablesOptions.md deleted file mode 100644 index b8dd8b5ef..000000000 --- a/docs/src/js/interfaces/ListTablesOptions.md +++ /dev/null @@ -1,35 +0,0 @@ -[**@lancedb/lancedb**](../README.md) • **Docs** - -*** - -[@lancedb/lancedb](../globals.md) / ListTablesOptions - -# Interface: ListTablesOptions - -## Properties - -### limit? - -```ts -optional limit: number; -``` - -An upper bound on how many tables to return. -Defaults to 100. Zero returns an empty page without a continuation token. - -A page may hold fewer than this and still not be the last one, so keep -going while the response carries a page token rather than while pages are -full. - -*** - -### pageToken? - -```ts -optional pageToken: string; -``` - -Token from a previous response, to resume listing where it left off. - -The token is opaque: it carries whatever the database needs to resume, and -callers should not construct or interpret one. diff --git a/docs/src/js/interfaces/ListTablesResponse.md b/docs/src/js/interfaces/ListTablesResponse.md deleted file mode 100644 index 76cac2b23..000000000 --- a/docs/src/js/interfaces/ListTablesResponse.md +++ /dev/null @@ -1,23 +0,0 @@ -[**@lancedb/lancedb**](../README.md) • **Docs** - -*** - -[@lancedb/lancedb](../globals.md) / ListTablesResponse - -# Interface: ListTablesResponse - -## Properties - -### pageToken? - -```ts -optional pageToken: string; -``` - -*** - -### tables - -```ts -tables: string[]; -``` diff --git a/docs/src/js/interfaces/TableNamesOptions.md b/docs/src/js/interfaces/TableNamesOptions.md index 123aa6c9f..e1fd8c382 100644 --- a/docs/src/js/interfaces/TableNamesOptions.md +++ b/docs/src/js/interfaces/TableNamesOptions.md @@ -8,7 +8,7 @@ ## Deprecated -Use [ListTablesOptions](ListTablesOptions.md) with [Connection.listTables](../classes/Connection.md#listtables) +Use [ListingOptions](ListingOptions.md) with [Connection.listTables](../classes/Connection.md#listtables) instead. ## Properties diff --git a/docs/src/python/python.md b/docs/src/python/python.md index f4c697497..c465d0ac4 100644 --- a/docs/src/python/python.md +++ b/docs/src/python/python.md @@ -26,14 +26,17 @@ 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. +Resource listings return iterators. Use `list(db.list_tables())` to collect all +table names. `db.list_tables().tables` is a list of the first page; `limit` +controls its page size and `page_token` resumes at a saved token. The same +patterns apply to `list_namespaces()`. Other resource listings use `page_limit`. ::: lancedb.listing.Listing +::: lancedb.listing.TableListing + +::: lancedb.listing.NamespaceListing + ::: lancedb.Session ## Catalogs (Synchronous) @@ -631,6 +634,16 @@ can be used to create, list, or open tables. ::: lancedb.listing.AsyncListing +`async for name in db.list_tables()` loads pages as needed. To inspect the +first page, use `page = await db.list_tables()` and then `page.tables`. Reading +the page list does not consume the iterator. `page.page_token` is the token for +the next request; drain cached results before using it to resume in a new +listing. The same behavior applies to namespaces. + +::: lancedb.listing.AsyncTableListing + +::: lancedb.listing.AsyncNamespaceListing + ## Namespaces (Asynchronous) ::: lancedb.connect_namespace_async diff --git a/docs/web-tests/ts/connection.test.ts b/docs/web-tests/ts/connection.test.ts index 891501d4c..651be7d36 100644 --- a/docs/web-tests/ts/connection.test.ts +++ b/docs/web-tests/ts/connection.test.ts @@ -4,6 +4,12 @@ import { afterEach, describe, expect, jest, test } from "@jest/globals"; import * as lancedb from "@lancedb/lancedb"; import { withTempDirectory } from "./util.ts"; +async function collectNames(items: AsyncIterable): Promise { + const names: string[] = []; + for await (const name of items) names.push(name); + return names; +} + test("connect to a local database", async () => { await withTempDirectory(async (tempDir) => { const cwd = process.cwd(); @@ -13,12 +19,12 @@ test("connect to a local database", async () => { const db = await lancedb.connect("ex_lancedb"); // --8<-- [end:connect] await db.createTable("probe", [{ id: 1 }]); - expect((await db.listTables()).tables).toEqual(["probe"]); + expect(await collectNames(db.listTables())).toEqual(["probe"]); } finally { process.chdir(cwd); } const reopened = await lancedb.connect(`${tempDir}/ex_lancedb`); - expect((await reopened.listTables()).tables).toEqual(["probe"]); + expect(await collectNames(reopened.listTables())).toEqual(["probe"]); }); }); @@ -109,8 +115,8 @@ async function namespaceTableOpsExample() { ); // Verify - console.log((await db.listNamespaces()).namespaces); // ["prod"] - console.log((await db.listNamespaces(["prod"])).namespaces); // ["recommendations", "search"] + console.log(await collectNames(db.listNamespaces())); // ["prod"] + console.log(await collectNames(db.listNamespaces(["prod"]))); // ["recommendations", "search"] console.log(await db.tableNames(["prod", "search"])); // ["user"] console.log(await db.tableNames(["prod", "recommendations"])); // ["user"] // --8<-- [end:namespace_table_ops] @@ -124,7 +130,7 @@ async function namespaceAdminOpsExample() { await db.createNamespace(["prod"]); await db.createNamespace(["prod", "search"]); - const childNamespaces = (await db.listNamespaces(["prod"])).namespaces; + const childNamespaces = await collectNames(db.listNamespaces(["prod"])); console.log(`Child namespaces under ${JSON.stringify(namespace)}:`, childNamespaces); // Child namespaces under ["prod","search"]: [ 'search' ] diff --git a/docs/web-tests/ts/quickstart.test.ts b/docs/web-tests/ts/quickstart.test.ts index e32a14f86..2c3ee8252 100644 --- a/docs/web-tests/ts/quickstart.test.ts +++ b/docs/web-tests/ts/quickstart.test.ts @@ -20,12 +20,7 @@ const connections = async function tableNames(db: lancedb.Connection) { const names: string[] = []; - let pageToken: string | undefined; - do { - const page = await db.listTables({ pageToken }); - names.push(...page.tables); - pageToken = page.pageToken; - } while (pageToken); + for await (const name of db.listTables()) names.push(name); return names; } diff --git a/nodejs/__test__/connection.test.ts b/nodejs/__test__/connection.test.ts index bbd0c3c52..73f99803a 100644 --- a/nodejs/__test__/connection.test.ts +++ b/nodejs/__test__/connection.test.ts @@ -5,15 +5,15 @@ import { mkdirSync, readdirSync } from "fs"; import { join } from "path"; import { Field, Float64, Schema } from "apache-arrow"; import * as tmp from "tmp"; -import { - Connection, - ListTablesResponse, - Table, - connect, - connectNamespace, -} from "../lancedb"; +import { Connection, Table, connect, connectNamespace } from "../lancedb"; import { LocalTable } from "../lancedb/table"; +async function collectNames(items: AsyncIterable): Promise { + const names: string[] = []; + for await (const name of items) names.push(name); + return names; +} + describe("when connecting", () => { let tmpDir: tmp.DirResult; beforeEach(() => { @@ -135,7 +135,7 @@ describe("given a connection", () => { await db.close(); expect(db.isOpen()).toBe(false); await expect(db.tableNames()).rejects.toThrow("Connection is closed"); - await expect(db.listTables()).rejects.toThrow("Connection is closed"); + expect(() => db.listTables()).toThrow("Connection is closed"); await expect(db.renameTable("a", "b")).rejects.toThrow( "Connection is closed", ); @@ -247,15 +247,15 @@ describe("given a connection", () => { await db.createTable(name, [{ id: 1 }]); } - const first = await db.listTables(namespacePath, { limit: 1 }); - expect(first.tables).toEqual(["a"]); - expect(first.pageToken).toBeDefined(); + const first = db.listTables(namespacePath, { pageLimit: 1 }); + expect(await first.next()).toEqual({ done: false, value: "a" }); + expect(first.pageToken()).toBeDefined(); - const second = await db.listTables(namespacePath, { - limit: 1, - pageToken: first.pageToken, + const second = db.listTables(namespacePath, { + pageLimit: 1, + pageToken: first.pageToken(), }); - expect(second.tables).toEqual(["b"]); + expect(await second.next()).toEqual({ done: false, value: "b" }); }); }); it("should keep the existing table when mode is exist_ok", async () => { @@ -299,19 +299,18 @@ describe("given a connection", () => { await db.createTable("a", [{ id: 1 }]); await db.createTable("c", [{ id: 1 }]); - const all = await db.listTables(); - expect(all.tables).toEqual(["a", "b", "c"]); - expect(all.pageToken).toBeUndefined(); + expect(await collectNames(db.listTables())).toEqual(["a", "b", "c"]); - const first = await db.listTables({ limit: 1 }); - expect(first.tables).toEqual(["a"]); - expect(first.pageToken).toBeDefined(); + const first = db.listTables([], { pageLimit: 1 }); + expect(first.numPageResults()).toBe(0); + expect(await first.next()).toEqual({ done: false, value: "a" }); + expect(first.pageToken()).toBeDefined(); - const second = await db.listTables({ - limit: 1, - pageToken: first.pageToken, + const second = db.listTables([], { + pageLimit: 1, + pageToken: first.pageToken(), }); - expect(second.tables).toEqual(["b"]); + expect(await second.next()).toEqual({ done: false, value: "b" }); }); it("should visit every table exactly once when walking pages", async () => { @@ -322,18 +321,9 @@ describe("given a connection", () => { await db.createTable(name, [{ id: 1 }]); } - const seen: string[] = []; - let pageToken: string | undefined = undefined; - do { - const page: ListTablesResponse = await db.listTables({ - limit: 2, - pageToken, - }); - seen.push(...page.tables); - pageToken = page.pageToken; - } while (pageToken); - - expect(seen).toEqual(created); + expect(await collectNames(db.listTables([], { pageLimit: 2 }))).toEqual( + created, + ); }); it.each([2 ** 32, 2 ** 32 + 1, Number.MAX_SAFE_INTEGER])( @@ -347,11 +337,10 @@ describe("given a connection", () => { ]) { await expect(names).resolves.toEqual(["a", "b"]); } - for (const page of [ - db.listTables({ limit }), - db.listTables([], { limit }), - ]) { - await expect(page).resolves.toEqual({ tables: ["a", "b"] }); + for (const namespacePath of [undefined, []]) { + expect(() => + db.listTables(namespacePath, { pageLimit: limit }), + ).toThrow(/limit/); } }, ); @@ -366,8 +355,11 @@ describe("given a connection", () => { ])("should reject an invalid table listing limit %s", async (limit) => { await expect(db.tableNames({ limit })).rejects.toThrow(/limit/); await expect(db.tableNames([], { limit })).rejects.toThrow(/limit/); - await expect(db.listTables({ limit })).rejects.toThrow(/limit/); - await expect(db.listTables([], { limit })).rejects.toThrow(/limit/); + for (const namespacePath of [undefined, []]) { + expect(() => db.listTables(namespacePath, { pageLimit: limit })).toThrow( + /limit/, + ); + } }); it("should preserve zero, small and unsigned-32-bit table listing limits", async () => { @@ -379,11 +371,21 @@ describe("given a connection", () => { [2 ** 32 - 1, ["a", "b"]], ] as const) { await expect(db.tableNames({ limit })).resolves.toEqual(names); - const page = await db.listTables({ limit }); - expect(page.tables).toEqual(names); } await expect(db.tableNames()).resolves.toEqual(["a", "b"]); - await expect(db.listTables()).resolves.toEqual({ tables: ["a", "b"] }); + for (const namespacePath of [undefined, []]) { + for (const pageLimit of [1, 2 ** 31 - 1]) { + expect( + await collectNames(db.listTables(namespacePath, { pageLimit })), + ).toEqual(["a", "b"]); + } + for (const pageLimit of [0, 2 ** 31, 2 ** 32 - 1]) { + expect(() => db.listTables(namespacePath, { pageLimit })).toThrow( + /limit/, + ); + } + } + expect(await collectNames(db.listTables())).toEqual(["a", "b"]); }); it("should list all names and use a default page size of 100", async () => { @@ -396,20 +398,17 @@ describe("given a connection", () => { } await expect(db.tableNames()).resolves.toEqual(names); await expect(db.tableNames({ limit: 0 })).resolves.toEqual([]); - const first = await db.listTables(); - expect(first.tables).toEqual(names.slice(0, 100)); - expect(first.pageToken).toBeDefined(); - const second = await db.listTables({ pageToken: first.pageToken }); - expect(second.tables).toEqual(names.slice(100)); - expect(second.pageToken).toBeUndefined(); - await expect(db.listTables({ limit: 200 })).resolves.toEqual({ - tables: names, - pageToken: undefined, - }); - await expect(db.listTables({ limit: 0 })).resolves.toEqual({ - tables: [], - pageToken: undefined, - }); + const first = db.listTables(); + const firstPage: string[] = []; + for (let i = 0; i < 100; i++) firstPage.push((await first.next()).value!); + expect(firstPage).toEqual(names.slice(0, 100)); + expect(first.pageToken()).toBeDefined(); + expect(first.numPageResults()).toBe(0); + expect(await collectNames(first)).toEqual(names.slice(100)); + expect(await collectNames(db.listTables([], { pageLimit: 200 }))).toEqual( + names, + ); + expect(() => db.listTables([], { pageLimit: 0 })).toThrow(/limit/); }); it("should list tables in a namespace", async () => { @@ -420,12 +419,8 @@ describe("given a connection", () => { await db.createNamespace(["child"]); await db.createTable("nested", [{ id: 1 }], ["child"]); - await expect(db.listTables(["child"])).resolves.toEqual( - expect.objectContaining({ tables: ["nested"] }), - ); - await expect(db.listTables()).resolves.toEqual( - expect.objectContaining({ tables: [] }), - ); + expect(await collectNames(db.listTables(["child"]))).toEqual(["nested"]); + expect(await collectNames(db.listTables())).toEqual([]); }); it("should create tables in v2 mode", async () => { @@ -719,22 +714,35 @@ describe("namespaces", () => { it("should list namespaces created at the root", async () => { await db.createNamespace(["alpha"]); await db.createNamespace(["beta"]); - const list = await db.listNamespaces(); - expect(list.namespaces).toEqual(expect.arrayContaining(["alpha", "beta"])); + const list = await collectNames(db.listNamespaces()); + expect(list).toEqual(expect.arrayContaining(["alpha", "beta"])); }); it("should list child namespaces under a parent", async () => { await db.createNamespace(["parent"]); await db.createNamespace(["parent", "child"]); - const list = await db.listNamespaces(["parent"]); - expect(list.namespaces).toContain("child"); + const list = await collectNames(db.listNamespaces(["parent"])); + expect(list).toContain("child"); + }); + + it("should resume a namespace listing from its page token", async () => { + await db.createNamespace(["alpha"]); + await db.createNamespace(["beta"]); + const first = db.listNamespaces([], { pageLimit: 1 }); + expect(first.numPageResults()).toBe(0); + expect((await first.next()).value).toBe("alpha"); + const second = db.listNamespaces([], { + pageLimit: 1, + pageToken: first.pageToken(), + }); + expect(await collectNames(second)).toEqual(["beta"]); }); it("should drop a namespace", async () => { await db.createNamespace(["ephemeral"]); await db.dropNamespace(["ephemeral"]); - const list = await db.listNamespaces(); - expect(list.namespaces).not.toContain("ephemeral"); + const list = await collectNames(db.listNamespaces()); + expect(list).not.toContain("ephemeral"); }); it("should raise an error on any namespace op after close", async () => { @@ -742,7 +750,7 @@ describe("namespaces", () => { await expect(db.describeNamespace(["foo"])).rejects.toThrow( "Connection is closed", ); - await expect(db.listNamespaces()).rejects.toThrow("Connection is closed"); + expect(() => db.listNamespaces()).toThrow("Connection is closed"); await expect(db.createNamespace(["foo"])).rejects.toThrow( "Connection is closed", ); diff --git a/nodejs/lancedb/connection.ts b/nodejs/lancedb/connection.ts index c93414fab..5f067eeef 100644 --- a/nodejs/lancedb/connection.ts +++ b/nodejs/lancedb/connection.ts @@ -30,16 +30,12 @@ import type { DescribeNamespaceResponse, DropNamespaceResponse, JobInfo, - ListNamespacesResponse, - ListTablesResponse, } from "./native"; import { ViewDescription, viewDescriptionFromNative } from "./view"; export type { CreateNamespaceResponse, DescribeNamespaceResponse, DropNamespaceResponse, - ListNamespacesResponse, - ListTablesResponse, }; import { sanitizeTable } from "./sanitize"; import { LocalTable, Table } from "./table"; @@ -141,7 +137,7 @@ export interface OpenTableOptions { } /** - * @deprecated Use {@link ListTablesOptions} with {@link Connection.listTables} + * @deprecated Use {@link ListingOptions} with {@link Connection.listTables} * instead. */ export interface TableNamesOptions { @@ -157,32 +153,6 @@ export interface TableNamesOptions { limit?: number; } -export interface ListTablesOptions { - /** - * Token from a previous response, to resume listing where it left off. - * - * The token is opaque: it carries whatever the database needs to resume, and - * callers should not construct or interpret one. - */ - pageToken?: string; - /** - * An upper bound on how many tables to return. - * Defaults to 100. Zero returns an empty page without a continuation token. - * - * A page may hold fewer than this and still not be the last one, so keep - * going while the response carries a page token rather than while pages are - * full. - */ - limit?: number; -} - -export interface ListNamespacesOptions { - /** Token from a previous response for pagination. */ - pageToken?: string; - /** An optional limit to the number of results to return. */ - limit?: number; -} - export interface CreateNamespaceOptions { /** Creation mode. */ mode?: "create" | "exist_ok" | "overwrite"; @@ -278,45 +248,18 @@ export abstract class Connection { options?: Partial, ): Promise; - /** - * List a page of the tables in this database. - * - * To retrieve the tables after the page, pass the `pageToken` the response - * carries back in. A page can be shorter than `limit` without being the last - * one, so walk until a response carries no page token: - * + /** Lazily list tables in a namespace, requesting 100 per page by default. + * @example * ```ts - * const names = []; - * let pageToken = undefined; - * do { - * const page = await conn.listTables({ pageToken, limit: 100 }); - * names.push(...page.tables); - * pageToken = page.pageToken; - * } while (pageToken); + * for await (const name of conn.listTables([], { pageLimit: 20 })) { + * console.log(name); + * } * ``` - * - * @param {Partial} options - Pagination options - * (`pageToken`, `limit`). - * @returns {Promise} A page of table names and an - * optional token for the tables after it. - */ - abstract listTables( - options?: Partial, - ): Promise; - /** - * List a page of the tables in this database. - * - * @param {string[]} namespacePath - The namespace path to list tables from - * (defaults to root namespace) - * @param {Partial} options - Pagination options - * (`pageToken`, `limit`). - * @returns {Promise} A page of table names and an - * optional token for the tables after it. */ abstract listTables( namespacePath?: string[], - options?: Partial, - ): Promise; + options?: ListingOptions, + ): Listing; /** * Open a table in the database. @@ -542,23 +485,11 @@ export abstract class Connection { namespacePath: string[], ): Promise; - /** - * List the immediate child namespaces under the given parent. - * - * Results may be paginated. To retrieve subsequent pages, pass the - * `pageToken` returned by a previous call. - * - * @param {string[]} namespacePath - The parent namespace path. Defaults - * to the root namespace if omitted. - * @param {Partial} options - Pagination options - * (`pageToken`, `limit`). - * @returns {Promise} Child namespace names and - * an optional token for fetching the next page. - */ + /** Lazily list immediate child namespaces, 100 per page by default. */ abstract listNamespaces( namespacePath?: string[], - options?: Partial, - ): Promise; + options?: ListingOptions, + ): Listing; /** * Create a new namespace at the given path. @@ -821,22 +752,17 @@ export class LocalConnection extends Connection { ); } - async listTables( - namespacePathOrOptions?: string[] | Partial, - options?: Partial, - ): Promise { - // Detect if first argument is namespacePath array or options object - const namespacePath = Array.isArray(namespacePathOrOptions) - ? namespacePathOrOptions - : undefined; - const listTablesOptions = Array.isArray(namespacePathOrOptions) - ? options - : (namespacePathOrOptions ?? options); - - return this.inner.listTables( - namespacePath ?? [], - listTablesOptions?.pageToken, - normalizeTableListingLimit(listTablesOptions?.limit), + listTables( + namespacePath?: string[], + options: ListingOptions = {}, + ): Listing { + validateListingOptions(options); + return new Listing( + this.inner.listTables( + namespacePath ?? [], + options.pageToken, + options.pageLimit, + ), ); } @@ -1039,12 +965,15 @@ export class LocalConnection extends Connection { listNamespaces( namespacePath?: string[], - options?: Partial, - ): Promise { - return this.inner.listNamespaces( - namespacePath ?? [], - options?.pageToken, - options?.limit, + options: ListingOptions = {}, + ): Listing { + validateListingOptions(options); + return new Listing( + this.inner.listNamespaces( + namespacePath ?? [], + options.pageToken, + options.pageLimit, + ), ); } diff --git a/nodejs/lancedb/index.ts b/nodejs/lancedb/index.ts index 4ad8b2a93..0f1b5f333 100644 --- a/nodejs/lancedb/index.ts +++ b/nodejs/lancedb/index.ts @@ -97,13 +97,9 @@ export { Connection, CreateTableOptions, TableNamesOptions, - ListTablesOptions, OpenTableOptions, - ListNamespacesOptions, CreateNamespaceOptions, DropNamespaceOptions, - ListNamespacesResponse, - ListTablesResponse, CreateNamespaceResponse, DropNamespaceResponse, DescribeNamespaceResponse, diff --git a/nodejs/src/connection.rs b/nodejs/src/connection.rs index 67c4706e4..914f9e4c9 100644 --- a/nodejs/src/connection.rs +++ b/nodejs/src/connection.rs @@ -16,8 +16,7 @@ use crate::table::Table; use lancedb::connection::{ConnectBuilder, Connection as LanceDBConnection, connect_namespace}; use lance_namespace::models::{ - CreateNamespaceRequest, DescribeNamespaceRequest, DropNamespaceRequest, ListNamespacesRequest, - ListTablesRequest, + CreateNamespaceRequest, DescribeNamespaceRequest, DropNamespaceRequest, }; use lancedb::ipc::{ipc_file_to_batches, ipc_file_to_schema}; @@ -31,18 +30,6 @@ pub struct DescribeNamespaceResponse { pub properties: Option>, } -#[napi(object)] -pub struct ListNamespacesResponse { - pub namespaces: Vec, - pub page_token: Option, -} - -#[napi(object)] -pub struct ListTablesResponse { - pub tables: Vec, - pub page_token: Option, -} - #[napi(object)] pub struct CreateNamespaceResponse { pub properties: Option>, @@ -239,31 +226,20 @@ impl Connection { op.execute().await.default_error() } - /// List a page of tables in the database. + /// Lazily list tables in the database. #[napi(catch_unwind)] - pub async fn list_tables( + pub fn list_tables( &self, namespace_path: Option>, page_token: Option, limit: Option, - ) -> napi::Result { - let request = ListTablesRequest { - // The root namespace is an empty path, not an absent one: a namespace-backed - // database rejects a request that names no namespace. - id: Some(namespace_path.unwrap_or_default()), - page_token, - limit: limit.map(|limit| i32::try_from(limit).unwrap_or(i32::MAX)), - ..Default::default() - }; - let response = self - .get_inner()? - .list_tables(request) - .await - .default_error()?; - Ok(ListTablesResponse { - tables: response.tables, - page_token: response.page_token, - }) + ) -> napi::Result { + Ok(crate::listing::NameListing::new( + self.get_inner()?.list_tables( + &namespace_path.unwrap_or_default(), + crate::listing::options(page_token, limit), + ), + )) } /// Create table from a Apache Arrow IPC (file) buffer. @@ -663,28 +639,19 @@ impl Connection { } #[napi(catch_unwind)] - /// List child namespaces under the given namespace path - pub async fn list_namespaces( + /// Lazily list child namespaces under the given namespace path + pub fn list_namespaces( &self, namespace_path: Option>, page_token: Option, limit: Option, - ) -> napi::Result { - let req = ListNamespacesRequest { - id: namespace_path, - page_token, - limit: checked_limit(limit)?, - ..Default::default() - }; - let resp = self - .get_inner()? - .list_namespaces(req) - .await - .default_error()?; - Ok(ListNamespacesResponse { - namespaces: resp.namespaces, - page_token: resp.page_token, - }) + ) -> napi::Result { + Ok(crate::listing::NameListing::new( + self.get_inner()?.list_namespaces( + &namespace_path.unwrap_or_default(), + crate::listing::options(page_token, limit), + ), + )) } #[napi(catch_unwind)] @@ -787,33 +754,3 @@ impl Connection { .default_error() } } - -fn checked_limit(limit: Option) -> napi::Result> { - limit - .map(i32::try_from) - .transpose() - .map_err(|_| napi::Error::from_reason("limit must not exceed 2147483647")) -} - -#[cfg(test)] -mod tests { - use super::checked_limit; - - #[test] - fn checked_limit_rejects_values_above_i32_max() { - // u32 values above i32::MAX wrapped negative through `as i32`, - // sending a negative page size to the server. - assert!(checked_limit(Some(3_000_000_000)).is_err()); - assert!(checked_limit(Some(u32::MAX)).is_err()); - } - - #[test] - fn checked_limit_passes_boundary_and_missing_values() { - assert_eq!( - checked_limit(Some(i32::MAX as u32)).unwrap(), - Some(i32::MAX) - ); - assert_eq!(checked_limit(Some(0)).unwrap(), Some(0)); - assert_eq!(checked_limit(None).unwrap(), None); - } -} diff --git a/python/python/lancedb/_lancedb.pyi b/python/python/lancedb/_lancedb.pyi index 96eef064c..6ef23abc7 100644 --- a/python/python/lancedb/_lancedb.pyi +++ b/python/python/lancedb/_lancedb.pyi @@ -19,11 +19,9 @@ from .index import ( FTS, ) from lance_namespace import ( - ListNamespacesResponse, CreateNamespaceResponse, DropNamespaceResponse, DescribeNamespaceResponse, - ListTablesResponse, ) from .remote import ClientConfig from .types import BaseTokenizerType @@ -114,12 +112,12 @@ class Connection(object): uri: str async def is_open(self): ... async def close(self): ... - async def list_namespaces( + def list_namespaces( self, namespace_path: Optional[List[str]] = None, page_token: Optional[str] = None, limit: Optional[int] = None, - ) -> ListNamespacesResponse: ... + ) -> NameListing: ... async def create_namespace( self, namespace_path: List[str], @@ -136,12 +134,12 @@ class Connection(object): self, namespace_path: List[str], ) -> DescribeNamespaceResponse: ... - async def list_tables( + def list_tables( self, namespace_path: Optional[List[str]] = None, page_token: Optional[str] = None, limit: Optional[int] = None, - ) -> ListTablesResponse: ... + ) -> NameListing: ... async def table_names( self, namespace_path: Optional[List[str]], @@ -1010,6 +1008,7 @@ async def connect_catalog( class NameListing: def num_page_results(self) -> int: ... def page_token(self) -> Optional[str]: ... + async def next_page(self) -> Optional[dict[str, Any]]: ... def __aiter__(self) -> NameListing: ... async def __anext__(self) -> str: ... diff --git a/python/python/lancedb/db.py b/python/python/lancedb/db.py index a74763afb..90dc1cccf 100644 --- a/python/python/lancedb/db.py +++ b/python/python/lancedb/db.py @@ -32,18 +32,24 @@ if sys.version_info >= (3, 12): else: from overrides import EnforceOverrides, override # type: ignore -from lancedb.listing import AsyncListing, Listing +from lancedb.listing import ( + AsyncListing, + AsyncNamespaceListing, + AsyncTableListing, + Listing, + NamespaceListing, + TableListing, + validate_page_limit, +) from lancedb.embeddings.registry import EmbeddingFunctionRegistry from lancedb.common import data_to_reader, sanitize_uri, validate_schema from lancedb.background_loop import LOOP from lance_namespace import ( LanceNamespace, - ListNamespacesResponse, CreateNamespaceResponse, DropNamespaceResponse, DescribeNamespaceResponse, - ListTablesResponse, connect as namespace_connect, ) from lance_namespace.errors import NamespaceNotEmptyError, TableNotFoundError @@ -171,7 +177,7 @@ class DBConnection(EnforceOverrides): namespace_path: Optional[List[str]] = None, page_token: Optional[str] = None, limit: Optional[int] = None, - ) -> ListNamespacesResponse: + ) -> NamespaceListing: """List immediate child namespace names in the given namespace. Parameters @@ -187,12 +193,12 @@ class DBConnection(EnforceOverrides): Returns ------- - ListNamespacesResponse + NamespaceListing Response containing namespace names and optional page_token for pagination. """ if namespace_path is None: namespace_path = [] - return ListNamespacesResponse(namespaces=[], page_token=None) + return NamespaceListing(AsyncNamespaceListing(None, empty=True)) def create_namespace( self, @@ -317,8 +323,8 @@ class DBConnection(EnforceOverrides): namespace_path: Optional[List[str]] = None, page_token: Optional[str] = None, limit: Optional[int] = None, - ) -> ListTablesResponse: - """List a page of tables in this database. + ) -> TableListing: + """List tables in this database, one page at a time. Parameters ---------- @@ -334,7 +340,7 @@ class DBConnection(EnforceOverrides): Returns ------- - ListTablesResponse + TableListing Response containing table names and optional page_token for pagination. """ raise NotImplementedError( @@ -1436,7 +1442,7 @@ class LanceDBConnection(DBConnection): namespace_path: Optional[List[str]] = None, page_token: Optional[str] = None, limit: Optional[int] = None, - ) -> ListNamespacesResponse: + ) -> NamespaceListing: """List immediate child namespace names in the given namespace. Parameters @@ -1452,16 +1458,14 @@ class LanceDBConnection(DBConnection): Returns ------- - ListNamespacesResponse + NamespaceListing Response containing namespace names and optional page_token for pagination. """ if namespace_path is None: namespace_path = [] - return LOOP.run( + return NamespaceListing( self._conn.list_namespaces( - namespace_path=namespace_path, - page_token=page_token, - limit=limit, + namespace_path=namespace_path, page_token=page_token, limit=limit ) ) @@ -1512,8 +1516,8 @@ class LanceDBConnection(DBConnection): namespace_path: Optional[List[str]] = None, page_token: Optional[str] = None, limit: Optional[int] = None, - ) -> ListTablesResponse: - """List a page of tables in this database. + ) -> TableListing: + """List tables in this database, one page at a time. Parameters ---------- @@ -1529,12 +1533,12 @@ class LanceDBConnection(DBConnection): Returns ------- - ListTablesResponse + TableListing Response containing table names and optional page_token for pagination. """ if namespace_path is None: namespace_path = [] - return LOOP.run( + return TableListing( self._conn.list_tables( namespace_path=namespace_path, page_token=page_token, limit=limit ) @@ -2286,12 +2290,12 @@ class AsyncConnection(object): else: return None - async def list_namespaces( + def list_namespaces( self, namespace_path: Optional[List[str]] = None, page_token: Optional[str] = None, limit: Optional[int] = None, - ) -> ListNamespacesResponse: + ) -> AsyncNamespaceListing: """List immediate child namespace names in the given namespace. Parameters @@ -2306,15 +2310,18 @@ class AsyncConnection(object): Returns ------- - ListNamespacesResponse + AsyncNamespaceListing Response containing namespace names and optional pagination token """ if namespace_path is None: namespace_path = [] - result = await self._inner.list_namespaces( + validate_page_limit(limit) + inner = self._inner.list_namespaces( namespace_path=namespace_path, page_token=page_token, limit=limit ) - return ListNamespacesResponse(**result) + if limit == 0: + return AsyncNamespaceListing(None, empty=True) + return AsyncNamespaceListing(inner, page_token) async def create_namespace( self, @@ -2393,13 +2400,13 @@ class AsyncConnection(object): result = await self._inner.describe_namespace(namespace_path) return DescribeNamespaceResponse(**result) - async def list_tables( + def list_tables( self, namespace_path: Optional[List[str]] = None, page_token: Optional[str] = None, limit: Optional[int] = None, - ) -> ListTablesResponse: - """List a page of tables in this database. + ) -> AsyncTableListing: + """List tables in this database, one page at a time. Parameters ---------- @@ -2415,15 +2422,18 @@ class AsyncConnection(object): Returns ------- - ListTablesResponse + AsyncTableListing Response containing table names and optional page_token for pagination. """ if namespace_path is None: namespace_path = [] - result = await self._inner.list_tables( + validate_page_limit(limit) + inner = self._inner.list_tables( namespace_path=namespace_path, page_token=page_token, limit=limit ) - return ListTablesResponse(**result) + if limit == 0: + return AsyncTableListing(None, empty=True) + return AsyncTableListing(inner, page_token) async def table_names( self, diff --git a/python/python/lancedb/listing.py b/python/python/lancedb/listing.py index 0620aa7d8..8abc9d85b 100644 --- a/python/python/lancedb/listing.py +++ b/python/python/lancedb/listing.py @@ -5,6 +5,7 @@ from __future__ import annotations +import asyncio from typing import Any, AsyncIterator, Callable, Generic, Iterator, Optional, TypeVar from .background_loop import LOOP @@ -83,3 +84,173 @@ class Listing(Iterator[T], Generic[T]): for initial, cached, final-page, and error behavior. """ return self._inner.page_token() + + +def validate_page_limit(limit: Optional[int]) -> None: + """Validate the legacy limit argument before passing it to the bindings.""" + if limit is not None and ( + not isinstance(limit, int) or limit < 0 or limit > 2**31 - 1 + ): + raise ValueError("limit must be between 0 and 2147483647") + + +class _AsyncPagedNameListing(AsyncIterator[str]): + """An async listing with an inspectable copy of its current backend page.""" + + def __init__( + self, inner: Any, page_token: Optional[str] = None, *, empty: bool = False + ): + self._inner = inner + self._token = page_token + self._page: Optional[list[str]] = [] if empty else None + self._position = 0 + self._done = empty + self._failure: Optional[BaseException] = None + self._lock = asyncio.Lock() + + @property + def page_token(self) -> Optional[str]: + """The continuation token for the next backend request.""" + return self._token + + def num_page_results(self) -> int: + """Count the remaining names in the current page.""" + return 0 if self._page is None else len(self._page) - self._position + + def _page_copy(self) -> list[str]: + if self._page is None: + raise RuntimeError("Await or iterate the listing before reading its page") + return self._page.copy() + + async def _fetch_page(self) -> None: + try: + page = await self._inner.next_page() + except asyncio.CancelledError: + raise + except BaseException as exc: + self._failure = exc + self._done = True + raise + if page is None: + self._page = [] + self._token = None + self._done = True + else: + self._page = page["items"] + self._token = page["page_token"] + self._done = self._token is None + self._position = 0 + + async def _load_initial(self): + async with self._lock: + if self._failure is not None: + raise self._failure + if self._page is None: + await self._fetch_page() + return self + + def __await__(self): + return self._load_initial().__await__() + + def __aiter__(self): + return self + + async def __anext__(self) -> str: + async with self._lock: + if self._failure is not None: + raise StopAsyncIteration + if self._page is None: + await self._fetch_page() + assert self._page is not None + while self._position >= len(self._page): + if self._done: + raise StopAsyncIteration + await self._fetch_page() + assert self._page is not None + item = self._page[self._position] + self._position += 1 + return item + + +class AsyncTableListing(_AsyncPagedNameListing): + """Awaitable table iterator; awaiting loads one page. + + For example:: + + page = await db.list_tables(limit=20) + print(page.tables, page.page_token) + async for name in db.list_tables(limit=20): + print(name) + """ + + @property + def tables(self) -> list[str]: + """A copy of every table name in the most recently fetched page.""" + return self._page_copy() + + +class AsyncNamespaceListing(_AsyncPagedNameListing): + """Awaitable namespace iterator; awaiting loads one page. + + For example:: + + page = await db.list_namespaces(limit=20) + print(page.namespaces, page.page_token) + """ + + @property + def namespaces(self) -> list[str]: + """A copy of every name in the most recently fetched page.""" + return self._page_copy() + + +class _PagedNameListing(Iterator[str]): + def __init__(self, inner: _AsyncPagedNameListing): + self._inner = inner + LOOP.run(inner._load_initial()) + + def __iter__(self): + return self + + def __next__(self) -> str: + try: + return LOOP.run(self._inner.__anext__()) + except StopAsyncIteration: + raise StopIteration from None + + @property + def page_token(self) -> Optional[str]: + return self._inner.page_token + + def num_page_results(self) -> int: + return self._inner.num_page_results() + + +class TableListing(_PagedNameListing): + """Table iterator retaining the current page as a list. + + For example:: + + page = db.list_tables(limit=20) + print(page.tables, page.page_token) + for name in db.list_tables(limit=20): + print(name) + """ + + @property + def tables(self) -> list[str]: + return self._inner.tables + + +class NamespaceListing(_PagedNameListing): + """Namespace iterator retaining the current page as a list. + + For example:: + + page = db.list_namespaces(limit=20) + print(page.namespaces, page.page_token) + """ + + @property + def namespaces(self) -> list[str]: + return self._inner.namespaces diff --git a/python/python/lancedb/namespace.py b/python/python/lancedb/namespace.py index a5d88eb1f..55baa39a3 100644 --- a/python/python/lancedb/namespace.py +++ b/python/python/lancedb/namespace.py @@ -52,7 +52,14 @@ 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.listing import ( + AsyncListing, + AsyncNamespaceListing, + AsyncTableListing, + Listing, + NamespaceListing, + TableListing, +) from lancedb.sql import AsyncQuery as AsyncSqlQuery from lancedb.sql import QueryDescription, QueryParameters from lance_namespace import ( @@ -61,8 +68,6 @@ from lance_namespace import ( CreateNamespaceResponse, DescribeNamespaceResponse, DropNamespaceResponse, - ListNamespacesResponse, - ListTablesResponse, NamespaceExistsRequest, TableExistsRequest, ) @@ -854,7 +859,7 @@ class LanceNamespaceDBConnection(DBConnection): namespace_path: Optional[List[str]] = None, page_token: Optional[str] = None, limit: Optional[int] = None, - ) -> ListNamespacesResponse: + ) -> NamespaceListing: """ List child namespaces under the given namespace. @@ -871,12 +876,12 @@ class LanceNamespaceDBConnection(DBConnection): Returns ------- - ListNamespacesResponse + NamespaceListing Response containing namespace names and optional page_token for pagination. """ if namespace_path is None: namespace_path = [] - return LOOP.run( + return NamespaceListing( self._inner.list_namespaces( namespace_path=namespace_path, page_token=page_token, limit=limit ) @@ -1023,7 +1028,7 @@ class LanceNamespaceDBConnection(DBConnection): namespace_path: Optional[List[str]] = None, page_token: Optional[str] = None, limit: Optional[int] = None, - ) -> ListTablesResponse: + ) -> TableListing: """ List all tables in this database with pagination support. @@ -1040,12 +1045,12 @@ class LanceNamespaceDBConnection(DBConnection): Returns ------- - ListTablesResponse + TableListing Response containing table names and optional page_token for pagination. """ if namespace_path is None: namespace_path = [] - return LOOP.run( + return TableListing( self._inner.list_tables( namespace_path=namespace_path, page_token=page_token, limit=limit ) @@ -1459,12 +1464,12 @@ class AsyncLanceNamespaceDBConnection: namespace_path = [] await self._inner.drop_all_tables(namespace_path=namespace_path) - async def list_namespaces( + def list_namespaces( self, namespace_path: Optional[List[str]] = None, page_token: Optional[str] = None, limit: Optional[int] = None, - ) -> ListNamespacesResponse: + ) -> AsyncNamespaceListing: """ List child namespaces under the given namespace. @@ -1481,12 +1486,12 @@ class AsyncLanceNamespaceDBConnection: Returns ------- - ListNamespacesResponse + NamespaceListing Response containing namespace names and optional page_token for pagination. """ if namespace_path is None: namespace_path = [] - return await self._inner.list_namespaces( + return self._inner.list_namespaces( namespace_path=namespace_path, page_token=page_token, limit=limit ) @@ -1616,12 +1621,12 @@ class AsyncLanceNamespaceDBConnection: except TableNotFoundError: return False - async def list_tables( + def list_tables( self, namespace_path: Optional[List[str]] = None, page_token: Optional[str] = None, limit: Optional[int] = None, - ) -> ListTablesResponse: + ) -> AsyncTableListing: """ List all tables in this database with pagination support. @@ -1638,12 +1643,12 @@ class AsyncLanceNamespaceDBConnection: Returns ------- - ListTablesResponse + TableListing Response containing table names and optional page_token for pagination. """ if namespace_path is None: namespace_path = [] - return await self._inner.list_tables( + return self._inner.list_tables( namespace_path=namespace_path, page_token=page_token, limit=limit ) diff --git a/python/python/lancedb/remote/db.py b/python/python/lancedb/remote/db.py index 89c7f8a78..08125cc42 100644 --- a/python/python/lancedb/remote/db.py +++ b/python/python/lancedb/remote/db.py @@ -36,7 +36,7 @@ from lancedb.remote import ClientConfig, RetryConfig, TimeoutConfig, TlsConfig import pyarrow as pa from ..common import DATA -from lancedb.listing import Listing +from lancedb.listing import Listing, NamespaceListing, TableListing from ..db import DBConnection, LOOP from ..functions import FunctionVersion, UdfDefinition from ..job import AsyncJob, Job @@ -54,8 +54,6 @@ from lance_namespace import ( CreateNamespaceResponse, DescribeNamespaceResponse, DropNamespaceResponse, - ListNamespacesResponse, - ListTablesResponse, ) from ..pydantic import LanceModel from ..table import Table @@ -268,7 +266,7 @@ class RemoteDBConnection(DBConnection): namespace_path: Optional[List[str]] = None, page_token: Optional[str] = None, limit: Optional[int] = None, - ) -> ListNamespacesResponse: + ) -> NamespaceListing: """List immediate child namespace names in the given namespace. Parameters @@ -284,12 +282,12 @@ class RemoteDBConnection(DBConnection): Returns ------- - ListNamespacesResponse + NamespaceListing Response containing namespace names and optional page_token for pagination. """ if namespace_path is None: namespace_path = [] - return LOOP.run( + return NamespaceListing( self._conn.list_namespaces( namespace_path=namespace_path, page_token=page_token, limit=limit ) @@ -379,7 +377,7 @@ class RemoteDBConnection(DBConnection): namespace_path: Optional[List[str]] = None, page_token: Optional[str] = None, limit: Optional[int] = None, - ) -> ListTablesResponse: + ) -> TableListing: """List a page of tables in this database. Parameters @@ -396,12 +394,12 @@ class RemoteDBConnection(DBConnection): Returns ------- - ListTablesResponse + TableListing Response containing table names and optional page_token for pagination. """ if namespace_path is None: namespace_path = [] - return LOOP.run( + return TableListing( self._conn.list_tables( namespace_path=namespace_path, page_token=page_token, limit=limit ) diff --git a/python/python/tests/test_embeddings.py b/python/python/tests/test_embeddings.py index e319d3048..207a3d253 100644 --- a/python/python/tests/test_embeddings.py +++ b/python/python/tests/test_embeddings.py @@ -896,6 +896,8 @@ def test_transient_embedding_error_is_retried_in_table_add(tmp_path): assert len(calls) == 3 assert mock_sleep.call_count == 2 assert all(call.args[0] <= 60 for call in mock_sleep.call_args_list) + + @patch("time.sleep") @pytest.mark.parametrize( "error", @@ -941,6 +943,8 @@ def test_retry_on_wrapped_provider_error(mock_sleep): assert wrapped_func() == [0.1] assert mock_func.call_count == 2 assert mock_sleep.call_count == 1 + + def test_openai_bad_request_partial_batch_recovery(): # Regression test for issue #1677 diff --git a/python/python/tests/test_listing.py b/python/python/tests/test_listing.py index 4a2745eb8..b0b973bbf 100644 --- a/python/python/tests/test_listing.py +++ b/python/python/tests/test_listing.py @@ -88,6 +88,143 @@ async def advance(items, asynchronous): KINDS = ["secrets", "views", "jobs", "functions", "materialized_views"] +@pytest.mark.asyncio +@pytest.mark.parametrize("kind", ["tables", "namespaces"]) +@pytest.mark.parametrize("asynchronous", [False, True]) +async def test_compatible_name_listing_pages(listing_server, kind, asynchronous): + endpoint, requests, responses = listing_server + responses.extend( + [ + (200, response_page(kind, ["a", "b"], "next")), + (200, response_page(kind, ["c"], None)), + ] + ) + db = await connection(endpoint, asynchronous) + listing = getattr(db, "list_" + kind)( + namespace_path=["team"], page_token="start", limit=2 + ) + field = kind + assert len(requests) == (0 if asynchronous else 1) + if asynchronous: + with pytest.raises(RuntimeError, match="Await or iterate"): + getattr(listing, field) + assert listing.page_token == "start" + assert await listing is listing + assert len(requests) == 1 + page_names = getattr(listing, field) + assert type(page_names) is list + assert page_names == ["a", "b"] + page_names.append("changed") + assert getattr(listing, field) == ["a", "b"] + assert listing.page_token == "next" + assert listing.num_page_results() == 2 + assert await advance(listing, asynchronous) == "a" + assert getattr(listing, field) == ["a", "b"] + assert await advance(listing, asynchronous) == "b" + assert len(requests) == 1 + assert await advance(listing, asynchronous) == "c" + assert getattr(listing, field) == ["c"] + assert listing.page_token is None + assert len(requests) == 2 + assert ([item async for item in listing] if asynchronous else list(listing)) == [] + if asynchronous: + assert await listing is listing + assert len(requests) == 2 + + +@pytest.mark.asyncio +@pytest.mark.parametrize("kind", ["tables", "namespaces"]) +async def test_name_listing_skips_empty_page_only_during_iteration( + listing_server, kind +): + endpoint, requests, responses = listing_server + responses.extend( + [ + (200, response_page(kind, [], "next")), + (200, response_page(kind, ["a"], None)), + ] + ) + db = await connection(endpoint, True) + listing = getattr(db, "list_" + kind)() + await listing + assert getattr(listing, kind) == [] + assert listing.page_token == "next" + assert len(requests) == 1 + assert [name async for name in listing] == ["a"] + assert len(requests) == 2 + + +@pytest.mark.asyncio +@pytest.mark.parametrize("kind", ["tables", "namespaces"]) +@pytest.mark.parametrize("asynchronous", [False, True]) +async def test_name_listing_zero_limit_uses_no_request( + listing_server, kind, asynchronous +): + endpoint, requests, _ = listing_server + db = await connection(endpoint, asynchronous) + listing = getattr(db, "list_" + kind)(limit=0) + if asynchronous: + assert await listing is listing + assert [name async for name in listing] == [] + else: + assert list(listing) == [] + assert getattr(listing, kind) == [] + assert listing.page_token is None + assert requests == [] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("kind", ["tables", "namespaces"]) +async def test_name_listing_failure_retains_token(listing_server, kind): + endpoint, requests, responses = listing_server + responses.append((401, {})) + db = await connection(endpoint, True) + listing = getattr(db, "list_" + kind)(page_token="retry") + with pytest.raises(HttpError): + await listing + assert listing.page_token == "retry" + with pytest.raises(HttpError): + await listing + assert [name async for name in listing] == [] + assert len(requests) == 1 + + +@pytest.mark.asyncio +@pytest.mark.parametrize("kind", ["tables", "namespaces"]) +async def test_name_listing_direct_async_iteration(listing_server, kind): + endpoint, requests, responses = listing_server + responses.extend( + [ + (200, response_page(kind, ["a"], "next")), + (200, response_page(kind, ["b"], None)), + ] + ) + db = await connection(endpoint, True) + listing = getattr(db, "list_" + kind)(limit=1) + assert requests == [] + assert [name async for name in listing] == ["a", "b"] + assert getattr(listing, kind) == ["b"] + assert len(requests) == 2 + + +@pytest.mark.asyncio +@pytest.mark.parametrize("kind", ["tables", "namespaces"]) +async def test_name_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) + listing = getattr(db, "list_" + kind)() + results = await asyncio.gather( + listing.__anext__(), + listing.__anext__(), + listing.__anext__(), + return_exceptions=True, + ) + assert results[:2] == ["a", "b"] + assert isinstance(results[2], StopAsyncIteration) + assert len(requests) == 1 + + @pytest.mark.asyncio @pytest.mark.parametrize("kind", KINDS) @pytest.mark.parametrize("asynchronous", [False, True]) diff --git a/python/python/tests/test_ollama_embeddings.py b/python/python/tests/test_ollama_embeddings.py index 9d8568ca7..de36dc24b 100644 --- a/python/python/tests/test_ollama_embeddings.py +++ b/python/python/tests/test_ollama_embeddings.py @@ -3,7 +3,6 @@ import json -import httpx import numpy as np import pyarrow as pa import pytest @@ -14,6 +13,7 @@ from lancedb.embeddings import get_registry @pytest.fixture def ollama_embeddings(): + httpx = pytest.importorskip("httpx") pytest.importorskip("ollama") requests = [] diff --git a/python/src/connection.rs b/python/src/connection.rs index baeb29780..2cdf71f4a 100644 --- a/python/src/connection.rs +++ b/python/src/connection.rs @@ -590,29 +590,17 @@ impl Connection { #[pyo3(signature = (namespace_path=None, page_token=None, limit=None))] pub fn list_namespaces( - self_: PyRef<'_, Self>, + &self, namespace_path: Option>, page_token: Option, limit: Option, - ) -> PyResult> { - let inner = self_.get_inner()?.clone(); - let py = self_.py(); - future_into_py(py, async move { - use lance_namespace::models::ListNamespacesRequest; - let request = ListNamespacesRequest { - id: namespace_path, - page_token, - limit: limit.map(|l| l as i32), - ..Default::default() - }; - let response = inner.list_namespaces(request).await.infer_error()?; - Python::attach(|py| -> PyResult> { - let dict = PyDict::new(py); - dict.set_item("namespaces", response.namespaces)?; - dict.set_item("page_token", response.page_token)?; - Ok(dict.unbind()) - }) - }) + ) -> PyResult { + Ok(crate::listing::NameListing::new( + self.get_inner()?.list_namespaces( + &namespace_path.unwrap_or_default(), + crate::listing::options(page_token, limit), + ), + )) } #[pyo3(signature = (namespace_path, mode=None, properties=None))] @@ -725,29 +713,17 @@ impl Connection { #[pyo3(signature = (namespace_path=None, page_token=None, limit=None))] pub fn list_tables( - self_: PyRef<'_, Self>, + &self, namespace_path: Option>, page_token: Option, limit: Option, - ) -> PyResult> { - let inner = self_.get_inner()?.clone(); - let py = self_.py(); - future_into_py(py, async move { - use lance_namespace::models::ListTablesRequest; - let request = ListTablesRequest { - id: namespace_path, - page_token, - limit: limit.map(|l| l as i32), - ..Default::default() - }; - let response = inner.list_tables(request).await.infer_error()?; - Python::attach(|py| -> PyResult> { - let dict = PyDict::new(py); - dict.set_item("tables", response.tables)?; - dict.set_item("page_token", response.page_token)?; - Ok(dict.unbind()) - }) - }) + ) -> PyResult { + Ok(crate::listing::NameListing::new( + self.get_inner()?.list_tables( + &namespace_path.unwrap_or_default(), + crate::listing::options(page_token, limit), + ), + )) } /// Get the configuration for constructing an equivalent namespace client. diff --git a/python/src/listing.rs b/python/src/listing.rs index 2d1a3916e..79a7effdd 100644 --- a/python/src/listing.rs +++ b/python/src/listing.rs @@ -3,10 +3,11 @@ use crate::error::PythonErrorExt; use crate::runtime::future_into_py; -use futures::{StreamExt, future::poll_fn}; +use futures::{Future, StreamExt, future::poll_fn}; use lancedb::listing::{Listing, ListingOptions}; use pyo3::exceptions::PyStopAsyncIteration; -use pyo3::{Bound, PyAny, PyRef, PyResult, Python, pyclass, pymethods}; +use pyo3::types::{PyDict, PyDictMethods}; +use pyo3::{Bound, Py, PyAny, PyRef, PyResult, Python, pyclass, pymethods}; use std::sync::{Arc, Mutex}; pub fn options(page_token: Option, page_limit: Option) -> ListingOptions { @@ -52,6 +53,31 @@ impl NameListing { Ok(value) }) } + + fn next_page(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 page = futures::future::poll_fn(|cx| { + let mut inner = inner.lock().unwrap(); + let future = inner.next_page(); + futures::pin_mut!(future); + future.poll(cx) + }) + .await + .infer_error()?; + Python::attach(|py| -> PyResult>> { + page.map(|page| { + let dict = PyDict::new(py); + dict.set_item("items", page.items)?; + dict.set_item("page_token", page.page_token)?; + Ok(dict.unbind()) + }) + .transpose() + }) + }) + } } #[pyclass] diff --git a/rust/lancedb/src/connection.rs b/rust/lancedb/src/connection.rs index 8a0e489c2..124f6a54a 100644 --- a/rust/lancedb/src/connection.rs +++ b/rust/lancedb/src/connection.rs @@ -13,7 +13,7 @@ use lance::dataset::refs::MAIN_BRANCH; use lance_namespace::models::{ CreateNamespaceRequest, CreateNamespaceResponse, DescribeNamespaceRequest, DescribeNamespaceResponse, DropNamespaceRequest, DropNamespaceResponse, ListNamespacesRequest, - ListNamespacesResponse, ListTablesRequest, ListTablesResponse, + ListTablesRequest, }; #[cfg(feature = "aws")] use object_store::aws::AwsCredential; @@ -28,7 +28,7 @@ use crate::database::{ }; use crate::embeddings::{EmbeddingRegistry, MemoryRegistry}; use crate::error::{Error, Result}; -use crate::listing::{Listing, ListingOptions}; +use crate::listing::{Listing, ListingOptions, ListingPage}; #[cfg(feature = "remote")] use crate::remote::{ client::ClientConfig, @@ -1126,12 +1126,43 @@ impl Connection { self.internal.drop_all_tables(namespace_path).await } - /// List immediate child namespace names in the given namespace - pub async fn list_namespaces( + /// Lazily list immediate child namespace names. Each request asks for at + /// most 100 names unless `options.page_limit` is set. + /// + /// ``` + /// # async fn example(connection: &lancedb::Connection) -> lancedb::Result<()> { + /// use futures::TryStreamExt; + /// let names: Vec = connection.list_namespaces(&[], Default::default()) + /// .try_collect().await?; + /// # Ok(()) + /// # } + /// ``` + pub fn list_namespaces( &self, - request: ListNamespacesRequest, - ) -> Result { - self.internal.list_namespaces(request).await + 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)?; + let response = database + .list_namespaces(ListNamespacesRequest { + id: Some(namespace_path), + page_token: options.page_token, + limit: Some(options.page_limit.unwrap_or(100) as i32), + ..Default::default() + }) + .await?; + Ok(ListingPage { + items: response.namespaces, + page_token: response.page_token, + }) + } + }) } /// Create a new namespace @@ -1182,23 +1213,43 @@ impl Connection { self.internal.namespace_client_config().await } - /// List tables with pagination support + /// Lazily list tables. Each request asks for at most 100 names unless + /// `options.page_limit` is set. /// - /// The default limit is 100 tables per page for both local and remote connections, - /// including namespaces. Zero returns an empty page without a continuation token. - /// Follow the response's opaque `page_token` until it is absent to retrieve every table; - /// a page can contain fewer than the limit even when more tables remain. - pub async fn list_tables(&self, mut request: ListTablesRequest) -> Result { - // Apply the SDK's page size here so every backend receives the same request. - let limit = request.limit.get_or_insert(100); - if *limit <= 0 { - return Ok(ListTablesResponse { - context: None, - tables: Vec::new(), - page_token: None, - }); - } - self.internal.list_tables(request).await + /// ``` + /// # async fn example(connection: &lancedb::Connection) -> lancedb::Result<()> { + /// use futures::TryStreamExt; + /// let names: Vec = connection.list_tables(&[], Default::default()) + /// .try_collect().await?; + /// # Ok(()) + /// # } + /// ``` + pub fn list_tables( + &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)?; + let response = database + .list_tables(ListTablesRequest { + id: Some(namespace_path), + page_token: options.page_token, + limit: Some(options.page_limit.unwrap_or(100) as i32), + ..Default::default() + }) + .await?; + Ok(ListingPage { + items: response.tables, + page_token: response.page_token, + }) + } + }) } /// Get the in-memory embedding registry. @@ -2260,26 +2311,11 @@ mod tests { } names.sort(); - // Walking in pages has to reach every table exactly once, with nothing lost at a - // page boundary. - let mut seen = Vec::with_capacity(names.len()); - let mut page_token = None; - loop { - let page = db - .list_tables(ListTablesRequest { - id: Some(Vec::new()), - limit: Some(2), - page_token, - ..Default::default() - }) - .await - .unwrap(); - seen.extend(page.tables); - page_token = page.page_token.filter(|token| !token.is_empty()); - if page_token.is_none() { - break; - } - } + let seen = futures::TryStreamExt::try_collect::>( + db.list_tables(&[], ListingOptions::default().page_limit(2)), + ) + .await + .unwrap(); assert_eq!(seen, names); } @@ -2305,37 +2341,20 @@ mod tests { .is_empty() ); - let first = db.list_tables(ListTablesRequest::default()).await.unwrap(); - assert_eq!(first.tables, names[..100]); + let mut listing = db.list_tables(&[], ListingOptions::default()); + assert_eq!(listing.num_page_results(), 0); + let first = listing.next_page().await.unwrap().unwrap(); + assert_eq!(first.items, names[..100]); assert!(first.page_token.is_some()); - let second = db - .list_tables(ListTablesRequest { - page_token: first.page_token, - ..Default::default() - }) - .await - .unwrap(); - assert_eq!(second.tables, names[100..]); + let second = listing.next_page().await.unwrap().unwrap(); + assert_eq!(second.items, names[100..]); assert!(second.page_token.is_none()); + assert!(listing.next_page().await.unwrap().is_none()); - let all = db - .list_tables(ListTablesRequest { - limit: Some(200), - ..Default::default() - }) - .await - .unwrap(); - assert_eq!(all.tables, names); - assert!(all.page_token.is_none()); - let empty = db - .list_tables(ListTablesRequest { - limit: Some(0), - ..Default::default() - }) - .await - .unwrap(); - assert!(empty.tables.is_empty()); - assert!(empty.page_token.is_none()); + let mut all = db.list_tables(&[], ListingOptions::default().page_limit(200)); + assert_eq!(all.next_page().await.unwrap().unwrap().items, names); + let mut invalid = db.list_tables(&[], ListingOptions::default().page_limit(0)); + assert!(invalid.next_page().await.is_err()); } #[tokio::test] diff --git a/rust/lancedb/src/listing.rs b/rust/lancedb/src/listing.rs index 710202d9d..a0f6e5d5e 100644 --- a/rust/lancedb/src/listing.rs +++ b/rust/lancedb/src/listing.rs @@ -51,8 +51,9 @@ type FetchPage = /// 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 +/// No requests are made until the stream is polled or [`Self::next_page`] is called. +/// Item iteration skips empty pages with continuation tokens; `next_page` returns +/// them. Errors terminate the listing 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`] @@ -100,6 +101,64 @@ impl Listing { pub fn page_token(&self) -> Option<&str> { self.options.page_token.as_deref() } + + /// Return exactly one page without consuming a following page. + /// + /// If item iteration has started, this returns the remaining items in its + /// cached page. An empty page with a continuation token is returned as is. + pub async fn next_page(&mut self) -> Result>> { + if self.items.len() == 0 { + if self.done { + return Ok(None); + } + futures::future::poll_fn(|cx| self.poll_fetch_page(cx)).await?; + } + let items = std::mem::replace(&mut self.items, Vec::new().into_iter()).collect(); + Ok(Some(ListingPage { + items, + page_token: self.options.page_token.clone(), + })) + } + + fn poll_fetch_page(&mut self, cx: &mut Context<'_>) -> Poll> { + if self + .options + .page_limit + .is_some_and(|limit| limit == 0 || limit > i32::MAX as u32) + { + self.done = true; + return Poll::Ready(Err(Error::InvalidInput { + message: "Listing page limit must be between 1 and 2147483647".into(), + })); + } + let pending = self + .pending + .get_or_insert_with(|| (self.fetch)(self.options.clone())); + let response = futures::ready!(pending.as_mut().poll(cx)); + self.pending = None; + match response { + Ok(response) => { + let token = response.page_token.filter(|token| !token.is_empty()); + if token + .as_ref() + .is_some_and(|token| !self.seen_tokens.insert(token.clone())) + { + self.done = true; + return Poll::Ready(Err(Error::Runtime { + message: "Listing response repeated a page_token".into(), + })); + } + self.done = token.is_none(); + self.options.page_token = token; + self.items = response.items.into_iter(); + Poll::Ready(Ok(())) + } + Err(error) => { + self.done = true; + Poll::Ready(Err(error)) + } + } + } } impl Stream for Listing { @@ -114,41 +173,8 @@ impl Stream for Listing { 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))); - } + if let Err(error) = futures::ready!(this.poll_fetch_page(cx)) { + return Poll::Ready(Some(Err(error))); } } } @@ -224,6 +250,36 @@ mod tests { assert!(listing.next().await.is_none()); } + #[tokio::test] + async fn next_page_preserves_empty_pages_and_remaining_items() { + let requests = Arc::new(Mutex::new(0)); + let count = requests.clone(); + let mut listing = Listing::new(ListingOptions::default(), move |options| { + *count.lock().unwrap() += 1; + async move { + let (items, page_token) = match options.page_token.as_deref() { + None => (vec![], Some("second".to_string())), + Some("second") => (vec!["a", "b"], Some("last".to_string())), + Some("last") => (vec!["c"], None), + _ => unreachable!(), + }; + Ok(ListingPage { items, page_token }) + } + }); + assert_eq!(*requests.lock().unwrap(), 0); + assert!(listing.next_page().await.unwrap().unwrap().items.is_empty()); + assert_eq!(*requests.lock().unwrap(), 1); + assert_eq!(listing.try_next().await.unwrap(), Some("a")); + assert_eq!(*requests.lock().unwrap(), 2); + let remaining = listing.next_page().await.unwrap().unwrap(); + assert_eq!(remaining.items, vec!["b"]); + assert_eq!(remaining.page_token.as_deref(), Some("last")); + assert_eq!(*requests.lock().unwrap(), 2); + assert_eq!(listing.next_page().await.unwrap().unwrap().items, vec!["c"]); + assert!(listing.next_page().await.unwrap().is_none()); + assert_eq!(*requests.lock().unwrap(), 3); + } + #[tokio::test] async fn invalid_limits_do_not_fetch() { for limit in [0, u32::MAX] { diff --git a/rust/lancedb/src/remote/db.rs b/rust/lancedb/src/remote/db.rs index 4a8cbf32e..8a456ab3b 100644 --- a/rust/lancedb/src/remote/db.rs +++ b/rust/lancedb/src/remote/db.rs @@ -2349,16 +2349,8 @@ mod tests { .unwrap() .is_empty() ); - let page = conn - .list_tables(lance_namespace::models::ListTablesRequest { - id: Some(namespace), - limit: Some(0), - ..Default::default() - }) - .await - .unwrap(); - assert!(page.tables.is_empty()); - assert!(page.page_token.is_none()); + let mut listing = conn.list_tables(&namespace, ListingOptions::default().page_limit(0)); + assert!(listing.next_page().await.is_err()); } #[rstest::rstest] @@ -2393,25 +2385,17 @@ mod tests { ) .unwrap() }); - let first = conn - .list_tables(lance_namespace::models::ListTablesRequest { - id: Some(namespace.clone()), - ..Default::default() - }) - .await - .unwrap(); - assert_eq!(first.tables.len(), 100); + let mut first_listing = conn.list_tables(&namespace, ListingOptions::default()); + let first = first_listing.next_page().await.unwrap().unwrap(); + assert_eq!(first.items.len(), 100); assert_eq!(first.page_token.as_deref(), Some("opaque-token")); - let second = conn - .list_tables(lance_namespace::models::ListTablesRequest { - id: Some(namespace), - page_token: first.page_token, - ..Default::default() - }) - .await - .unwrap(); + let mut second_listing = conn.list_tables( + &namespace, + ListingOptions::default().page_token(first.page_token.unwrap()), + ); + let second = second_listing.next_page().await.unwrap().unwrap(); let expected: Vec<_> = (100..130).map(|i| format!("t{i:03}")).collect(); - assert_eq!(second.tables, expected); + assert_eq!(second.items, expected); assert!(second.page_token.is_none()); } @@ -3680,14 +3664,9 @@ mod tests { assert!(table.is_ok(), "Failed to create table: {:?}", table.err()); // List tables in the child namespace - let list_response = conn - .list_tables(ListTablesRequest { - id: Some(namespace.clone()), - ..Default::default() - }) - .await - .expect("Failed to list tables"); - assert_eq!(list_response.tables, vec!["test_table"]); + let list_response = conn.list_tables(&namespace, ListingOptions::default()); + let names = list_response.try_collect::>().await.unwrap(); + assert_eq!(names, vec!["test_table"]); // Get namespace client and verify it can also list tables let namespace_client = conn.namespace_client().await.unwrap(); @@ -3754,17 +3733,12 @@ mod tests { } // List tables in the child namespace - let list_response = conn - .list_tables(ListTablesRequest { - id: Some(namespace.clone()), - ..Default::default() - }) - .await - .unwrap(); - assert_eq!(list_response.tables.len(), 3); - assert!(list_response.tables.contains(&"table1".to_string())); - assert!(list_response.tables.contains(&"table2".to_string())); - assert!(list_response.tables.contains(&"table3".to_string())); + let list_response = conn.list_tables(&namespace, ListingOptions::default()); + let names = list_response.try_collect::>().await.unwrap(); + assert_eq!(names.len(), 3); + assert!(names.contains(&"table1".to_string())); + assert!(names.contains(&"table2".to_string())); + assert!(names.contains(&"table3".to_string())); } }