mirror of
https://github.com/lancedb/lancedb.git
synced 2026-10-14 07:42:28 +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()
|
||||
|
||||
```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
|
||||
|
||||
* **options** = `{}`
|
||||
|
||||
* **options.limit?**: `number`
|
||||
* **options.pageLimit?**: `number`
|
||||
|
||||
* **options.pageToken?**: `string`
|
||||
|
||||
#### 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
|
||||
const catalog = await connectCatalog("https://my-server.example", { apiKey: "secret" });
|
||||
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)
|
||||
- [Catalog](classes/Catalog.md)
|
||||
- [Connection](classes/Connection.md)
|
||||
- [DatabaseNames](classes/DatabaseNames.md)
|
||||
- [HeaderProvider](classes/HeaderProvider.md)
|
||||
- [Index](classes/Index.md)
|
||||
- [Job](classes/Job.md)
|
||||
@@ -108,7 +109,6 @@
|
||||
- [JobEventsOptions](interfaces/JobEventsOptions.md)
|
||||
- [JobFailureInfo](interfaces/JobFailureInfo.md)
|
||||
- [JobInfo](interfaces/JobInfo.md)
|
||||
- [ListDatabasesResponse](interfaces/ListDatabasesResponse.md)
|
||||
- [ListNamespacesOptions](interfaces/ListNamespacesOptions.md)
|
||||
- [ListNamespacesResponse](interfaces/ListNamespacesResponse.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.ListDatabasesResponse
|
||||
::: lancedb.catalog.DatabaseNames
|
||||
|
||||
|
||||
## Remote SQL
|
||||
|
||||
@@ -440,6 +441,8 @@ latency on your workload before opting in.
|
||||
|
||||
::: lancedb.catalog.AsyncCatalog
|
||||
|
||||
::: lancedb.catalog.AsyncDatabaseNames
|
||||
|
||||
Connections represent a connection to a LanceDb database and
|
||||
can be used to create, list, or open tables.
|
||||
|
||||
|
||||
@@ -80,9 +80,12 @@ describe("remote catalog", () => {
|
||||
const second = await catalog.connectDatabase("other");
|
||||
expect(await second.tableNames()).toEqual([]);
|
||||
expect(await first.tableNames()).toEqual([]);
|
||||
expect(
|
||||
await catalog.listDatabases({ limit: 1, pageToken: "a/b" }),
|
||||
).toEqual({ databases: ["team/search"], pageToken: "next" });
|
||||
const names = catalog.listDatabases({ pageLimit: 1, pageToken: "a/b" });
|
||||
expect(await names.next()).toEqual({
|
||||
done: false,
|
||||
value: "team/search",
|
||||
});
|
||||
expect(names.pageToken()).toBe("next");
|
||||
await catalog.dropDatabase("team/search", { ignoreMissing: true });
|
||||
expect(requests[0].url).toBe("/v1/namespace/team%2Fsearch/create");
|
||||
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 () => {
|
||||
await expect(connectCatalog("/tmp/catalog")).rejects.toThrow();
|
||||
const catalog = await connectCatalog("http://127.0.0.1:1");
|
||||
for (const limit of [0, -1, 1.5, 2147483648]) {
|
||||
await expect(catalog.listDatabases({ limit })).rejects.toThrow("limit");
|
||||
for (const limit of [
|
||||
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(
|
||||
"Invalid database name",
|
||||
|
||||
+62
-12
@@ -5,9 +5,9 @@ import { Connection, LocalConnection } from "./connection";
|
||||
import { HeaderProvider } from "./header";
|
||||
import {
|
||||
JsHeaderProvider,
|
||||
ListDatabasesResponse,
|
||||
Catalog as NativeCatalog,
|
||||
CatalogOptions as NativeCatalogOptions,
|
||||
DatabaseNames as NativeDatabaseNames,
|
||||
} from "./native.js";
|
||||
import { OAuthConfig } from "./oauth";
|
||||
|
||||
@@ -21,7 +21,41 @@ export interface CatalogOptions
|
||||
| (() => 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. */
|
||||
export class Catalog {
|
||||
@@ -56,21 +90,35 @@ export class Catalog {
|
||||
await this.inner.dropDatabase(name, options.ignoreMissing);
|
||||
}
|
||||
|
||||
/** List one page of databases; pass pageToken from a response for the next page. */
|
||||
async listDatabases(
|
||||
options: { limit?: number; pageToken?: string } = {},
|
||||
): Promise<ListDatabasesResponse> {
|
||||
/**
|
||||
* 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.
|
||||
*
|
||||
* @example
|
||||
* ```ts
|
||||
* for await (const name of catalog.listDatabases({ pageLimit: 20 })) {
|
||||
* console.log(name);
|
||||
* }
|
||||
* ```
|
||||
*/
|
||||
listDatabases(
|
||||
options: { pageLimit?: number; pageToken?: string } = {},
|
||||
): DatabaseNames {
|
||||
if (
|
||||
options.limit !== undefined &&
|
||||
(!Number.isInteger(options.limit) ||
|
||||
options.limit <= 0 ||
|
||||
options.limit > 2147483647)
|
||||
options.pageLimit !== undefined &&
|
||||
(!Number.isInteger(options.pageLimit) ||
|
||||
options.pageLimit <= 0 ||
|
||||
options.pageLimit > 2147483647)
|
||||
) {
|
||||
throw new Error(
|
||||
"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
|
||||
* const catalog = await connectCatalog("https://my-server.example", { apiKey: "secret" });
|
||||
* 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(
|
||||
|
||||
@@ -646,6 +646,6 @@ export async function connectNamespace(
|
||||
export {
|
||||
Catalog,
|
||||
CatalogOptions,
|
||||
ListDatabasesResponse,
|
||||
DatabaseNames,
|
||||
connectCatalog,
|
||||
} from "./catalog";
|
||||
|
||||
+42
-19
@@ -1,12 +1,12 @@
|
||||
// SPDX-License-Identifier: Apache-2.0
|
||||
// SPDX-FileCopyrightText: Copyright The LanceDB Authors
|
||||
|
||||
use std::sync::Arc;
|
||||
use std::sync::{Arc, Mutex};
|
||||
use std::time::Duration;
|
||||
|
||||
use lancedb::catalog::{
|
||||
CatalogConnection, CreateDatabaseRequest, DropDatabaseRequest, ListDatabasesRequest,
|
||||
};
|
||||
use futures::{StreamExt, future::poll_fn};
|
||||
|
||||
use lancedb::catalog::{CatalogConnection, CreateDatabaseRequest, DropDatabaseRequest};
|
||||
use napi::bindgen_prelude::*;
|
||||
use napi_derive::napi;
|
||||
|
||||
@@ -25,10 +25,37 @@ pub struct CatalogOptions {
|
||||
pub oauth_config: Option<OAuthConfig>,
|
||||
}
|
||||
|
||||
#[napi(object)]
|
||||
pub struct ListDatabasesResponse {
|
||||
pub databases: Vec<String>,
|
||||
pub page_token: Option<String>,
|
||||
/// A lazy iterator over database names with inspectable pagination state.
|
||||
#[napi]
|
||||
pub struct DatabaseNames {
|
||||
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]
|
||||
@@ -106,18 +133,14 @@ impl Catalog {
|
||||
.default_error()
|
||||
}
|
||||
#[napi]
|
||||
pub async fn list_databases(
|
||||
pub fn list_databases(
|
||||
&self,
|
||||
limit: Option<u32>,
|
||||
page_token: Option<String>,
|
||||
) -> Result<ListDatabasesResponse> {
|
||||
let mut request = ListDatabasesRequest::default();
|
||||
request.limit = limit;
|
||||
request.page_token = page_token;
|
||||
let response = self.inner.list_databases(request).await.default_error()?;
|
||||
Ok(ListDatabasesResponse {
|
||||
databases: response.databases,
|
||||
page_token: response.page_token,
|
||||
})
|
||||
page_limit: Option<u32>,
|
||||
) -> DatabaseNames {
|
||||
DatabaseNames {
|
||||
inner: Mutex::new(self.inner.list_databases(page_token, page_limit)),
|
||||
next_lock: futures::lock::Mutex::new(()),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -62,7 +62,6 @@ from .namespace import (
|
||||
from .catalog import (
|
||||
AsyncCatalog,
|
||||
Catalog,
|
||||
ListDatabasesResponse,
|
||||
connect_catalog,
|
||||
connect_catalog_async,
|
||||
)
|
||||
@@ -578,7 +577,6 @@ async def connect_async(
|
||||
__all__ = [
|
||||
"Catalog",
|
||||
"AsyncCatalog",
|
||||
"ListDatabasesResponse",
|
||||
"connect_catalog",
|
||||
"connect_catalog_async",
|
||||
"AsyncMaterializedView",
|
||||
|
||||
@@ -946,9 +946,15 @@ class Catalog:
|
||||
async def drop_database(
|
||||
self, name: str, *, ignore_missing: bool = False
|
||||
) -> None: ...
|
||||
async def list_databases(
|
||||
self, *, limit: Optional[int] = None, page_token: Optional[str] = None
|
||||
) -> tuple[list[str], Optional[str]]: ...
|
||||
def list_databases(
|
||||
self, *, page_token: Optional[str] = None, page_limit: Optional[int] = None
|
||||
) -> 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(
|
||||
endpoint: str,
|
||||
|
||||
@@ -3,9 +3,8 @@
|
||||
|
||||
"""Remote catalogs manage databases through a server's root namespace."""
|
||||
|
||||
from dataclasses import dataclass
|
||||
from datetime import timedelta
|
||||
from typing import Any, Optional, Union
|
||||
from typing import Any, AsyncIterator, Iterator, Optional, Union
|
||||
|
||||
from . import _lancedb
|
||||
from .background_loop import LOOP
|
||||
@@ -14,12 +13,66 @@ from .remote import ClientConfig, OAuthConfig
|
||||
from .remote.db import RemoteDBConnection
|
||||
|
||||
|
||||
@dataclass
|
||||
class ListDatabasesResponse:
|
||||
"""A page of database names and an optional continuation token."""
|
||||
class AsyncDatabaseNames(AsyncIterator[str]):
|
||||
"""An async iterator of database names with pagination state.
|
||||
|
||||
databases: list[str]
|
||||
page_token: Optional[str] = None
|
||||
Returned by [list_databases][lancedb.catalog.AsyncCatalog.list_databases].
|
||||
"""
|
||||
|
||||
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:
|
||||
@@ -50,14 +103,25 @@ class AsyncCatalog:
|
||||
"""Connect to an existing database by its logical name."""
|
||||
return AsyncConnection(await self._inner.connect_database(name))
|
||||
|
||||
async def list_databases(
|
||||
self, *, limit: Optional[int] = None, page_token: Optional[str] = None
|
||||
) -> ListDatabasesResponse:
|
||||
"""List a page of databases. Pass the returned token for the next page."""
|
||||
names, token = await self._inner.list_databases(
|
||||
limit=limit, page_token=page_token
|
||||
def list_databases(
|
||||
self, *, page_token: Optional[str] = None, page_limit: Optional[int] = None
|
||||
) -> AsyncDatabaseNames:
|
||||
"""Iterate lazily over all database names using ``async for``.
|
||||
|
||||
``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:
|
||||
"""Drop an empty database. A nonempty database is an error."""
|
||||
@@ -73,7 +137,8 @@ class Catalog:
|
||||
```python
|
||||
catalog = lancedb.connect_catalog("https://my-server.example", api_key="secret")
|
||||
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(
|
||||
self, *, limit: Optional[int] = None, page_token: Optional[str] = None
|
||||
) -> ListDatabasesResponse:
|
||||
"""List a page of databases. Pass the returned token for the next page."""
|
||||
return LOOP.run(self._inner.list_databases(limit=limit, page_token=page_token))
|
||||
self, *, page_token: Optional[str] = None, page_limit: Optional[int] = None
|
||||
) -> DatabaseNames:
|
||||
"""Iterate lazily over database names, fetching more results as needed.
|
||||
|
||||
``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:
|
||||
"""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, {"tables": []}),
|
||||
(200, {"tables": []}),
|
||||
(200, {"namespaces": ["team/search"], "page_token": "next"}),
|
||||
(200, {"namespaces": ["team/search"]}),
|
||||
(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"):
|
||||
connection.execute_query_async("SELECT 1")
|
||||
assert restored.table_names() == []
|
||||
page = catalog.list_databases(limit=1, page_token="a/b")
|
||||
assert page == lancedb.ListDatabasesResponse(["team/search"], "next")
|
||||
assert list(catalog.list_databases()) == ["team/search"]
|
||||
catalog.drop_database("team/search", ignore_missing=True)
|
||||
assert requests[0][0] == "/v1/namespace/team%2Fsearch/create"
|
||||
assert requests[0][2] == {"mode": "ExistOk"}
|
||||
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"}
|
||||
for i, (_, headers, _) in enumerate(requests):
|
||||
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):
|
||||
with pytest.raises(ValueError, match="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-FileCopyrightText: Copyright The LanceDB Authors
|
||||
|
||||
use std::sync::{Arc, Mutex};
|
||||
use std::time::Duration;
|
||||
|
||||
use lancedb::catalog::{
|
||||
CatalogConnection, CreateDatabaseRequest, DropDatabaseRequest, ListDatabasesRequest,
|
||||
};
|
||||
use pyo3::exceptions::PyValueError;
|
||||
use futures::{StreamExt, future::poll_fn};
|
||||
|
||||
use lancedb::catalog::{CatalogConnection, CreateDatabaseRequest, DropDatabaseRequest};
|
||||
use pyo3::exceptions::{PyStopAsyncIteration, PyValueError};
|
||||
use pyo3::{Bound, PyAny, PyRef, PyResult, Python, pyclass, pyfunction, pymethods};
|
||||
|
||||
use crate::connection::{Connection, PyClientConfig};
|
||||
@@ -67,19 +68,48 @@ impl Catalog {
|
||||
})
|
||||
}
|
||||
|
||||
#[pyo3(signature = (*, limit=None, page_token=None))]
|
||||
fn list_databases<'py>(
|
||||
self_: PyRef<'py, Self>,
|
||||
limit: Option<u32>,
|
||||
page_token: Option<String>,
|
||||
) -> PyResult<Bound<'py, PyAny>> {
|
||||
#[pyo3(signature = (*, page_token=None, page_limit=None))]
|
||||
fn list_databases(&self, page_token: Option<String>, page_limit: Option<u32>) -> DatabaseNames {
|
||||
DatabaseNames {
|
||||
inner: Arc::new(Mutex::new(
|
||||
self.inner.list_databases(page_token, page_limit),
|
||||
)),
|
||||
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 next_lock = self_.next_lock.clone();
|
||||
future_into_py(self_.py(), async move {
|
||||
let mut request = ListDatabasesRequest::default();
|
||||
request.limit = limit;
|
||||
request.page_token = page_token;
|
||||
let response = inner.list_databases(request).await.infer_error()?;
|
||||
Ok((response.databases, response.page_token))
|
||||
// Serialize advances while allowing synchronous state inspection
|
||||
// between polls, including while a REST request is pending.
|
||||
let _guard = next_lock.lock().await;
|
||||
poll_fn(|cx| inner.lock().unwrap().poll_next_unpin(cx))
|
||||
.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);
|
||||
m.add_class::<Connection>()?;
|
||||
m.add_class::<catalog::Catalog>()?;
|
||||
m.add_class::<catalog::DatabaseNames>()?;
|
||||
m.add_function(wrap_pyfunction!(catalog::connect_catalog, m)?)?;
|
||||
m.add_class::<Session>()?;
|
||||
m.add_class::<Table>()?;
|
||||
|
||||
+100
-5
@@ -15,7 +15,12 @@
|
||||
//! ```
|
||||
|
||||
use std::fmt;
|
||||
use std::pin::Pin;
|
||||
use std::sync::Arc;
|
||||
use std::task::{Context, Poll};
|
||||
|
||||
use futures::future::BoxFuture;
|
||||
use futures::{FutureExt, Stream};
|
||||
|
||||
use crate::Result;
|
||||
use crate::connection::Connection;
|
||||
@@ -118,6 +123,73 @@ pub struct ListDatabasesResponse {
|
||||
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.
|
||||
#[async_trait::async_trait]
|
||||
pub trait Catalog: Send + Sync + std::fmt::Debug + 'static {
|
||||
@@ -188,12 +260,35 @@ impl CatalogConnection {
|
||||
self.catalog.drop_database(request.into()).await
|
||||
}
|
||||
|
||||
/// List a page of databases. Pass [`ListDatabasesRequest::default`] for the first page.
|
||||
pub async fn list_databases(
|
||||
/// Iterate lazily over all database names, fetching more results as needed.
|
||||
///
|
||||
/// `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,
|
||||
request: ListDatabasesRequest,
|
||||
) -> Result<ListDatabasesResponse> {
|
||||
self.catalog.list_databases(request).await
|
||||
page_token: Option<String>,
|
||||
page_limit: Option<u32>,
|
||||
) -> 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.
|
||||
|
||||
@@ -416,6 +416,7 @@ mod tests {
|
||||
RemoteCatalog::try_new(&endpoint, options).unwrap(),
|
||||
));
|
||||
let page = catalog
|
||||
.catalog()
|
||||
.list_databases(ListDatabasesRequest::default().limit(1).page_token("a/b"))
|
||||
.await
|
||||
.unwrap();
|
||||
@@ -443,6 +444,7 @@ mod tests {
|
||||
.unwrap();
|
||||
assert!(
|
||||
catalog
|
||||
.catalog()
|
||||
.list_databases(ListDatabasesRequest::default())
|
||||
.await
|
||||
.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]
|
||||
async fn catalog_validates_before_sending_requests() {
|
||||
for endpoint in [
|
||||
|
||||
Reference in New Issue
Block a user