refactor!: implement lazy iterators for several APIs (#4459)

Implement lazy iterators for:
- secrets
- views
- jobs
- functions
- materialized views
This commit is contained in:
Colin Patrick McCabe
2026-10-07 10:24:20 -07:00
committed by GitHub
parent 069b2f7d31
commit de005f1d61
32 changed files with 2185 additions and 454 deletions
+24 -8
View File
@@ -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`&lt;[`JobInfo`](../interfaces/JobInfo.md)[]&gt;
[`Listing`](Listing.md)&lt;[`JobInfo`](../interfaces/JobInfo.md)&gt;
***
### 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`&lt;`string`[]&gt;
[`Listing`](Listing.md)&lt;`string`&gt;
***
@@ -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`&lt;`string`[]&gt;
[`Listing`](Listing.md)&lt;`string`&gt;
***
+91
View File
@@ -0,0 +1,91 @@
[**@lancedb/lancedb**](../README.md) • **Docs**
***
[@lancedb/lancedb](../globals.md) / Listing
# Class: Listing&lt;T&gt;
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`&lt;`T`&gt;
## Methods
### \[asyncIterator\]()
```ts
asyncIterator: AsyncIterableIterator<T>
```
#### Returns
`AsyncIterableIterator`&lt;`T`&gt;
#### Implementation of
`AsyncIterableIterator.[asyncIterator]`
***
### next()
```ts
next(): Promise<IteratorResult<T, any>>
```
Fetch the next item, requesting another page only when needed.
#### Returns
`Promise`&lt;`IteratorResult`&lt;`T`, `any`&gt;&gt;
#### 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`
+2
View File
@@ -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)
+29
View File
@@ -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.
+10
View File
@@ -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
+10 -4
View File
@@ -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;
}
+149 -3
View File
@@ -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);
},
);
});
}
});
+39 -11
View File
@@ -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> {
+2
View File
@@ -649,3 +649,5 @@ export {
DatabaseNames,
connectCatalog,
} from "./catalog";
export { Listing, ListingOptions } from "./listing";
+73
View File
@@ -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
View File
@@ -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
+1
View File
@@ -18,6 +18,7 @@ mod header;
mod index;
mod iterator;
mod job;
mod listing;
pub mod merge;
pub mod otel;
pub mod permutation;
+83
View File
@@ -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()))
}
}
+45 -11
View File
@@ -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
View File
@@ -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.
+85
View File
@@ -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()
+31 -9
View File
@@ -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
+88 -12
View File
@@ -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"] == []
+194
View File
@@ -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 == []
+26 -10
View File
@@ -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) == []
+25
View File
@@ -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
):
+2 -2
View File
@@ -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
View File
@@ -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>> {
+5
View File
@@ -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>()?;
+131
View File
@@ -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
View File
@@ -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
+55 -49
View File
@@ -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
+1
View File
@@ -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;
+261
View File
@@ -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<_>>()
);
}
}
+21 -2
View File
@@ -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
View File
@@ -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());
}