mirror of
https://github.com/lancedb/lancedb.git
synced 2026-10-12 23:02:28 +00:00
refactor!: implement lazy iterators for several APIs (#4459)
Implement lazy iterators for: - secrets - views - jobs - functions - materialized views
This commit is contained in:
@@ -611,30 +611,42 @@ Return true if the connection has not been closed
|
||||
### listJobs()
|
||||
|
||||
```ts
|
||||
abstract listJobs(): Promise<JobInfo[]>
|
||||
abstract listJobs(options?): Listing<JobInfo>
|
||||
```
|
||||
|
||||
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<string[]>
|
||||
abstract listMaterializedViews(options?): Listing<string>
|
||||
```
|
||||
|
||||
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<string[]>
|
||||
abstract listViews(namespacePath?, options?): Listing<string>
|
||||
```
|
||||
|
||||
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`>
|
||||
|
||||
***
|
||||
|
||||
|
||||
@@ -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<T>
|
||||
```
|
||||
|
||||
#### Returns
|
||||
|
||||
`AsyncIterableIterator`<`T`>
|
||||
|
||||
#### Implementation of
|
||||
|
||||
`AsyncIterableIterator.[asyncIterator]`
|
||||
|
||||
***
|
||||
|
||||
### next()
|
||||
|
||||
```ts
|
||||
next(): Promise<IteratorResult<T, any>>
|
||||
```
|
||||
|
||||
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`
|
||||
@@ -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)
|
||||
|
||||
@@ -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.
|
||||
@@ -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
|
||||
|
||||
@@ -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<T>(items: AsyncIterable<T>): Promise<T[]> {
|
||||
const result: T[] = [];
|
||||
for await (const item of items) result.push(item);
|
||||
return result;
|
||||
}
|
||||
|
||||
@@ -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<T>(items: AsyncIterable<T>): Promise<T[]> {
|
||||
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<string, unknown> }[] = [];
|
||||
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);
|
||||
},
|
||||
);
|
||||
});
|
||||
}
|
||||
});
|
||||
|
||||
@@ -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<string[]>;
|
||||
abstract listMaterializedViews(options?: ListingOptions): Listing<string>;
|
||||
|
||||
/**
|
||||
* 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<string[]>;
|
||||
abstract listViews(
|
||||
namespacePath?: string[],
|
||||
options?: ListingOptions,
|
||||
): Listing<string>;
|
||||
|
||||
abstract openTable(
|
||||
name: string,
|
||||
@@ -637,8 +645,12 @@ export abstract class Connection {
|
||||
*/
|
||||
abstract openJob(jobId: string): Promise<Job>;
|
||||
|
||||
/** List server-side jobs across the database's tables. */
|
||||
abstract listJobs(): Promise<JobInfo[]>;
|
||||
/**
|
||||
* 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<JobInfo>;
|
||||
|
||||
/**
|
||||
* 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<string[]> {
|
||||
return await this.inner.listMaterializedViews();
|
||||
listMaterializedViews(options: ListingOptions = {}): Listing<string> {
|
||||
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<string[]> {
|
||||
return this.inner.listViews(namespacePath ?? []);
|
||||
listViews(
|
||||
namespacePath?: string[],
|
||||
options: ListingOptions = {},
|
||||
): Listing<string> {
|
||||
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<JobInfo[]> {
|
||||
return this.inner.listJobs();
|
||||
listJobs(options: ListingOptions = {}): Listing<JobInfo> {
|
||||
validateListingOptions(options);
|
||||
return new Listing(
|
||||
this.inner.listJobs(options.pageToken, options.pageLimit),
|
||||
);
|
||||
}
|
||||
|
||||
async cancelJob(jobId: string): Promise<boolean> {
|
||||
|
||||
@@ -649,3 +649,5 @@ export {
|
||||
DatabaseNames,
|
||||
connectCatalog,
|
||||
} from "./catalog";
|
||||
|
||||
export { Listing, ListingOptions } from "./listing";
|
||||
|
||||
@@ -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<T> implements AsyncIterableIterator<T> {
|
||||
/** @hidden */
|
||||
constructor(
|
||||
private readonly inner: {
|
||||
next(): Promise<T | null | undefined>;
|
||||
numPageResults(): number;
|
||||
pageToken(): string | null | undefined;
|
||||
},
|
||||
) {}
|
||||
|
||||
[Symbol.asyncIterator](): AsyncIterableIterator<T> {
|
||||
return this;
|
||||
}
|
||||
|
||||
/** Fetch the next item, requesting another page only when needed. */
|
||||
async next(): Promise<IteratorResult<T>> {
|
||||
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;
|
||||
}
|
||||
}
|
||||
+28
-14
@@ -371,13 +371,15 @@ impl Connection {
|
||||
}
|
||||
|
||||
#[napi(catch_unwind)]
|
||||
pub async fn list_materialized_views(&self) -> napi::Result<Vec<String>> {
|
||||
let views = self
|
||||
.get_inner()?
|
||||
.list_materialized_views()
|
||||
.await
|
||||
.default_error()?;
|
||||
Ok(views)
|
||||
pub fn list_materialized_views(
|
||||
&self,
|
||||
page_token: Option<String>,
|
||||
page_limit: Option<u32>,
|
||||
) -> napi::Result<crate::listing::NameListing> {
|
||||
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<Vec<String>>,
|
||||
) -> napi::Result<Vec<String>> {
|
||||
let ns = namespace_path.unwrap_or_default();
|
||||
self.get_inner()?.list_views(&ns).await.default_error()
|
||||
page_token: Option<String>,
|
||||
page_limit: Option<u32>,
|
||||
) -> napi::Result<crate::listing::NameListing> {
|
||||
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<Vec<crate::job::JobInfo>> {
|
||||
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<String>,
|
||||
page_limit: Option<u32>,
|
||||
) -> napi::Result<crate::listing::JobListing> {
|
||||
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
|
||||
|
||||
@@ -18,6 +18,7 @@ mod header;
|
||||
mod index;
|
||||
mod iterator;
|
||||
mod job;
|
||||
mod listing;
|
||||
pub mod merge;
|
||||
pub mod otel;
|
||||
pub mod permutation;
|
||||
|
||||
@@ -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<String>, page_limit: Option<u32>) -> ListingOptions {
|
||||
let mut options = ListingOptions::default();
|
||||
options.page_token = page_token;
|
||||
options.page_limit = page_limit;
|
||||
options
|
||||
}
|
||||
|
||||
#[napi]
|
||||
pub struct NameListing {
|
||||
inner: Mutex<Listing<String>>,
|
||||
next_lock: futures::lock::Mutex<()>,
|
||||
}
|
||||
impl NameListing {
|
||||
pub fn new(inner: Listing<String>) -> 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<String> {
|
||||
self.inner.lock().unwrap().page_token().map(str::to_owned)
|
||||
}
|
||||
#[napi]
|
||||
pub async fn next(&self) -> napi::Result<Option<String>> {
|
||||
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<Listing<lancedb::database::JobInfo>>,
|
||||
next_lock: futures::lock::Mutex<()>,
|
||||
}
|
||||
impl JobListing {
|
||||
pub fn new(inner: Listing<lancedb::database::JobInfo>) -> 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<String> {
|
||||
self.inner.lock().unwrap().page_token().map(str::to_owned)
|
||||
}
|
||||
#[napi]
|
||||
pub async fn next(&self) -> napi::Result<Option<crate::job::JobInfo>> {
|
||||
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()))
|
||||
}
|
||||
}
|
||||
@@ -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: ...
|
||||
|
||||
+218
-47
@@ -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.
|
||||
|
||||
@@ -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()
|
||||
@@ -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
|
||||
|
||||
@@ -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:
|
||||
|
||||
@@ -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"] == []
|
||||
|
||||
@@ -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 == []
|
||||
@@ -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) == []
|
||||
|
||||
@@ -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
|
||||
):
|
||||
|
||||
@@ -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 == [
|
||||
|
||||
+53
-45
@@ -464,12 +464,16 @@ impl Connection {
|
||||
})
|
||||
}
|
||||
|
||||
pub fn list_materialized_views(self_: PyRef<'_, Self>) -> PyResult<Bound<'_, PyAny>> {
|
||||
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<String>,
|
||||
page_limit: Option<u32>,
|
||||
) -> PyResult<crate::listing::NameListing> {
|
||||
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<Vec<String>>,
|
||||
) -> PyResult<Bound<'_, PyAny>> {
|
||||
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::<PyResult<Vec<_>>>()
|
||||
})
|
||||
page_token: Option<String>,
|
||||
page_limit: Option<u32>,
|
||||
) -> PyResult<crate::listing::FunctionListing> {
|
||||
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<Vec<String>>,
|
||||
) -> PyResult<Bound<'_, PyAny>> {
|
||||
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<String>,
|
||||
page_limit: Option<u32>,
|
||||
) -> PyResult<crate::listing::NameListing> {
|
||||
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<Vec<String>>,
|
||||
) -> PyResult<Bound<'_, PyAny>> {
|
||||
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<String>,
|
||||
page_limit: Option<u32>,
|
||||
) -> PyResult<crate::listing::NameListing> {
|
||||
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<Bound<'_, PyAny>> {
|
||||
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::<Vec<_>>())
|
||||
})
|
||||
#[pyo3(signature = (*, page_token=None, page_limit=None))]
|
||||
pub fn list_jobs(
|
||||
&self,
|
||||
page_token: Option<String>,
|
||||
page_limit: Option<u32>,
|
||||
) -> PyResult<crate::listing::JobListing> {
|
||||
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<Bound<'_, PyAny>> {
|
||||
|
||||
@@ -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::<Connection>()?;
|
||||
m.add_class::<catalog::Catalog>()?;
|
||||
m.add_class::<catalog::DatabaseNames>()?;
|
||||
m.add_class::<listing::NameListing>()?;
|
||||
m.add_class::<listing::FunctionListing>()?;
|
||||
m.add_class::<listing::JobListing>()?;
|
||||
|
||||
m.add_function(wrap_pyfunction!(catalog::connect_catalog, m)?)?;
|
||||
m.add_class::<Session>()?;
|
||||
m.add_class::<Table>()?;
|
||||
|
||||
@@ -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<String>, page_limit: Option<u32>) -> ListingOptions {
|
||||
let mut options = ListingOptions::default();
|
||||
options.page_token = page_token;
|
||||
options.page_limit = page_limit;
|
||||
options
|
||||
}
|
||||
|
||||
#[pyclass]
|
||||
pub struct NameListing {
|
||||
inner: Arc<Mutex<Listing<String>>>,
|
||||
next_lock: Arc<tokio::sync::Mutex<()>>,
|
||||
}
|
||||
impl NameListing {
|
||||
pub fn new(inner: Listing<String>) -> 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<String> {
|
||||
py.detach(|| self.inner.lock().unwrap().page_token().map(str::to_owned))
|
||||
}
|
||||
fn __aiter__(self_: PyRef<'_, Self>) -> PyRef<'_, Self> {
|
||||
self_
|
||||
}
|
||||
fn __anext__(self_: PyRef<'_, Self>) -> PyResult<Bound<'_, PyAny>> {
|
||||
let inner = self_.inner.clone();
|
||||
let next_lock = self_.next_lock.clone();
|
||||
future_into_py(self_.py(), async move {
|
||||
let _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<Mutex<Listing<lancedb::function::FunctionVersion>>>,
|
||||
next_lock: Arc<tokio::sync::Mutex<()>>,
|
||||
}
|
||||
impl FunctionListing {
|
||||
pub fn new(inner: Listing<lancedb::function::FunctionVersion>) -> 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<String> {
|
||||
py.detach(|| self.inner.lock().unwrap().page_token().map(str::to_owned))
|
||||
}
|
||||
fn __aiter__(self_: PyRef<'_, Self>) -> PyRef<'_, Self> {
|
||||
self_
|
||||
}
|
||||
fn __anext__(self_: PyRef<'_, Self>) -> PyResult<Bound<'_, PyAny>> {
|
||||
let inner = self_.inner.clone();
|
||||
let next_lock = self_.next_lock.clone();
|
||||
future_into_py(self_.py(), async move {
|
||||
let _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<Mutex<Listing<lancedb::database::JobInfo>>>,
|
||||
next_lock: Arc<tokio::sync::Mutex<()>>,
|
||||
}
|
||||
impl JobListing {
|
||||
pub fn new(inner: Listing<lancedb::database::JobInfo>) -> 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<String> {
|
||||
py.detach(|| self.inner.lock().unwrap().page_token().map(str::to_owned))
|
||||
}
|
||||
fn __aiter__(self_: PyRef<'_, Self>) -> PyRef<'_, Self> {
|
||||
self_
|
||||
}
|
||||
fn __anext__(self_: PyRef<'_, Self>) -> PyResult<Bound<'_, PyAny>> {
|
||||
let inner = self_.inner.clone();
|
||||
let next_lock = self_.next_lock.clone();
|
||||
future_into_py(self_.py(), async move {
|
||||
let _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))
|
||||
})
|
||||
}
|
||||
}
|
||||
+102
-19
@@ -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<dyn std::error::Error>> {
|
||||
/// 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<Vec<crate::function::FunctionVersion>> {
|
||||
self.internal.list_functions(namespace_path).await
|
||||
options: ListingOptions,
|
||||
) -> Listing<crate::function::FunctionVersion> {
|
||||
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<Vec<String>> {
|
||||
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<String> {
|
||||
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<Vec<String>> {
|
||||
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<String> {
|
||||
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<Vec<JobInfo>> {
|
||||
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<JobInfo> {
|
||||
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
|
||||
|
||||
@@ -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> {
|
||||
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<Vec<String>> {
|
||||
async fn list_materialized_views(
|
||||
&self,
|
||||
namespace_path: &[String],
|
||||
options: ListingOptions,
|
||||
) -> Result<ListingPage<String>> {
|
||||
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<crate::function::FunctionVersion> {
|
||||
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<Vec<crate::function::FunctionVersion>> {
|
||||
_options: ListingOptions,
|
||||
) -> Result<ListingPage<crate::function::FunctionVersion>> {
|
||||
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<Vec<String>> {
|
||||
async fn list_secrets(
|
||||
&self,
|
||||
_namespace_path: &[String],
|
||||
_options: ListingOptions,
|
||||
) -> Result<ListingPage<String>> {
|
||||
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<Job> {
|
||||
view_ops_not_supported()
|
||||
}
|
||||
/// The names of the views in one namespace.
|
||||
async fn list_views(&self, _namespace_path: &[String]) -> Result<Vec<String>> {
|
||||
/// Fetch one page of view names in one namespace.
|
||||
async fn list_views(
|
||||
&self,
|
||||
_namespace_path: &[String],
|
||||
_options: ListingOptions,
|
||||
) -> Result<ListingPage<String>> {
|
||||
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<crate::job::Job> {
|
||||
job_op_not_supported("open_job")
|
||||
}
|
||||
/// List server-side jobs across the database's tables.
|
||||
async fn list_jobs(&self) -> Result<Vec<JobInfo>> {
|
||||
/// Fetch one page of server-side jobs across the database's tables.
|
||||
async fn list_jobs(&self, _options: ListingOptions) -> Result<ListingPage<JobInfo>> {
|
||||
job_op_not_supported("list_jobs")
|
||||
}
|
||||
/// Request cancellation of a job by id. Returns true if the server
|
||||
|
||||
@@ -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;
|
||||
|
||||
@@ -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<String>,
|
||||
/// 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<u32>,
|
||||
}
|
||||
|
||||
impl ListingOptions {
|
||||
/// Resume at a saved continuation token.
|
||||
pub fn page_token(mut self, token: impl Into<String>) -> 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<T> {
|
||||
/// Items returned on this page, possibly empty even when more pages exist.
|
||||
pub items: Vec<T>,
|
||||
/// Token for the next page; `None` or an empty string ends the listing.
|
||||
pub page_token: Option<String>,
|
||||
}
|
||||
|
||||
type FetchPage<T> =
|
||||
Box<dyn Fn(ListingOptions) -> BoxFuture<'static, Result<ListingPage<T>>> + 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<T> {
|
||||
fetch: FetchPage<T>,
|
||||
items: std::vec::IntoIter<T>,
|
||||
options: ListingOptions,
|
||||
done: bool,
|
||||
seen_tokens: HashSet<String>,
|
||||
pending: Option<BoxFuture<'static, Result<ListingPage<T>>>>,
|
||||
}
|
||||
|
||||
impl<T> Listing<T> {
|
||||
pub(crate) fn new<F, Fut>(options: ListingOptions, fetch: F) -> Self
|
||||
where
|
||||
F: Fn(ListingOptions) -> Fut + Send + Sync + 'static,
|
||||
Fut: Future<Output = Result<ListingPage<T>>> + 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<T: Unpin> Stream for Listing<T> {
|
||||
type Item = Result<T>;
|
||||
|
||||
fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
|
||||
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::<String>::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::<String>::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::<Vec<_>>().await.unwrap(),
|
||||
(0..=110).collect::<Vec<_>>()
|
||||
);
|
||||
}
|
||||
}
|
||||
@@ -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<Vec<String>> {
|
||||
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<String> {
|
||||
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.
|
||||
|
||||
+294
-201
@@ -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<String>,
|
||||
}
|
||||
|
||||
/// 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<S: HttpSend> Database for RemoteDatabase<S> {
|
||||
fn uri(&self) -> &str {
|
||||
@@ -919,45 +916,28 @@ impl<S: HttpSend> Database for RemoteDatabase<S> {
|
||||
Ok(job)
|
||||
}
|
||||
|
||||
async fn list_materialized_views(&self, namespace_path: &[String]) -> Result<Vec<String>> {
|
||||
#[derive(serde::Deserialize)]
|
||||
struct ListMaterializedViewsResponse {
|
||||
#[serde(default)]
|
||||
views: Vec<String>,
|
||||
#[serde(default)]
|
||||
page_token: Option<String>,
|
||||
}
|
||||
|
||||
async fn list_materialized_views(
|
||||
&self,
|
||||
namespace_path: &[String],
|
||||
options: ListingOptions,
|
||||
) -> Result<ListingPage<String>> {
|
||||
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<String> = 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<S: HttpSend> Database for RemoteDatabase<S> {
|
||||
response.json().await.err_to_http(request_id)
|
||||
}
|
||||
|
||||
async fn list_functions(&self, namespace_path: &[String]) -> Result<Vec<FunctionVersion>> {
|
||||
async fn list_functions(
|
||||
&self,
|
||||
namespace_path: &[String],
|
||||
options: ListingOptions,
|
||||
) -> Result<ListingPage<crate::function::FunctionVersion>> {
|
||||
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<String> = 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<S: HttpSend> Database for RemoteDatabase<S> {
|
||||
.await
|
||||
}
|
||||
|
||||
async fn list_secrets(&self, namespace_path: &[String]) -> Result<Vec<String>> {
|
||||
async fn list_secrets(
|
||||
&self,
|
||||
namespace_path: &[String],
|
||||
options: ListingOptions,
|
||||
) -> Result<ListingPage<String>> {
|
||||
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<String> = 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<S: HttpSend> Database for RemoteDatabase<S> {
|
||||
}
|
||||
}
|
||||
|
||||
async fn list_views(&self, namespace_path: &[String]) -> Result<Vec<String>> {
|
||||
async fn list_views(
|
||||
&self,
|
||||
namespace_path: &[String],
|
||||
options: ListingOptions,
|
||||
) -> Result<ListingPage<String>> {
|
||||
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<String> = 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<Job> {
|
||||
@@ -1295,47 +1250,33 @@ impl<S: HttpSend> Database for RemoteDatabase<S> {
|
||||
}
|
||||
}
|
||||
|
||||
async fn list_jobs(&self) -> Result<Vec<JobInfo>> {
|
||||
let mut out = Vec::new();
|
||||
let mut page_token: Option<String> = 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<ListingPage<JobInfo>> {
|
||||
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<bool> {
|
||||
@@ -1905,6 +1846,8 @@ impl From<StorageOptions> 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::<Vec<_>>()
|
||||
.await
|
||||
.unwrap(),
|
||||
["adults"]
|
||||
);
|
||||
}
|
||||
@@ -3534,6 +3487,110 @@ mod tests {
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_resource_listings_are_lazy_and_forward_pagination() {
|
||||
async fn check<T: Unpin + std::fmt::Debug>(mut items: Listing<T>, count: Arc<AtomicUsize>) {
|
||||
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::<HashMap<_, _>>();
|
||||
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::<serde_json::Value>(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::<Vec<_>>()
|
||||
.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::<Vec<_>>()
|
||||
.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::<Vec<_>>().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::<Vec<_>>()
|
||||
.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::<Vec<_>>()
|
||||
.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::<Vec<_>>()
|
||||
.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::<Vec<_>>()
|
||||
.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::<Vec<_>>()
|
||||
.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::<Vec<_>>()
|
||||
.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::<Vec<_>>()
|
||||
.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::<Vec<_>>()
|
||||
.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::<Vec<_>>().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::<Vec<_>>()
|
||||
.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::<Vec<_>>()
|
||||
.await
|
||||
.is_err()
|
||||
);
|
||||
assert!(conn.drop_function("embed", "1", &namespace).await.is_err());
|
||||
}
|
||||
|
||||
|
||||
Reference in New Issue
Block a user