mirror of
https://github.com/lancedb/lancedb.git
synced 2026-10-23 20:22:37 +00:00
feat!: list_databases returns an iterator (#4442)
CatalogConnection.list_databases should return an iterator so that users don't have to manually re-implement pagination. For those who do want more control over pagination, add num_page_results and page_token to the Iterator object, and also allow the iterator object to be created at a specific page token.
This commit is contained in:
@@ -89,19 +89,30 @@ Drop an empty database. The server rejects nonempty databases.
|
|||||||
### listDatabases()
|
### listDatabases()
|
||||||
|
|
||||||
```ts
|
```ts
|
||||||
listDatabases(options): Promise<ListDatabasesResponse>
|
listDatabases(options): DatabaseNames
|
||||||
```
|
```
|
||||||
|
|
||||||
List one page of databases; pass pageToken from a response for the next page.
|
Iterate lazily over all database names using `for await...of`.
|
||||||
|
pageToken resumes from a saved token; omitted starts at the beginning.
|
||||||
|
pageLimit limits each REST response, not the total; omitted uses the server default.
|
||||||
|
Request errors are raised during iteration and terminate the iterator.
|
||||||
|
|
||||||
#### Parameters
|
#### Parameters
|
||||||
|
|
||||||
* **options** = `{}`
|
* **options** = `{}`
|
||||||
|
|
||||||
* **options.limit?**: `number`
|
* **options.pageLimit?**: `number`
|
||||||
|
|
||||||
* **options.pageToken?**: `string`
|
* **options.pageToken?**: `string`
|
||||||
|
|
||||||
#### Returns
|
#### Returns
|
||||||
|
|
||||||
`Promise`<[`ListDatabasesResponse`](../interfaces/ListDatabasesResponse.md)>
|
[`DatabaseNames`](DatabaseNames.md)
|
||||||
|
|
||||||
|
#### Example
|
||||||
|
|
||||||
|
```ts
|
||||||
|
for await (const name of catalog.listDatabases({ pageLimit: 20 })) {
|
||||||
|
console.log(name);
|
||||||
|
}
|
||||||
|
```
|
||||||
|
|||||||
@@ -0,0 +1,79 @@
|
|||||||
|
[**@lancedb/lancedb**](../README.md) • **Docs**
|
||||||
|
|
||||||
|
***
|
||||||
|
|
||||||
|
[@lancedb/lancedb](../globals.md) / DatabaseNames
|
||||||
|
|
||||||
|
# Class: DatabaseNames
|
||||||
|
|
||||||
|
A lazy async iterator of database names with pagination state.
|
||||||
|
Returned by [Catalog.listDatabases](Catalog.md#listdatabases).
|
||||||
|
|
||||||
|
## Implements
|
||||||
|
|
||||||
|
- `AsyncIterableIterator`<`string`>
|
||||||
|
|
||||||
|
## Methods
|
||||||
|
|
||||||
|
### \[asyncIterator\]()
|
||||||
|
|
||||||
|
```ts
|
||||||
|
asyncIterator: AsyncIterableIterator<string>
|
||||||
|
```
|
||||||
|
|
||||||
|
#### Returns
|
||||||
|
|
||||||
|
`AsyncIterableIterator`<`string`>
|
||||||
|
|
||||||
|
#### Implementation of
|
||||||
|
|
||||||
|
`AsyncIterableIterator.[asyncIterator]`
|
||||||
|
|
||||||
|
***
|
||||||
|
|
||||||
|
### next()
|
||||||
|
|
||||||
|
```ts
|
||||||
|
next(): Promise<IteratorResult<string, any>>
|
||||||
|
```
|
||||||
|
|
||||||
|
Fetch the next name, requesting another page only when needed.
|
||||||
|
|
||||||
|
#### Returns
|
||||||
|
|
||||||
|
`Promise`<`IteratorResult`<`string`, `any`>>
|
||||||
|
|
||||||
|
#### Implementation of
|
||||||
|
|
||||||
|
`AsyncIterableIterator.next`
|
||||||
|
|
||||||
|
***
|
||||||
|
|
||||||
|
### numPageResults()
|
||||||
|
|
||||||
|
```ts
|
||||||
|
numPageResults(): number
|
||||||
|
```
|
||||||
|
|
||||||
|
Number of names available without another REST request.
|
||||||
|
|
||||||
|
#### Returns
|
||||||
|
|
||||||
|
`number`
|
||||||
|
|
||||||
|
***
|
||||||
|
|
||||||
|
### pageToken()
|
||||||
|
|
||||||
|
```ts
|
||||||
|
pageToken(): undefined | string
|
||||||
|
```
|
||||||
|
|
||||||
|
Token for the next REST request. Initially this is the supplied starting token.
|
||||||
|
Returns undefined after the final page is fetched, even if names remain cached.
|
||||||
|
Drain the cache before saving a token to avoid skipping names when resuming.
|
||||||
|
A failed request terminates iteration and retains its token for a new iterator.
|
||||||
|
|
||||||
|
#### Returns
|
||||||
|
|
||||||
|
`undefined` \| `string`
|
||||||
@@ -28,5 +28,7 @@ headers; opened database connections inherit authentication and client options.
|
|||||||
```ts
|
```ts
|
||||||
const catalog = await connectCatalog("https://my-server.example", { apiKey: "secret" });
|
const catalog = await connectCatalog("https://my-server.example", { apiKey: "secret" });
|
||||||
const db = await catalog.createDatabase("analytics", { existOk: true });
|
const db = await catalog.createDatabase("analytics", { existOk: true });
|
||||||
const page = await catalog.listDatabases({ limit: 20 });
|
for await (const name of catalog.listDatabases({ pageLimit: 20 })) {
|
||||||
|
console.log(name);
|
||||||
|
}
|
||||||
```
|
```
|
||||||
|
|||||||
@@ -27,6 +27,7 @@
|
|||||||
- [Branches](classes/Branches.md)
|
- [Branches](classes/Branches.md)
|
||||||
- [Catalog](classes/Catalog.md)
|
- [Catalog](classes/Catalog.md)
|
||||||
- [Connection](classes/Connection.md)
|
- [Connection](classes/Connection.md)
|
||||||
|
- [DatabaseNames](classes/DatabaseNames.md)
|
||||||
- [HeaderProvider](classes/HeaderProvider.md)
|
- [HeaderProvider](classes/HeaderProvider.md)
|
||||||
- [Index](classes/Index.md)
|
- [Index](classes/Index.md)
|
||||||
- [Job](classes/Job.md)
|
- [Job](classes/Job.md)
|
||||||
@@ -108,7 +109,6 @@
|
|||||||
- [JobEventsOptions](interfaces/JobEventsOptions.md)
|
- [JobEventsOptions](interfaces/JobEventsOptions.md)
|
||||||
- [JobFailureInfo](interfaces/JobFailureInfo.md)
|
- [JobFailureInfo](interfaces/JobFailureInfo.md)
|
||||||
- [JobInfo](interfaces/JobInfo.md)
|
- [JobInfo](interfaces/JobInfo.md)
|
||||||
- [ListDatabasesResponse](interfaces/ListDatabasesResponse.md)
|
|
||||||
- [ListNamespacesOptions](interfaces/ListNamespacesOptions.md)
|
- [ListNamespacesOptions](interfaces/ListNamespacesOptions.md)
|
||||||
- [ListNamespacesResponse](interfaces/ListNamespacesResponse.md)
|
- [ListNamespacesResponse](interfaces/ListNamespacesResponse.md)
|
||||||
- [ListTablesOptions](interfaces/ListTablesOptions.md)
|
- [ListTablesOptions](interfaces/ListTablesOptions.md)
|
||||||
|
|||||||
@@ -1,23 +0,0 @@
|
|||||||
[**@lancedb/lancedb**](../README.md) • **Docs**
|
|
||||||
|
|
||||||
***
|
|
||||||
|
|
||||||
[@lancedb/lancedb](../globals.md) / ListDatabasesResponse
|
|
||||||
|
|
||||||
# Interface: ListDatabasesResponse
|
|
||||||
|
|
||||||
## Properties
|
|
||||||
|
|
||||||
### databases
|
|
||||||
|
|
||||||
```ts
|
|
||||||
databases: string[];
|
|
||||||
```
|
|
||||||
|
|
||||||
***
|
|
||||||
|
|
||||||
### pageToken?
|
|
||||||
|
|
||||||
```ts
|
|
||||||
optional pageToken: string;
|
|
||||||
```
|
|
||||||
@@ -37,7 +37,8 @@ are ordinary connections. Dropping a database requires it to be empty.
|
|||||||
|
|
||||||
::: lancedb.catalog.Catalog
|
::: lancedb.catalog.Catalog
|
||||||
|
|
||||||
::: lancedb.catalog.ListDatabasesResponse
|
::: lancedb.catalog.DatabaseNames
|
||||||
|
|
||||||
|
|
||||||
## Remote SQL
|
## Remote SQL
|
||||||
|
|
||||||
@@ -440,6 +441,8 @@ latency on your workload before opting in.
|
|||||||
|
|
||||||
::: lancedb.catalog.AsyncCatalog
|
::: lancedb.catalog.AsyncCatalog
|
||||||
|
|
||||||
|
::: lancedb.catalog.AsyncDatabaseNames
|
||||||
|
|
||||||
Connections represent a connection to a LanceDb database and
|
Connections represent a connection to a LanceDb database and
|
||||||
can be used to create, list, or open tables.
|
can be used to create, list, or open tables.
|
||||||
|
|
||||||
|
|||||||
@@ -80,9 +80,12 @@ describe("remote catalog", () => {
|
|||||||
const second = await catalog.connectDatabase("other");
|
const second = await catalog.connectDatabase("other");
|
||||||
expect(await second.tableNames()).toEqual([]);
|
expect(await second.tableNames()).toEqual([]);
|
||||||
expect(await first.tableNames()).toEqual([]);
|
expect(await first.tableNames()).toEqual([]);
|
||||||
expect(
|
const names = catalog.listDatabases({ pageLimit: 1, pageToken: "a/b" });
|
||||||
await catalog.listDatabases({ limit: 1, pageToken: "a/b" }),
|
expect(await names.next()).toEqual({
|
||||||
).toEqual({ databases: ["team/search"], pageToken: "next" });
|
done: false,
|
||||||
|
value: "team/search",
|
||||||
|
});
|
||||||
|
expect(names.pageToken()).toBe("next");
|
||||||
await catalog.dropDatabase("team/search", { ignoreMissing: true });
|
await catalog.dropDatabase("team/search", { ignoreMissing: true });
|
||||||
expect(requests[0].url).toBe("/v1/namespace/team%2Fsearch/create");
|
expect(requests[0].url).toBe("/v1/namespace/team%2Fsearch/create");
|
||||||
expect(requests[0].body).toEqual({ mode: "ExistOk" });
|
expect(requests[0].body).toEqual({ mode: "ExistOk" });
|
||||||
@@ -129,11 +132,139 @@ describe("remote catalog", () => {
|
|||||||
);
|
);
|
||||||
});
|
});
|
||||||
|
|
||||||
|
it("lazily traverses pages, including empty pages, and exposes cached state", async () => {
|
||||||
|
await withCatalog(
|
||||||
|
[
|
||||||
|
// biome-ignore lint/style/useNamingConvention: server wire format
|
||||||
|
[200, { namespaces: ["a", "b"], page_token: "empty" }],
|
||||||
|
// biome-ignore lint/style/useNamingConvention: server wire format
|
||||||
|
[200, { namespaces: [], page_token: "last" }],
|
||||||
|
// biome-ignore lint/style/useNamingConvention: server wire format
|
||||||
|
[200, { namespaces: ["c", "d"], page_token: "" }],
|
||||||
|
],
|
||||||
|
async (catalog, requests) => {
|
||||||
|
const names = catalog.listDatabases({ pageLimit: 2 });
|
||||||
|
expect(names[Symbol.asyncIterator]()).toBe(names);
|
||||||
|
expect(requests).toHaveLength(0);
|
||||||
|
expect(names.numPageResults()).toBe(0);
|
||||||
|
expect(names.pageToken()).toBeUndefined();
|
||||||
|
expect(await names.next()).toEqual({ done: false, value: "a" });
|
||||||
|
expect(names.numPageResults()).toBe(1);
|
||||||
|
expect(names.pageToken()).toBe("empty");
|
||||||
|
expect(await names.next()).toEqual({ done: false, value: "b" });
|
||||||
|
expect(requests).toHaveLength(1);
|
||||||
|
expect(await names.next()).toEqual({ done: false, value: "c" });
|
||||||
|
expect(requests).toHaveLength(3);
|
||||||
|
expect(names.pageToken()).toBeUndefined();
|
||||||
|
expect(names.numPageResults()).toBe(1);
|
||||||
|
const remaining = [];
|
||||||
|
for await (const name of names) remaining.push(name);
|
||||||
|
expect(remaining).toEqual(["d"]);
|
||||||
|
expect(names.numPageResults()).toBe(0);
|
||||||
|
expect(await names.next()).toEqual({ done: true, value: undefined });
|
||||||
|
expect(requests.map((request) => request.url)).toEqual([
|
||||||
|
"/v1/namespace/$/list?limit=2",
|
||||||
|
"/v1/namespace/$/list?limit=2&page_token=empty",
|
||||||
|
"/v1/namespace/$/list?limit=2&page_token=last",
|
||||||
|
]);
|
||||||
|
},
|
||||||
|
);
|
||||||
|
});
|
||||||
|
|
||||||
|
it("resumes from a saved token after draining the cache", async () => {
|
||||||
|
await withCatalog(
|
||||||
|
[
|
||||||
|
// biome-ignore lint/style/useNamingConvention: server wire format
|
||||||
|
[200, { namespaces: ["a", "b"], page_token: "next/token" }],
|
||||||
|
[200, { namespaces: ["c"] }],
|
||||||
|
],
|
||||||
|
async (catalog, requests) => {
|
||||||
|
const names = catalog.listDatabases({
|
||||||
|
pageToken: "start/token",
|
||||||
|
pageLimit: 2,
|
||||||
|
});
|
||||||
|
expect(names.pageToken()).toBe("start/token");
|
||||||
|
expect(await names.next()).toEqual({ done: false, value: "a" });
|
||||||
|
while (names.numPageResults() > 0) await names.next();
|
||||||
|
const resumed = catalog.listDatabases({
|
||||||
|
pageToken: names.pageToken(),
|
||||||
|
pageLimit: 2,
|
||||||
|
});
|
||||||
|
const results = [];
|
||||||
|
for await (const name of resumed) results.push(name);
|
||||||
|
expect(results).toEqual(["c"]);
|
||||||
|
expect(requests.map((request) => request.url)).toEqual([
|
||||||
|
"/v1/namespace/$/list?limit=2&page_token=start%2Ftoken",
|
||||||
|
"/v1/namespace/$/list?limit=2&page_token=next%2Ftoken",
|
||||||
|
]);
|
||||||
|
},
|
||||||
|
);
|
||||||
|
});
|
||||||
|
|
||||||
|
it("terminates on errors and retains the failed request token", async () => {
|
||||||
|
await withCatalog(
|
||||||
|
[
|
||||||
|
// biome-ignore lint/style/useNamingConvention: server wire format
|
||||||
|
[200, { namespaces: ["a"], page_token: "retry" }],
|
||||||
|
[400, {}],
|
||||||
|
[200, { namespaces: ["b"] }],
|
||||||
|
],
|
||||||
|
async (catalog, requests) => {
|
||||||
|
const names = catalog.listDatabases();
|
||||||
|
await names.next();
|
||||||
|
await expect(names.next()).rejects.toThrow();
|
||||||
|
expect(names.pageToken()).toBe("retry");
|
||||||
|
expect(names.numPageResults()).toBe(0);
|
||||||
|
expect(await names.next()).toEqual({ done: true, value: undefined });
|
||||||
|
expect(requests).toHaveLength(2);
|
||||||
|
const resumed = catalog.listDatabases({ pageToken: names.pageToken() });
|
||||||
|
expect(await resumed.next()).toEqual({ done: false, value: "b" });
|
||||||
|
expect(await resumed.next()).toEqual({ done: true, value: undefined });
|
||||||
|
},
|
||||||
|
);
|
||||||
|
});
|
||||||
|
|
||||||
|
it("handles an empty listing and concurrent advances", async () => {
|
||||||
|
await withCatalog(
|
||||||
|
[
|
||||||
|
[200, { namespaces: [] }],
|
||||||
|
[200, { namespaces: ["a", "b"] }],
|
||||||
|
],
|
||||||
|
async (catalog, requests) => {
|
||||||
|
const empty = catalog.listDatabases();
|
||||||
|
expect(await empty.next()).toEqual({ done: true, value: undefined });
|
||||||
|
expect(await empty.next()).toEqual({ done: true, value: undefined });
|
||||||
|
const names = catalog.listDatabases();
|
||||||
|
const pending = [names.next(), names.next(), names.next()];
|
||||||
|
expect(names.numPageResults()).toBe(0);
|
||||||
|
expect(names.pageToken()).toBeUndefined();
|
||||||
|
const results = await Promise.all(pending);
|
||||||
|
expect(
|
||||||
|
results
|
||||||
|
.filter((result) => !result.done)
|
||||||
|
.map((result) => result.value)
|
||||||
|
.sort(),
|
||||||
|
).toEqual(["a", "b"]);
|
||||||
|
expect(results.filter((result) => result.done)).toHaveLength(1);
|
||||||
|
expect(requests).toHaveLength(2);
|
||||||
|
},
|
||||||
|
);
|
||||||
|
});
|
||||||
|
|
||||||
it("validates endpoints and pagination", async () => {
|
it("validates endpoints and pagination", async () => {
|
||||||
await expect(connectCatalog("/tmp/catalog")).rejects.toThrow();
|
await expect(connectCatalog("/tmp/catalog")).rejects.toThrow();
|
||||||
const catalog = await connectCatalog("http://127.0.0.1:1");
|
const catalog = await connectCatalog("http://127.0.0.1:1");
|
||||||
for (const limit of [0, -1, 1.5, 2147483648]) {
|
for (const limit of [
|
||||||
await expect(catalog.listDatabases({ limit })).rejects.toThrow("limit");
|
0,
|
||||||
|
-1,
|
||||||
|
1.5,
|
||||||
|
2147483648,
|
||||||
|
Number.NaN,
|
||||||
|
Number.POSITIVE_INFINITY,
|
||||||
|
]) {
|
||||||
|
expect(() => catalog.listDatabases({ pageLimit: limit })).toThrow(
|
||||||
|
"limit",
|
||||||
|
);
|
||||||
}
|
}
|
||||||
await expect(catalog.connectDatabase("a$b")).rejects.toThrow(
|
await expect(catalog.connectDatabase("a$b")).rejects.toThrow(
|
||||||
"Invalid database name",
|
"Invalid database name",
|
||||||
|
|||||||
+62
-12
@@ -5,9 +5,9 @@ import { Connection, LocalConnection } from "./connection";
|
|||||||
import { HeaderProvider } from "./header";
|
import { HeaderProvider } from "./header";
|
||||||
import {
|
import {
|
||||||
JsHeaderProvider,
|
JsHeaderProvider,
|
||||||
ListDatabasesResponse,
|
|
||||||
Catalog as NativeCatalog,
|
Catalog as NativeCatalog,
|
||||||
CatalogOptions as NativeCatalogOptions,
|
CatalogOptions as NativeCatalogOptions,
|
||||||
|
DatabaseNames as NativeDatabaseNames,
|
||||||
} from "./native.js";
|
} from "./native.js";
|
||||||
import { OAuthConfig } from "./oauth";
|
import { OAuthConfig } from "./oauth";
|
||||||
|
|
||||||
@@ -21,7 +21,41 @@ export interface CatalogOptions
|
|||||||
| (() => Record<string, string> | Promise<Record<string, string>>);
|
| (() => Record<string, string> | Promise<Record<string, string>>);
|
||||||
}
|
}
|
||||||
|
|
||||||
export type { ListDatabasesResponse } from "./native.js";
|
/**
|
||||||
|
* A lazy async iterator of database names with pagination state.
|
||||||
|
* Returned by {@link Catalog.listDatabases}.
|
||||||
|
*/
|
||||||
|
export class DatabaseNames implements AsyncIterableIterator<string> {
|
||||||
|
/** @hidden */
|
||||||
|
constructor(private readonly inner: NativeDatabaseNames) {}
|
||||||
|
|
||||||
|
[Symbol.asyncIterator](): AsyncIterableIterator<string> {
|
||||||
|
return this;
|
||||||
|
}
|
||||||
|
|
||||||
|
/** Fetch the next name, requesting another page only when needed. */
|
||||||
|
async next(): Promise<IteratorResult<string>> {
|
||||||
|
const value = await this.inner.next();
|
||||||
|
return value == null
|
||||||
|
? { done: true, value: undefined }
|
||||||
|
: { done: false, value };
|
||||||
|
}
|
||||||
|
|
||||||
|
/** Number of names available without another REST request. */
|
||||||
|
numPageResults(): number {
|
||||||
|
return this.inner.numPageResults();
|
||||||
|
}
|
||||||
|
|
||||||
|
/**
|
||||||
|
* Token for the next REST request. Initially this is the supplied starting token.
|
||||||
|
* Returns undefined after the final page is fetched, even if names remain cached.
|
||||||
|
* Drain the cache before saving a token to avoid skipping names when resuming.
|
||||||
|
* A failed request terminates iteration and retains its token for a new iterator.
|
||||||
|
*/
|
||||||
|
pageToken(): string | undefined {
|
||||||
|
return this.inner.pageToken() ?? undefined;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
/** A remote catalog manages databases through the server's root namespace. */
|
/** A remote catalog manages databases through the server's root namespace. */
|
||||||
export class Catalog {
|
export class Catalog {
|
||||||
@@ -56,21 +90,35 @@ export class Catalog {
|
|||||||
await this.inner.dropDatabase(name, options.ignoreMissing);
|
await this.inner.dropDatabase(name, options.ignoreMissing);
|
||||||
}
|
}
|
||||||
|
|
||||||
/** List one page of databases; pass pageToken from a response for the next page. */
|
/**
|
||||||
async listDatabases(
|
* Iterate lazily over all database names using `for await...of`.
|
||||||
options: { limit?: number; pageToken?: string } = {},
|
* pageToken resumes from a saved token; omitted starts at the beginning.
|
||||||
): Promise<ListDatabasesResponse> {
|
* pageLimit limits each REST response, not the total; omitted uses the server default.
|
||||||
|
* Request errors are raised during iteration and terminate the iterator.
|
||||||
|
*
|
||||||
|
* @example
|
||||||
|
* ```ts
|
||||||
|
* for await (const name of catalog.listDatabases({ pageLimit: 20 })) {
|
||||||
|
* console.log(name);
|
||||||
|
* }
|
||||||
|
* ```
|
||||||
|
*/
|
||||||
|
listDatabases(
|
||||||
|
options: { pageLimit?: number; pageToken?: string } = {},
|
||||||
|
): DatabaseNames {
|
||||||
if (
|
if (
|
||||||
options.limit !== undefined &&
|
options.pageLimit !== undefined &&
|
||||||
(!Number.isInteger(options.limit) ||
|
(!Number.isInteger(options.pageLimit) ||
|
||||||
options.limit <= 0 ||
|
options.pageLimit <= 0 ||
|
||||||
options.limit > 2147483647)
|
options.pageLimit > 2147483647)
|
||||||
) {
|
) {
|
||||||
throw new Error(
|
throw new Error(
|
||||||
"Database list limit must be an integer between 1 and 2147483647",
|
"Database list limit must be an integer between 1 and 2147483647",
|
||||||
);
|
);
|
||||||
}
|
}
|
||||||
return this.inner.listDatabases(options.limit, options.pageToken);
|
return new DatabaseNames(
|
||||||
|
this.inner.listDatabases(options.pageToken, options.pageLimit),
|
||||||
|
);
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -82,7 +130,9 @@ export class Catalog {
|
|||||||
* ```ts
|
* ```ts
|
||||||
* const catalog = await connectCatalog("https://my-server.example", { apiKey: "secret" });
|
* const catalog = await connectCatalog("https://my-server.example", { apiKey: "secret" });
|
||||||
* const db = await catalog.createDatabase("analytics", { existOk: true });
|
* const db = await catalog.createDatabase("analytics", { existOk: true });
|
||||||
* const page = await catalog.listDatabases({ limit: 20 });
|
* for await (const name of catalog.listDatabases({ pageLimit: 20 })) {
|
||||||
|
* console.log(name);
|
||||||
|
* }
|
||||||
* ```
|
* ```
|
||||||
*/
|
*/
|
||||||
export async function connectCatalog(
|
export async function connectCatalog(
|
||||||
|
|||||||
@@ -646,6 +646,6 @@ export async function connectNamespace(
|
|||||||
export {
|
export {
|
||||||
Catalog,
|
Catalog,
|
||||||
CatalogOptions,
|
CatalogOptions,
|
||||||
ListDatabasesResponse,
|
DatabaseNames,
|
||||||
connectCatalog,
|
connectCatalog,
|
||||||
} from "./catalog";
|
} from "./catalog";
|
||||||
|
|||||||
+42
-19
@@ -1,12 +1,12 @@
|
|||||||
// SPDX-License-Identifier: Apache-2.0
|
// SPDX-License-Identifier: Apache-2.0
|
||||||
// SPDX-FileCopyrightText: Copyright The LanceDB Authors
|
// SPDX-FileCopyrightText: Copyright The LanceDB Authors
|
||||||
|
|
||||||
use std::sync::Arc;
|
use std::sync::{Arc, Mutex};
|
||||||
use std::time::Duration;
|
use std::time::Duration;
|
||||||
|
|
||||||
use lancedb::catalog::{
|
use futures::{StreamExt, future::poll_fn};
|
||||||
CatalogConnection, CreateDatabaseRequest, DropDatabaseRequest, ListDatabasesRequest,
|
|
||||||
};
|
use lancedb::catalog::{CatalogConnection, CreateDatabaseRequest, DropDatabaseRequest};
|
||||||
use napi::bindgen_prelude::*;
|
use napi::bindgen_prelude::*;
|
||||||
use napi_derive::napi;
|
use napi_derive::napi;
|
||||||
|
|
||||||
@@ -25,10 +25,37 @@ pub struct CatalogOptions {
|
|||||||
pub oauth_config: Option<OAuthConfig>,
|
pub oauth_config: Option<OAuthConfig>,
|
||||||
}
|
}
|
||||||
|
|
||||||
#[napi(object)]
|
/// A lazy iterator over database names with inspectable pagination state.
|
||||||
pub struct ListDatabasesResponse {
|
#[napi]
|
||||||
pub databases: Vec<String>,
|
pub struct DatabaseNames {
|
||||||
pub page_token: Option<String>,
|
inner: Mutex<lancedb::catalog::DatabaseNames>,
|
||||||
|
next_lock: futures::lock::Mutex<()>,
|
||||||
|
}
|
||||||
|
|
||||||
|
#[napi]
|
||||||
|
impl DatabaseNames {
|
||||||
|
/// Number of names cached without another REST request.
|
||||||
|
#[napi]
|
||||||
|
pub fn num_page_results(&self) -> u32 {
|
||||||
|
self.inner.lock().unwrap().num_page_results() as u32
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Continuation token for the next REST request.
|
||||||
|
#[napi]
|
||||||
|
pub fn page_token(&self) -> Option<String> {
|
||||||
|
self.inner.lock().unwrap().page_token().map(str::to_owned)
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Fetch the next name, returning None when exhausted.
|
||||||
|
#[napi]
|
||||||
|
pub async fn next(&self) -> Result<Option<String>> {
|
||||||
|
// Serialize advances but allow synchronous inspection between polls.
|
||||||
|
let _guard = self.next_lock.lock().await;
|
||||||
|
poll_fn(|cx| self.inner.lock().unwrap().poll_next_unpin(cx))
|
||||||
|
.await
|
||||||
|
.transpose()
|
||||||
|
.default_error()
|
||||||
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
#[napi]
|
#[napi]
|
||||||
@@ -106,18 +133,14 @@ impl Catalog {
|
|||||||
.default_error()
|
.default_error()
|
||||||
}
|
}
|
||||||
#[napi]
|
#[napi]
|
||||||
pub async fn list_databases(
|
pub fn list_databases(
|
||||||
&self,
|
&self,
|
||||||
limit: Option<u32>,
|
|
||||||
page_token: Option<String>,
|
page_token: Option<String>,
|
||||||
) -> Result<ListDatabasesResponse> {
|
page_limit: Option<u32>,
|
||||||
let mut request = ListDatabasesRequest::default();
|
) -> DatabaseNames {
|
||||||
request.limit = limit;
|
DatabaseNames {
|
||||||
request.page_token = page_token;
|
inner: Mutex::new(self.inner.list_databases(page_token, page_limit)),
|
||||||
let response = self.inner.list_databases(request).await.default_error()?;
|
next_lock: futures::lock::Mutex::new(()),
|
||||||
Ok(ListDatabasesResponse {
|
}
|
||||||
databases: response.databases,
|
|
||||||
page_token: response.page_token,
|
|
||||||
})
|
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -62,7 +62,6 @@ from .namespace import (
|
|||||||
from .catalog import (
|
from .catalog import (
|
||||||
AsyncCatalog,
|
AsyncCatalog,
|
||||||
Catalog,
|
Catalog,
|
||||||
ListDatabasesResponse,
|
|
||||||
connect_catalog,
|
connect_catalog,
|
||||||
connect_catalog_async,
|
connect_catalog_async,
|
||||||
)
|
)
|
||||||
@@ -578,7 +577,6 @@ async def connect_async(
|
|||||||
__all__ = [
|
__all__ = [
|
||||||
"Catalog",
|
"Catalog",
|
||||||
"AsyncCatalog",
|
"AsyncCatalog",
|
||||||
"ListDatabasesResponse",
|
|
||||||
"connect_catalog",
|
"connect_catalog",
|
||||||
"connect_catalog_async",
|
"connect_catalog_async",
|
||||||
"AsyncMaterializedView",
|
"AsyncMaterializedView",
|
||||||
|
|||||||
@@ -946,9 +946,15 @@ class Catalog:
|
|||||||
async def drop_database(
|
async def drop_database(
|
||||||
self, name: str, *, ignore_missing: bool = False
|
self, name: str, *, ignore_missing: bool = False
|
||||||
) -> None: ...
|
) -> None: ...
|
||||||
async def list_databases(
|
def list_databases(
|
||||||
self, *, limit: Optional[int] = None, page_token: Optional[str] = None
|
self, *, page_token: Optional[str] = None, page_limit: Optional[int] = None
|
||||||
) -> tuple[list[str], Optional[str]]: ...
|
) -> DatabaseNames: ...
|
||||||
|
|
||||||
|
class DatabaseNames:
|
||||||
|
def num_page_results(self) -> int: ...
|
||||||
|
def page_token(self) -> Optional[str]: ...
|
||||||
|
def __aiter__(self) -> DatabaseNames: ...
|
||||||
|
async def __anext__(self) -> str: ...
|
||||||
|
|
||||||
async def connect_catalog(
|
async def connect_catalog(
|
||||||
endpoint: str,
|
endpoint: str,
|
||||||
|
|||||||
@@ -3,9 +3,8 @@
|
|||||||
|
|
||||||
"""Remote catalogs manage databases through a server's root namespace."""
|
"""Remote catalogs manage databases through a server's root namespace."""
|
||||||
|
|
||||||
from dataclasses import dataclass
|
|
||||||
from datetime import timedelta
|
from datetime import timedelta
|
||||||
from typing import Any, Optional, Union
|
from typing import Any, AsyncIterator, Iterator, Optional, Union
|
||||||
|
|
||||||
from . import _lancedb
|
from . import _lancedb
|
||||||
from .background_loop import LOOP
|
from .background_loop import LOOP
|
||||||
@@ -14,12 +13,66 @@ from .remote import ClientConfig, OAuthConfig
|
|||||||
from .remote.db import RemoteDBConnection
|
from .remote.db import RemoteDBConnection
|
||||||
|
|
||||||
|
|
||||||
@dataclass
|
class AsyncDatabaseNames(AsyncIterator[str]):
|
||||||
class ListDatabasesResponse:
|
"""An async iterator of database names with pagination state.
|
||||||
"""A page of database names and an optional continuation token."""
|
|
||||||
|
|
||||||
databases: list[str]
|
Returned by [list_databases][lancedb.catalog.AsyncCatalog.list_databases].
|
||||||
page_token: Optional[str] = None
|
"""
|
||||||
|
|
||||||
|
def __init__(self, inner: _lancedb.DatabaseNames):
|
||||||
|
self._inner = inner
|
||||||
|
|
||||||
|
def __aiter__(self) -> "AsyncDatabaseNames":
|
||||||
|
return self
|
||||||
|
|
||||||
|
async def __anext__(self) -> str:
|
||||||
|
return await self._inner.__anext__()
|
||||||
|
|
||||||
|
def num_page_results(self) -> int:
|
||||||
|
"""Return the number of names available without another REST request."""
|
||||||
|
return self._inner.num_page_results()
|
||||||
|
|
||||||
|
def page_token(self) -> Optional[str]:
|
||||||
|
"""Return the token for the next REST request.
|
||||||
|
|
||||||
|
Before the first request, this is the supplied starting token (``None``
|
||||||
|
starts at the beginning). After the final page is fetched, it is ``None``,
|
||||||
|
even if names remain cached. Drain the cache before saving a token to
|
||||||
|
avoid skipping those names when resuming. A failed request terminates
|
||||||
|
iteration and leaves its token available for resuming a new iterator.
|
||||||
|
"""
|
||||||
|
return self._inner.page_token()
|
||||||
|
|
||||||
|
|
||||||
|
class DatabaseNames(Iterator[str]):
|
||||||
|
"""A synchronous iterator of database names with pagination state.
|
||||||
|
|
||||||
|
Returned by [Catalog.list_databases][lancedb.catalog.Catalog.list_databases].
|
||||||
|
"""
|
||||||
|
|
||||||
|
def __init__(self, inner: AsyncDatabaseNames):
|
||||||
|
self._inner = inner
|
||||||
|
|
||||||
|
def __iter__(self) -> "DatabaseNames":
|
||||||
|
return self
|
||||||
|
|
||||||
|
def __next__(self) -> str:
|
||||||
|
try:
|
||||||
|
return LOOP.run(self._inner.__anext__())
|
||||||
|
except StopAsyncIteration:
|
||||||
|
raise StopIteration from None
|
||||||
|
|
||||||
|
def num_page_results(self) -> int:
|
||||||
|
"""Return the number of names available without another REST request."""
|
||||||
|
return self._inner.num_page_results()
|
||||||
|
|
||||||
|
def page_token(self) -> Optional[str]:
|
||||||
|
"""Return the token for the next REST request.
|
||||||
|
|
||||||
|
See [page_token][lancedb.catalog.AsyncDatabaseNames.page_token]
|
||||||
|
for initial, cached, final-page, and error behavior.
|
||||||
|
"""
|
||||||
|
return self._inner.page_token()
|
||||||
|
|
||||||
|
|
||||||
class AsyncCatalog:
|
class AsyncCatalog:
|
||||||
@@ -50,14 +103,25 @@ class AsyncCatalog:
|
|||||||
"""Connect to an existing database by its logical name."""
|
"""Connect to an existing database by its logical name."""
|
||||||
return AsyncConnection(await self._inner.connect_database(name))
|
return AsyncConnection(await self._inner.connect_database(name))
|
||||||
|
|
||||||
async def list_databases(
|
def list_databases(
|
||||||
self, *, limit: Optional[int] = None, page_token: Optional[str] = None
|
self, *, page_token: Optional[str] = None, page_limit: Optional[int] = None
|
||||||
) -> ListDatabasesResponse:
|
) -> AsyncDatabaseNames:
|
||||||
"""List a page of databases. Pass the returned token for the next page."""
|
"""Iterate lazily over all database names using ``async for``.
|
||||||
names, token = await self._inner.list_databases(
|
|
||||||
limit=limit, page_token=page_token
|
``page_token`` resumes from a saved token; ``None`` starts at the beginning.
|
||||||
|
``page_limit`` sets the maximum names per REST response (not the total);
|
||||||
|
``None`` uses the server default. Request errors are raised during iteration.
|
||||||
|
|
||||||
|
Examples
|
||||||
|
--------
|
||||||
|
```python
|
||||||
|
async for name in catalog.list_databases():
|
||||||
|
print(name)
|
||||||
|
```
|
||||||
|
"""
|
||||||
|
return AsyncDatabaseNames(
|
||||||
|
self._inner.list_databases(page_token=page_token, page_limit=page_limit)
|
||||||
)
|
)
|
||||||
return ListDatabasesResponse(names, token)
|
|
||||||
|
|
||||||
async def drop_database(self, name: str, *, ignore_missing: bool = False) -> None:
|
async def drop_database(self, name: str, *, ignore_missing: bool = False) -> None:
|
||||||
"""Drop an empty database. A nonempty database is an error."""
|
"""Drop an empty database. A nonempty database is an error."""
|
||||||
@@ -73,7 +137,8 @@ class Catalog:
|
|||||||
```python
|
```python
|
||||||
catalog = lancedb.connect_catalog("https://my-server.example", api_key="secret")
|
catalog = lancedb.connect_catalog("https://my-server.example", api_key="secret")
|
||||||
db = catalog.create_database("analytics", exist_ok=True)
|
db = catalog.create_database("analytics", exist_ok=True)
|
||||||
page = catalog.list_databases(limit=20)
|
for name in catalog.list_databases():
|
||||||
|
print(name)
|
||||||
```
|
```
|
||||||
"""
|
"""
|
||||||
|
|
||||||
@@ -118,10 +183,17 @@ class Catalog:
|
|||||||
)
|
)
|
||||||
|
|
||||||
def list_databases(
|
def list_databases(
|
||||||
self, *, limit: Optional[int] = None, page_token: Optional[str] = None
|
self, *, page_token: Optional[str] = None, page_limit: Optional[int] = None
|
||||||
) -> ListDatabasesResponse:
|
) -> DatabaseNames:
|
||||||
"""List a page of databases. Pass the returned token for the next page."""
|
"""Iterate lazily over database names, fetching more results as needed.
|
||||||
return LOOP.run(self._inner.list_databases(limit=limit, page_token=page_token))
|
|
||||||
|
``page_token`` resumes from a saved token; ``None`` starts at the beginning.
|
||||||
|
``page_limit`` sets the maximum names per REST response (not the total);
|
||||||
|
``None`` uses the server default. Request errors are raised during iteration.
|
||||||
|
"""
|
||||||
|
return DatabaseNames(
|
||||||
|
self._inner.list_databases(page_token=page_token, page_limit=page_limit)
|
||||||
|
)
|
||||||
|
|
||||||
def drop_database(self, name: str, *, ignore_missing: bool = False) -> None:
|
def drop_database(self, name: str, *, ignore_missing: bool = False) -> None:
|
||||||
"""Drop an empty database. A nonempty database is an error."""
|
"""Drop an empty database. A nonempty database is an error."""
|
||||||
|
|||||||
@@ -55,7 +55,7 @@ def test_catalog_sync_scope_and_serialization(catalog_server):
|
|||||||
(200, {}),
|
(200, {}),
|
||||||
(200, {"tables": []}),
|
(200, {"tables": []}),
|
||||||
(200, {"tables": []}),
|
(200, {"tables": []}),
|
||||||
(200, {"namespaces": ["team/search"], "page_token": "next"}),
|
(200, {"namespaces": ["team/search"]}),
|
||||||
(204, None),
|
(204, None),
|
||||||
]
|
]
|
||||||
)
|
)
|
||||||
@@ -85,13 +85,12 @@ def test_catalog_sync_scope_and_serialization(catalog_server):
|
|||||||
with pytest.raises(ValueError, match="sql_host_override must use"):
|
with pytest.raises(ValueError, match="sql_host_override must use"):
|
||||||
connection.execute_query_async("SELECT 1")
|
connection.execute_query_async("SELECT 1")
|
||||||
assert restored.table_names() == []
|
assert restored.table_names() == []
|
||||||
page = catalog.list_databases(limit=1, page_token="a/b")
|
assert list(catalog.list_databases()) == ["team/search"]
|
||||||
assert page == lancedb.ListDatabasesResponse(["team/search"], "next")
|
|
||||||
catalog.drop_database("team/search", ignore_missing=True)
|
catalog.drop_database("team/search", ignore_missing=True)
|
||||||
assert requests[0][0] == "/v1/namespace/team%2Fsearch/create"
|
assert requests[0][0] == "/v1/namespace/team%2Fsearch/create"
|
||||||
assert requests[0][2] == {"mode": "ExistOk"}
|
assert requests[0][2] == {"mode": "ExistOk"}
|
||||||
assert requests[1][0] == "/v1/namespace/team%2Fsearch/describe"
|
assert requests[1][0] == "/v1/namespace/team%2Fsearch/describe"
|
||||||
assert requests[4][0] == "/v1/namespace/$/list?limit=1&page_token=a%2Fb"
|
assert requests[4][0] == "/v1/namespace/$/list"
|
||||||
assert requests[5][2] == {"mode": "Skip", "behavior": "Restrict"}
|
assert requests[5][2] == {"mode": "Skip", "behavior": "Restrict"}
|
||||||
for i, (_, headers, _) in enumerate(requests):
|
for i, (_, headers, _) in enumerate(requests):
|
||||||
headers = {key.lower(): value for key, value in headers.items()}
|
headers = {key.lower(): value for key, value in headers.items()}
|
||||||
@@ -138,3 +137,129 @@ async def test_catalog_async_and_errors(catalog_server):
|
|||||||
def test_catalog_requires_remote_endpoint(endpoint):
|
def test_catalog_requires_remote_endpoint(endpoint):
|
||||||
with pytest.raises(ValueError, match="endpoint"):
|
with pytest.raises(ValueError, match="endpoint"):
|
||||||
lancedb.connect_catalog(endpoint)
|
lancedb.connect_catalog(endpoint)
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.parametrize("asynchronous", [False, True])
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
@pytest.mark.parametrize("page_token,page_limit", [(None, None), ("start/token", 2)])
|
||||||
|
async def test_database_names_iteration(
|
||||||
|
catalog_server, asynchronous, page_token, page_limit
|
||||||
|
):
|
||||||
|
endpoint, requests, responses = catalog_server
|
||||||
|
responses.extend(
|
||||||
|
[
|
||||||
|
(200, {"namespaces": ["first", "second"], "page_token": "a/b"}),
|
||||||
|
(200, {"namespaces": [], "page_token": "next"}),
|
||||||
|
(200, {"namespaces": ["last", "final"], "page_token": ""}),
|
||||||
|
]
|
||||||
|
)
|
||||||
|
catalog = (
|
||||||
|
await lancedb.connect_catalog_async(endpoint)
|
||||||
|
if asynchronous
|
||||||
|
else lancedb.connect_catalog(endpoint)
|
||||||
|
)
|
||||||
|
names = catalog.list_databases(page_token=page_token, page_limit=page_limit)
|
||||||
|
assert requests == []
|
||||||
|
assert names.num_page_results() == 0
|
||||||
|
assert names.page_token() == page_token
|
||||||
|
for cached, expected in [(1, "first"), (0, "second")]:
|
||||||
|
assert (await names.__anext__() if asynchronous else next(names)) == expected
|
||||||
|
assert names.num_page_results() == cached
|
||||||
|
assert names.page_token() == "a/b"
|
||||||
|
assert len(requests) == 1
|
||||||
|
assert (await names.__anext__() if asynchronous else next(names)) == "last"
|
||||||
|
assert names.num_page_results() == 1
|
||||||
|
assert names.page_token() is None
|
||||||
|
assert (await names.__anext__() if asynchronous else next(names)) == "final"
|
||||||
|
assert names.num_page_results() == 0
|
||||||
|
assert names.page_token() is None
|
||||||
|
for _ in range(2):
|
||||||
|
with pytest.raises(StopAsyncIteration if asynchronous else StopIteration):
|
||||||
|
if asynchronous:
|
||||||
|
await names.__anext__()
|
||||||
|
else:
|
||||||
|
next(names)
|
||||||
|
from urllib.parse import parse_qs, urlsplit
|
||||||
|
|
||||||
|
for request, token in zip(requests, [page_token, "a/b", "next"]):
|
||||||
|
url = urlsplit(request[0])
|
||||||
|
assert url.path == "/v1/namespace/$/list"
|
||||||
|
expected = {}
|
||||||
|
if token is not None:
|
||||||
|
expected["page_token"] = [token]
|
||||||
|
if page_limit is not None:
|
||||||
|
expected["limit"] = [str(page_limit)]
|
||||||
|
assert parse_qs(url.query) == expected
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.parametrize("asynchronous", [False, True])
|
||||||
|
@pytest.mark.parametrize("fail", [False, True])
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_database_names_empty_or_error(catalog_server, asynchronous, fail):
|
||||||
|
endpoint, requests, responses = catalog_server
|
||||||
|
responses.append(
|
||||||
|
(401, {"error": "unauthorized"}) if fail else (200, {"namespaces": []})
|
||||||
|
)
|
||||||
|
catalog = (
|
||||||
|
await lancedb.connect_catalog_async(endpoint)
|
||||||
|
if asynchronous
|
||||||
|
else lancedb.connect_catalog(endpoint)
|
||||||
|
)
|
||||||
|
names = catalog.list_databases()
|
||||||
|
if fail:
|
||||||
|
with pytest.raises(HttpError):
|
||||||
|
if asynchronous:
|
||||||
|
await names.__anext__()
|
||||||
|
else:
|
||||||
|
next(names)
|
||||||
|
assert ([name async for name in names] if asynchronous else list(names)) == []
|
||||||
|
assert len(requests) == 1
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.parametrize("asynchronous", [False, True])
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_database_names_resume(catalog_server, asynchronous):
|
||||||
|
endpoint, requests, responses = catalog_server
|
||||||
|
responses.extend(
|
||||||
|
[
|
||||||
|
(200, {"namespaces": ["first"], "page_token": "resume/token"}),
|
||||||
|
(200, {"namespaces": ["second"]}),
|
||||||
|
]
|
||||||
|
)
|
||||||
|
catalog = (
|
||||||
|
await lancedb.connect_catalog_async(endpoint)
|
||||||
|
if asynchronous
|
||||||
|
else lancedb.connect_catalog(endpoint)
|
||||||
|
)
|
||||||
|
names = catalog.list_databases(page_limit=1)
|
||||||
|
assert (await names.__anext__() if asynchronous else next(names)) == "first"
|
||||||
|
assert names.num_page_results() == 0
|
||||||
|
resumed = catalog.list_databases(page_token=names.page_token(), page_limit=1)
|
||||||
|
assert resumed.page_token() == "resume/token"
|
||||||
|
assert len(requests) == 1
|
||||||
|
assert ([name async for name in resumed] if asynchronous else list(resumed)) == [
|
||||||
|
"second"
|
||||||
|
]
|
||||||
|
assert resumed.page_token() is None
|
||||||
|
assert resumed.num_page_results() == 0
|
||||||
|
assert requests[1][0] == "/v1/namespace/$/list?limit=1&page_token=resume%2Ftoken"
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.parametrize("asynchronous", [False, True])
|
||||||
|
@pytest.mark.parametrize("page_limit", [0, 2**32 - 1])
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_database_names_invalid_limit(catalog_server, asynchronous, page_limit):
|
||||||
|
endpoint, requests, _ = catalog_server
|
||||||
|
catalog = (
|
||||||
|
await lancedb.connect_catalog_async(endpoint)
|
||||||
|
if asynchronous
|
||||||
|
else lancedb.connect_catalog(endpoint)
|
||||||
|
)
|
||||||
|
names = catalog.list_databases(page_limit=page_limit)
|
||||||
|
with pytest.raises(ValueError, match="limit"):
|
||||||
|
if asynchronous:
|
||||||
|
await names.__anext__()
|
||||||
|
else:
|
||||||
|
next(names)
|
||||||
|
assert names.num_page_results() == 0
|
||||||
|
assert requests == []
|
||||||
|
|||||||
+45
-15
@@ -1,12 +1,13 @@
|
|||||||
// SPDX-License-Identifier: Apache-2.0
|
// SPDX-License-Identifier: Apache-2.0
|
||||||
// SPDX-FileCopyrightText: Copyright The LanceDB Authors
|
// SPDX-FileCopyrightText: Copyright The LanceDB Authors
|
||||||
|
|
||||||
|
use std::sync::{Arc, Mutex};
|
||||||
use std::time::Duration;
|
use std::time::Duration;
|
||||||
|
|
||||||
use lancedb::catalog::{
|
use futures::{StreamExt, future::poll_fn};
|
||||||
CatalogConnection, CreateDatabaseRequest, DropDatabaseRequest, ListDatabasesRequest,
|
|
||||||
};
|
use lancedb::catalog::{CatalogConnection, CreateDatabaseRequest, DropDatabaseRequest};
|
||||||
use pyo3::exceptions::PyValueError;
|
use pyo3::exceptions::{PyStopAsyncIteration, PyValueError};
|
||||||
use pyo3::{Bound, PyAny, PyRef, PyResult, Python, pyclass, pyfunction, pymethods};
|
use pyo3::{Bound, PyAny, PyRef, PyResult, Python, pyclass, pyfunction, pymethods};
|
||||||
|
|
||||||
use crate::connection::{Connection, PyClientConfig};
|
use crate::connection::{Connection, PyClientConfig};
|
||||||
@@ -67,19 +68,48 @@ impl Catalog {
|
|||||||
})
|
})
|
||||||
}
|
}
|
||||||
|
|
||||||
#[pyo3(signature = (*, limit=None, page_token=None))]
|
#[pyo3(signature = (*, page_token=None, page_limit=None))]
|
||||||
fn list_databases<'py>(
|
fn list_databases(&self, page_token: Option<String>, page_limit: Option<u32>) -> DatabaseNames {
|
||||||
self_: PyRef<'py, Self>,
|
DatabaseNames {
|
||||||
limit: Option<u32>,
|
inner: Arc::new(Mutex::new(
|
||||||
page_token: Option<String>,
|
self.inner.list_databases(page_token, page_limit),
|
||||||
) -> PyResult<Bound<'py, PyAny>> {
|
)),
|
||||||
|
next_lock: Arc::new(tokio::sync::Mutex::new(())),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
#[pyclass]
|
||||||
|
pub struct DatabaseNames {
|
||||||
|
inner: Arc<Mutex<lancedb::catalog::DatabaseNames>>,
|
||||||
|
next_lock: Arc<tokio::sync::Mutex<()>>,
|
||||||
|
}
|
||||||
|
|
||||||
|
#[pymethods]
|
||||||
|
impl DatabaseNames {
|
||||||
|
fn num_page_results(&self, py: Python<'_>) -> usize {
|
||||||
|
py.detach(|| self.inner.lock().unwrap().num_page_results())
|
||||||
|
}
|
||||||
|
|
||||||
|
fn page_token(&self, py: Python<'_>) -> Option<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 inner = self_.inner.clone();
|
||||||
|
let next_lock = self_.next_lock.clone();
|
||||||
future_into_py(self_.py(), async move {
|
future_into_py(self_.py(), async move {
|
||||||
let mut request = ListDatabasesRequest::default();
|
// Serialize advances while allowing synchronous state inspection
|
||||||
request.limit = limit;
|
// between polls, including while a REST request is pending.
|
||||||
request.page_token = page_token;
|
let _guard = next_lock.lock().await;
|
||||||
let response = inner.list_databases(request).await.infer_error()?;
|
poll_fn(|cx| inner.lock().unwrap().poll_next_unpin(cx))
|
||||||
Ok((response.databases, response.page_token))
|
.await
|
||||||
|
.ok_or_else(|| PyStopAsyncIteration::new_err(""))?
|
||||||
|
.infer_error()
|
||||||
})
|
})
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -64,6 +64,7 @@ pub fn _lancedb(py: Python, m: &Bound<'_, PyModule>) -> PyResult<()> {
|
|||||||
env_logger::init_from_env(env);
|
env_logger::init_from_env(env);
|
||||||
m.add_class::<Connection>()?;
|
m.add_class::<Connection>()?;
|
||||||
m.add_class::<catalog::Catalog>()?;
|
m.add_class::<catalog::Catalog>()?;
|
||||||
|
m.add_class::<catalog::DatabaseNames>()?;
|
||||||
m.add_function(wrap_pyfunction!(catalog::connect_catalog, m)?)?;
|
m.add_function(wrap_pyfunction!(catalog::connect_catalog, m)?)?;
|
||||||
m.add_class::<Session>()?;
|
m.add_class::<Session>()?;
|
||||||
m.add_class::<Table>()?;
|
m.add_class::<Table>()?;
|
||||||
|
|||||||
+100
-5
@@ -15,7 +15,12 @@
|
|||||||
//! ```
|
//! ```
|
||||||
|
|
||||||
use std::fmt;
|
use std::fmt;
|
||||||
|
use std::pin::Pin;
|
||||||
use std::sync::Arc;
|
use std::sync::Arc;
|
||||||
|
use std::task::{Context, Poll};
|
||||||
|
|
||||||
|
use futures::future::BoxFuture;
|
||||||
|
use futures::{FutureExt, Stream};
|
||||||
|
|
||||||
use crate::Result;
|
use crate::Result;
|
||||||
use crate::connection::Connection;
|
use crate::connection::Connection;
|
||||||
@@ -118,6 +123,73 @@ pub struct ListDatabasesResponse {
|
|||||||
pub page_token: Option<String>,
|
pub page_token: Option<String>,
|
||||||
}
|
}
|
||||||
|
|
||||||
|
/// A lazy stream of database names with inspectable pagination state.
|
||||||
|
///
|
||||||
|
/// The continuation token advances when a page is fetched, so resuming from it
|
||||||
|
/// skips any names still cached in this stream. Drain the cache before saving a
|
||||||
|
/// token if those names must be included. Errors terminate the stream; the token
|
||||||
|
/// for the failed request remains available for resuming with a new stream.
|
||||||
|
pub struct DatabaseNames {
|
||||||
|
catalog: Arc<dyn Catalog>,
|
||||||
|
names: std::vec::IntoIter<String>,
|
||||||
|
page_token: Option<String>,
|
||||||
|
page_limit: Option<u32>,
|
||||||
|
done: bool,
|
||||||
|
pending: Option<BoxFuture<'static, Result<ListDatabasesResponse>>>,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl DatabaseNames {
|
||||||
|
/// Number of names available without another REST request.
|
||||||
|
pub fn num_page_results(&self) -> usize {
|
||||||
|
self.names.len()
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Token for the next REST request, or `None` after the last page is fetched.
|
||||||
|
///
|
||||||
|
/// Before the first request, returns the supplied starting token (`None`
|
||||||
|
/// means start from the beginning). A final page can still have cached names.
|
||||||
|
pub fn page_token(&self) -> Option<&str> {
|
||||||
|
self.page_token.as_deref()
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
impl Stream for DatabaseNames {
|
||||||
|
type Item = Result<String>;
|
||||||
|
|
||||||
|
fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
|
||||||
|
let this = self.get_mut();
|
||||||
|
loop {
|
||||||
|
if let Some(name) = this.names.next() {
|
||||||
|
return Poll::Ready(Some(Ok(name)));
|
||||||
|
}
|
||||||
|
if this.done {
|
||||||
|
return Poll::Ready(None);
|
||||||
|
}
|
||||||
|
let pending = this.pending.get_or_insert_with(|| {
|
||||||
|
let catalog = this.catalog.clone();
|
||||||
|
let request = ListDatabasesRequest {
|
||||||
|
page_token: this.page_token.clone(),
|
||||||
|
limit: this.page_limit,
|
||||||
|
};
|
||||||
|
async move { catalog.list_databases(request).await }.boxed()
|
||||||
|
});
|
||||||
|
let response = futures::ready!(pending.as_mut().poll(cx));
|
||||||
|
this.pending = None;
|
||||||
|
match response {
|
||||||
|
Ok(response) => {
|
||||||
|
this.page_token = response.page_token.filter(|token| !token.is_empty());
|
||||||
|
this.done = this.page_token.is_none();
|
||||||
|
this.names = response.databases.into_iter();
|
||||||
|
}
|
||||||
|
Err(error) => {
|
||||||
|
this.done = true;
|
||||||
|
return Poll::Ready(Some(Err(error)));
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
/// A backend that manages databases. Implementations own database lifecycle semantics.
|
/// A backend that manages databases. Implementations own database lifecycle semantics.
|
||||||
#[async_trait::async_trait]
|
#[async_trait::async_trait]
|
||||||
pub trait Catalog: Send + Sync + std::fmt::Debug + 'static {
|
pub trait Catalog: Send + Sync + std::fmt::Debug + 'static {
|
||||||
@@ -188,12 +260,35 @@ impl CatalogConnection {
|
|||||||
self.catalog.drop_database(request.into()).await
|
self.catalog.drop_database(request.into()).await
|
||||||
}
|
}
|
||||||
|
|
||||||
/// List a page of databases. Pass [`ListDatabasesRequest::default`] for the first page.
|
/// Iterate lazily over all database names, fetching more results as needed.
|
||||||
pub async fn list_databases(
|
///
|
||||||
|
/// `page_token` starts at a saved continuation token; `None` starts at the
|
||||||
|
/// beginning. `page_limit` limits each REST response (not the whole stream);
|
||||||
|
/// `None` uses the server default. Errors terminate the listing.
|
||||||
|
///
|
||||||
|
/// ```
|
||||||
|
/// # async fn example(catalog: &lancedb::catalog::CatalogConnection) -> lancedb::Result<()> {
|
||||||
|
/// use futures::TryStreamExt;
|
||||||
|
/// let mut names = catalog.list_databases(None, None);
|
||||||
|
/// while let Some(name) = names.try_next().await? {
|
||||||
|
/// println!("{name}");
|
||||||
|
/// }
|
||||||
|
/// # Ok(())
|
||||||
|
/// # }
|
||||||
|
/// ```
|
||||||
|
pub fn list_databases(
|
||||||
&self,
|
&self,
|
||||||
request: ListDatabasesRequest,
|
page_token: Option<String>,
|
||||||
) -> Result<ListDatabasesResponse> {
|
page_limit: Option<u32>,
|
||||||
self.catalog.list_databases(request).await
|
) -> DatabaseNames {
|
||||||
|
DatabaseNames {
|
||||||
|
catalog: self.catalog.clone(),
|
||||||
|
names: Vec::new().into_iter(),
|
||||||
|
page_token,
|
||||||
|
page_limit,
|
||||||
|
done: false,
|
||||||
|
pending: None,
|
||||||
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
/// Connect to a database without creating it if it is missing.
|
/// Connect to a database without creating it if it is missing.
|
||||||
|
|||||||
@@ -416,6 +416,7 @@ mod tests {
|
|||||||
RemoteCatalog::try_new(&endpoint, options).unwrap(),
|
RemoteCatalog::try_new(&endpoint, options).unwrap(),
|
||||||
));
|
));
|
||||||
let page = catalog
|
let page = catalog
|
||||||
|
.catalog()
|
||||||
.list_databases(ListDatabasesRequest::default().limit(1).page_token("a/b"))
|
.list_databases(ListDatabasesRequest::default().limit(1).page_token("a/b"))
|
||||||
.await
|
.await
|
||||||
.unwrap();
|
.unwrap();
|
||||||
@@ -443,6 +444,7 @@ mod tests {
|
|||||||
.unwrap();
|
.unwrap();
|
||||||
assert!(
|
assert!(
|
||||||
catalog
|
catalog
|
||||||
|
.catalog()
|
||||||
.list_databases(ListDatabasesRequest::default())
|
.list_databases(ListDatabasesRequest::default())
|
||||||
.await
|
.await
|
||||||
.unwrap()
|
.unwrap()
|
||||||
@@ -538,6 +540,128 @@ mod tests {
|
|||||||
);
|
);
|
||||||
}
|
}
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
async fn database_names_follow_pages_lazily() {
|
||||||
|
use futures::TryStreamExt;
|
||||||
|
|
||||||
|
let (endpoint, task) = server(vec![
|
||||||
|
(
|
||||||
|
200,
|
||||||
|
json!({"namespaces": ["first", "second"], "page_token": "a/b"}),
|
||||||
|
),
|
||||||
|
(200, json!({"namespaces": [], "page_token": "next"})),
|
||||||
|
(200, json!({"namespaces": ["last"], "page_token": ""})),
|
||||||
|
])
|
||||||
|
.await;
|
||||||
|
let catalog = CatalogConnection::new(Arc::new(
|
||||||
|
RemoteCatalog::try_new(endpoint, RemoteCatalogOptions::default()).unwrap(),
|
||||||
|
));
|
||||||
|
let mut names = catalog.list_databases(None, None);
|
||||||
|
drop(catalog);
|
||||||
|
assert_eq!(names.num_page_results(), 0);
|
||||||
|
assert_eq!(names.page_token(), None);
|
||||||
|
assert_eq!(names.try_next().await.unwrap().as_deref(), Some("first"));
|
||||||
|
assert_eq!(names.num_page_results(), 1);
|
||||||
|
assert_eq!(names.page_token(), Some("a/b"));
|
||||||
|
assert_eq!(names.try_next().await.unwrap().as_deref(), Some("second"));
|
||||||
|
assert_eq!(names.num_page_results(), 0);
|
||||||
|
assert_eq!(names.page_token(), Some("a/b"));
|
||||||
|
assert!(!task.is_finished());
|
||||||
|
assert_eq!(names.try_next().await.unwrap().as_deref(), Some("last"));
|
||||||
|
assert_eq!(names.try_next().await.unwrap(), None);
|
||||||
|
assert_eq!(names.try_next().await.unwrap(), None);
|
||||||
|
let requests = task.await.unwrap();
|
||||||
|
assert_eq!(requests[0].line, "GET /v1/namespace/$/list HTTP/1.1");
|
||||||
|
assert_eq!(
|
||||||
|
requests[1].line,
|
||||||
|
"GET /v1/namespace/$/list?page_token=a%2Fb HTTP/1.1"
|
||||||
|
);
|
||||||
|
assert_eq!(
|
||||||
|
requests[2].line,
|
||||||
|
"GET /v1/namespace/$/list?page_token=next HTTP/1.1"
|
||||||
|
);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
async fn database_names_stop_after_error() {
|
||||||
|
use futures::StreamExt;
|
||||||
|
|
||||||
|
let (endpoint, task) = server(vec![
|
||||||
|
(200, json!({"namespaces": ["first"], "page_token": "next"})),
|
||||||
|
(401, json!({"error": "unauthorized"})),
|
||||||
|
])
|
||||||
|
.await;
|
||||||
|
let catalog = CatalogConnection::new(Arc::new(
|
||||||
|
RemoteCatalog::try_new(endpoint, RemoteCatalogOptions::default()).unwrap(),
|
||||||
|
));
|
||||||
|
let mut names = catalog.list_databases(None, None);
|
||||||
|
assert_eq!(names.next().await.unwrap().unwrap(), "first");
|
||||||
|
assert!(names.next().await.unwrap().is_err());
|
||||||
|
assert_eq!(names.num_page_results(), 0);
|
||||||
|
assert_eq!(names.page_token(), Some("next"));
|
||||||
|
assert!(names.next().await.is_none());
|
||||||
|
task.await.unwrap();
|
||||||
|
}
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
async fn database_names_can_resume_with_page_limit() {
|
||||||
|
use futures::TryStreamExt;
|
||||||
|
let (endpoint, task) = server(vec![
|
||||||
|
(
|
||||||
|
200,
|
||||||
|
json!({"namespaces": ["first", "second"], "page_token": "next/token"}),
|
||||||
|
),
|
||||||
|
(200, json!({"namespaces": ["third", "fourth"]})),
|
||||||
|
])
|
||||||
|
.await;
|
||||||
|
let catalog = CatalogConnection::new(Arc::new(
|
||||||
|
RemoteCatalog::try_new(endpoint, RemoteCatalogOptions::default()).unwrap(),
|
||||||
|
));
|
||||||
|
let mut names = catalog.list_databases(Some("start/token".into()), Some(2));
|
||||||
|
assert_eq!(names.page_token(), Some("start/token"));
|
||||||
|
assert_eq!(names.num_page_results(), 0);
|
||||||
|
assert_eq!(names.try_next().await.unwrap().as_deref(), Some("first"));
|
||||||
|
assert_eq!(names.page_token(), Some("next/token"));
|
||||||
|
assert_eq!(names.num_page_results(), 1);
|
||||||
|
assert_eq!(names.try_next().await.unwrap().as_deref(), Some("second"));
|
||||||
|
assert_eq!(names.num_page_results(), 0);
|
||||||
|
let token = names.page_token().map(str::to_owned);
|
||||||
|
drop(names);
|
||||||
|
let mut resumed = catalog.list_databases(token, Some(2));
|
||||||
|
assert_eq!(resumed.try_next().await.unwrap().as_deref(), Some("third"));
|
||||||
|
assert_eq!(resumed.page_token(), None);
|
||||||
|
assert_eq!(resumed.num_page_results(), 1);
|
||||||
|
assert_eq!(resumed.try_next().await.unwrap().as_deref(), Some("fourth"));
|
||||||
|
assert_eq!(resumed.num_page_results(), 0);
|
||||||
|
assert_eq!(resumed.try_next().await.unwrap(), None);
|
||||||
|
let requests = task.await.unwrap();
|
||||||
|
assert_eq!(
|
||||||
|
requests[0].line,
|
||||||
|
"GET /v1/namespace/$/list?limit=2&page_token=start%2Ftoken HTTP/1.1"
|
||||||
|
);
|
||||||
|
assert_eq!(
|
||||||
|
requests[1].line,
|
||||||
|
"GET /v1/namespace/$/list?limit=2&page_token=next%2Ftoken HTTP/1.1"
|
||||||
|
);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
async fn database_names_reject_invalid_page_limit() {
|
||||||
|
use futures::TryStreamExt;
|
||||||
|
let catalog = CatalogConnection::new(Arc::new(
|
||||||
|
RemoteCatalog::try_new("http://127.0.0.1:1", RemoteCatalogOptions::default()).unwrap(),
|
||||||
|
));
|
||||||
|
for limit in [0, u32::MAX] {
|
||||||
|
let mut names = catalog.list_databases(None, Some(limit));
|
||||||
|
assert!(matches!(
|
||||||
|
names.try_next().await,
|
||||||
|
Err(Error::InvalidInput { .. })
|
||||||
|
));
|
||||||
|
assert_eq!(names.num_page_results(), 0);
|
||||||
|
assert_eq!(names.try_next().await.unwrap(), None);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
#[tokio::test]
|
#[tokio::test]
|
||||||
async fn catalog_validates_before_sending_requests() {
|
async fn catalog_validates_before_sending_requests() {
|
||||||
for endpoint in [
|
for endpoint in [
|
||||||
|
|||||||
Reference in New Issue
Block a user