mirror of
https://github.com/lancedb/lancedb.git
synced 2026-09-21 20:45:57 +00:00
feat: add remote catalogs and Python and TypeScript bindings (#4195)
Add a `Catalog` trait and `RemoteCatalog` for managing databases,
exposed through Rust, synchronous/asynchronous Python, and TypeScript. A
remote catalog represents the server's root namespace, and each database
is one child namespace. Create/connect return ordinary LanceDB
connections, so existing table APIs work unchanged.
## Rust API
`Catalog` is an object-safe async trait with `create_database`,
`connect_database`, `list_databases`, and `drop_database`. Backend
create/connect methods return `Arc<dyn Database>`; the public
`CatalogConnection` wraps them as `Connection` values and shares its
embedding registry with those connections. `RemoteCatalog` implements
the trait; `connect_catalog` is the convenience builder, available with
the `remote` feature.
```rust
use lancedb::catalog::{
CreateDatabaseRequest, DropDatabaseRequest, ListDatabasesRequest,
};
let catalog = lancedb::connect_catalog("https://my-server.example")
.api_key("my-api-key")
.execute()
.await?;
let db = catalog.create_database(
CreateDatabaseRequest::new("analytics").exist_ok(true),
).await?;
let connected = catalog.connect_database("analytics").await?;
let page = catalog.list_databases(
ListDatabasesRequest::default().limit(20),
).await?;
// page.databases: Vec<String>; page.page_token: Option<String>
catalog.drop_database(
DropDatabaseRequest::new("analytics").ignore_missing(true),
).await?;
```
Create/drop also accept a plain name for default behavior, e.g.
`catalog.create_database("analytics").await?`. Existing names fail
creation unless `exist_ok` is enabled; missing names fail drop unless
`ignore_missing` is enabled. Drop always requires an empty database.
## Python API
```python
import lancedb
catalog = lancedb.connect_catalog(
"https://my-server.example", api_key="my-api-key"
)
db = catalog.create_database("analytics", exist_ok=True)
connected = catalog.connect_database("analytics")
page = catalog.list_databases(limit=20)
# page.databases: list[str]; page.page_token: Optional[str]
if page.page_token is not None:
next_page = catalog.list_databases(limit=20, page_token=page.page_token)
catalog.drop_database("analytics", ignore_missing=True)
```
`connect_catalog` returns `Catalog`; create/connect return the existing
`DBConnection` API. The async equivalent is `catalog = await
lancedb.connect_catalog_async(...)`, returning `AsyncCatalog`; await
each of the same four methods, with create/connect returning
`AsyncConnection`.
## TypeScript API
```typescript
import { connectCatalog } from "@lancedb/lancedb";
const catalog = await connectCatalog("https://my-server.example", {
apiKey: "my-api-key",
});
const db = await catalog.createDatabase("analytics", { existOk: true });
const connected = await catalog.connectDatabase("analytics");
const page = await catalog.listDatabases({ limit: 20 });
// page.databases: string[]; page.pageToken?: string
if (page.pageToken !== undefined) {
const nextPage = await catalog.listDatabases({
limit: 20, pageToken: page.pageToken,
});
}
await catalog.dropDatabase("analytics", { ignoreMissing: true });
```
Create/connect return the existing `Connection` API. All four methods
are asynchronous.
## REST mapping
All paths below are relative to the catalog endpoint. `{name}` is the
logical database name encoded as one URL path component. The default
namespace delimiter is `$`, so the root identifier is encoded as `%24`.
| Catalog operation | Existing REST route | Request |
| --- | --- | --- |
| `create_database(name)` | `POST /v1/namespace/{name}/create` |
`{"mode":"Create"}`; `exist_ok=true` sends `{"mode":"ExistOk"}` |
| `connect_database(name)` | `POST /v1/namespace/{name}/describe` |
`{}`; verifies existence before returning a scoped connection |
| `list_databases(...)` | `GET /v1/namespace/%24/list` | Optional
`limit` and `page_token` query parameters |
| `drop_database(name)` | `POST /v1/namespace/{name}/drop` |
`{"mode":"Fail","behavior":"Restrict"}`; `ignore_missing=true` changes
mode to `"Skip"` |
For example, database `team/search` uses
`/v1/namespace/team%2Fsearch/create`. A paginated root listing can use
`/v1/namespace/%24/list?limit=20&page_token=a%2Fb`. The list response
retains the existing namespace wire shape,
`{"namespaces":["analytics"],"page_token":"next"}`; the SDK exposes
`namespaces` as `databases` and preserves the opaque continuation token.
An absent or empty token ends pagination. Page limits must be between 1
and 2147483647. Create/drop accept a namespace JSON response or HTTP
204.
Catalog management requests omit both `x-lancedb-database` and
`x-lancedb-database-prefix`, including values supplied through static or
dynamic headers. Returned database connections set `x-lancedb-database`
to the exact logical name and keep independent scope. API keys, OAuth or
dynamic authentication, client settings, table read consistency
settings, and an optional SQL endpoint override carry over to those
connections. OAuth cannot be combined with an API key or a custom header
provider.
For SQL through an HTTPS catalog, configure the existing SQL endpoint
contract with Rust
`.sql_host_override("grpc+tls://sql.example.com:10026")` or Python
`sql_host_override="grpc+tls://sql.example.com:10026"`. TypeScript
catalog options expose the same setting as `sqlHostOverride`. It is
inherited by created/connected databases, retained by Python connection
serialization, and initialized lazily when SQL is executed.
Create HTTP 409 maps to `DatabaseAlreadyExists`; connect/drop HTTP 404
maps to `DatabaseNotFound`, except that `ignore_missing` suppresses a
missing-database drop. Other server errors propagate. The server
enforces restricted deletion; the client never requests cascading
deletion.
Database names preserve literal slashes as part of one name. They must
be nonempty ASCII, with no control characters, surrounding whitespace,
or configured namespace delimiter, and cannot be `.` or `..`. Endpoints
must be HTTP(S) URLs without embedded credentials, query parameters, or
fragments.
## Scope
This PR adds the client API and reuses existing namespace endpoints.
Local catalogs, `__catalog` storage, location generation/sanitization,
and `__manifest` lifecycle support remain deferred; the Lance dependency
is unchanged.
The PR also runs macOS Node tests serially to avoid existing
resource-contention timeouts reproduced across recent main runs.
---------
Co-authored-by: Xuanwo <github@xuanwo.io>
This commit is contained in:
@@ -207,4 +207,4 @@ jobs:
|
||||
pnpm tsc
|
||||
- name: Test
|
||||
run: |
|
||||
pnpm test
|
||||
pnpm test --runInBand
|
||||
|
||||
@@ -0,0 +1,107 @@
|
||||
[**@lancedb/lancedb**](../README.md) • **Docs**
|
||||
|
||||
***
|
||||
|
||||
[@lancedb/lancedb](../globals.md) / Catalog
|
||||
|
||||
# Class: Catalog
|
||||
|
||||
A remote catalog manages databases through the server's root namespace.
|
||||
|
||||
## Accessors
|
||||
|
||||
### uri
|
||||
|
||||
```ts
|
||||
get uri(): string
|
||||
```
|
||||
|
||||
The root namespace endpoint.
|
||||
|
||||
#### Returns
|
||||
|
||||
`string`
|
||||
|
||||
## Methods
|
||||
|
||||
### connectDatabase()
|
||||
|
||||
```ts
|
||||
connectDatabase(name): Promise<Connection>
|
||||
```
|
||||
|
||||
Connect to an existing database by its logical name.
|
||||
|
||||
#### Parameters
|
||||
|
||||
* **name**: `string`
|
||||
|
||||
#### Returns
|
||||
|
||||
`Promise`<[`Connection`](Connection.md)>
|
||||
|
||||
***
|
||||
|
||||
### createDatabase()
|
||||
|
||||
```ts
|
||||
createDatabase(name, options): Promise<Connection>
|
||||
```
|
||||
|
||||
Create a database, or open an existing database when existOk is true.
|
||||
|
||||
#### Parameters
|
||||
|
||||
* **name**: `string`
|
||||
|
||||
* **options** = `{}`
|
||||
|
||||
* **options.existOk?**: `boolean`
|
||||
|
||||
#### Returns
|
||||
|
||||
`Promise`<[`Connection`](Connection.md)>
|
||||
|
||||
***
|
||||
|
||||
### dropDatabase()
|
||||
|
||||
```ts
|
||||
dropDatabase(name, options): Promise<void>
|
||||
```
|
||||
|
||||
Drop an empty database. The server rejects nonempty databases.
|
||||
|
||||
#### Parameters
|
||||
|
||||
* **name**: `string`
|
||||
|
||||
* **options** = `{}`
|
||||
|
||||
* **options.ignoreMissing?**: `boolean`
|
||||
|
||||
#### Returns
|
||||
|
||||
`Promise`<`void`>
|
||||
|
||||
***
|
||||
|
||||
### listDatabases()
|
||||
|
||||
```ts
|
||||
listDatabases(options): Promise<ListDatabasesResponse>
|
||||
```
|
||||
|
||||
List one page of databases; pass pageToken from a response for the next page.
|
||||
|
||||
#### Parameters
|
||||
|
||||
* **options** = `{}`
|
||||
|
||||
* **options.limit?**: `number`
|
||||
|
||||
* **options.pageToken?**: `string`
|
||||
|
||||
#### Returns
|
||||
|
||||
`Promise`<[`ListDatabasesResponse`](../interfaces/ListDatabasesResponse.md)>
|
||||
@@ -0,0 +1,32 @@
|
||||
[**@lancedb/lancedb**](../README.md) • **Docs**
|
||||
|
||||
***
|
||||
|
||||
[@lancedb/lancedb](../globals.md) / connectCatalog
|
||||
|
||||
# Function: connectCatalog()
|
||||
|
||||
```ts
|
||||
function connectCatalog(endpoint, options): Promise<Catalog>
|
||||
```
|
||||
|
||||
Connect to an HTTP(S) catalog endpoint. Catalog requests omit database-selection
|
||||
headers; opened database connections inherit authentication and client options.
|
||||
|
||||
## Parameters
|
||||
|
||||
* **endpoint**: `string`
|
||||
|
||||
* **options**: [`CatalogOptions`](../interfaces/CatalogOptions.md) = `{}`
|
||||
|
||||
## Returns
|
||||
|
||||
`Promise`<[`Catalog`](../classes/Catalog.md)>
|
||||
|
||||
## Example
|
||||
|
||||
```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 });
|
||||
```
|
||||
@@ -25,6 +25,7 @@
|
||||
- [BoostQuery](classes/BoostQuery.md)
|
||||
- [BranchContents](classes/BranchContents.md)
|
||||
- [Branches](classes/Branches.md)
|
||||
- [Catalog](classes/Catalog.md)
|
||||
- [Connection](classes/Connection.md)
|
||||
- [HeaderProvider](classes/HeaderProvider.md)
|
||||
- [Index](classes/Index.md)
|
||||
@@ -64,6 +65,7 @@
|
||||
- [BranchIndexSummary](interfaces/BranchIndexSummary.md)
|
||||
- [BranchRowCountSummary](interfaces/BranchRowCountSummary.md)
|
||||
- [BucketStats](interfaces/BucketStats.md)
|
||||
- [CatalogOptions](interfaces/CatalogOptions.md)
|
||||
- [CherryPickError](interfaces/CherryPickError.md)
|
||||
- [CherryPickPreview](interfaces/CherryPickPreview.md)
|
||||
- [CherryPickResult](interfaces/CherryPickResult.md)
|
||||
@@ -102,6 +104,7 @@
|
||||
- [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)
|
||||
@@ -167,6 +170,7 @@
|
||||
- [RecordBatchIterator](functions/RecordBatchIterator.md)
|
||||
- [blob](functions/blob.md)
|
||||
- [connect](functions/connect.md)
|
||||
- [connectCatalog](functions/connectCatalog.md)
|
||||
- [connectNamespace](functions/connectNamespace.md)
|
||||
- [instrumentLanceDbMetrics](functions/instrumentLanceDbMetrics.md)
|
||||
- [isBlobField](functions/isBlobField.md)
|
||||
|
||||
@@ -0,0 +1,81 @@
|
||||
[**@lancedb/lancedb**](../README.md) • **Docs**
|
||||
|
||||
***
|
||||
|
||||
[@lancedb/lancedb](../globals.md) / CatalogOptions
|
||||
|
||||
# Interface: CatalogOptions
|
||||
|
||||
Options shared by a catalog and the database connections it returns.
|
||||
|
||||
## Extends
|
||||
|
||||
- `Omit`<`NativeCatalogOptions`, `"oauthConfig"`>
|
||||
|
||||
## Properties
|
||||
|
||||
### apiKey?
|
||||
|
||||
```ts
|
||||
optional apiKey: string;
|
||||
```
|
||||
|
||||
#### Inherited from
|
||||
|
||||
`Omit.apiKey`
|
||||
|
||||
***
|
||||
|
||||
### clientConfig?
|
||||
|
||||
```ts
|
||||
optional clientConfig: ClientConfig;
|
||||
```
|
||||
|
||||
#### Inherited from
|
||||
|
||||
`Omit.clientConfig`
|
||||
|
||||
***
|
||||
|
||||
### headerProvider?
|
||||
|
||||
```ts
|
||||
optional headerProvider: HeaderProvider | () => Record<string, string> | Promise<Record<string, string>>;
|
||||
```
|
||||
|
||||
Called for each request to supply authentication headers.
|
||||
|
||||
***
|
||||
|
||||
### oauthConfig?
|
||||
|
||||
```ts
|
||||
optional oauthConfig: OAuthConfig;
|
||||
```
|
||||
|
||||
***
|
||||
|
||||
### readConsistencyInterval?
|
||||
|
||||
```ts
|
||||
optional readConsistencyInterval: number;
|
||||
```
|
||||
|
||||
#### Inherited from
|
||||
|
||||
`Omit.readConsistencyInterval`
|
||||
|
||||
***
|
||||
|
||||
### sqlHostOverride?
|
||||
|
||||
```ts
|
||||
optional sqlHostOverride: string;
|
||||
```
|
||||
|
||||
SQL service endpoint inherited by database connections.
|
||||
|
||||
#### Inherited from
|
||||
|
||||
`Omit.sqlHostOverride`
|
||||
@@ -0,0 +1,23 @@
|
||||
[**@lancedb/lancedb**](../README.md) • **Docs**
|
||||
|
||||
***
|
||||
|
||||
[@lancedb/lancedb](../globals.md) / ListDatabasesResponse
|
||||
|
||||
# Interface: ListDatabasesResponse
|
||||
|
||||
## Properties
|
||||
|
||||
### databases
|
||||
|
||||
```ts
|
||||
databases: string[];
|
||||
```
|
||||
|
||||
***
|
||||
|
||||
### pageToken?
|
||||
|
||||
```ts
|
||||
optional pageToken: string;
|
||||
```
|
||||
@@ -28,6 +28,17 @@ is also an [asynchronous API client](#connections-asynchronous).
|
||||
|
||||
::: lancedb.Session
|
||||
|
||||
## Catalogs (Synchronous)
|
||||
|
||||
Remote catalogs manage databases through a server's root namespace. Opened databases
|
||||
are ordinary connections. Dropping a database requires it to be empty.
|
||||
|
||||
::: lancedb.connect_catalog
|
||||
|
||||
::: lancedb.catalog.Catalog
|
||||
|
||||
::: lancedb.catalog.ListDatabasesResponse
|
||||
|
||||
## Remote SQL
|
||||
|
||||
Submit SQL against a remote LanceDB database through the connection.
|
||||
@@ -361,6 +372,10 @@ still work. Queries return descriptors. Call
|
||||
|
||||
## Connections (Asynchronous)
|
||||
|
||||
::: lancedb.connect_catalog_async
|
||||
|
||||
::: lancedb.catalog.AsyncCatalog
|
||||
|
||||
Connections represent a connection to a LanceDb database and
|
||||
can be used to create, list, or open tables.
|
||||
|
||||
|
||||
@@ -0,0 +1,142 @@
|
||||
// SPDX-License-Identifier: Apache-2.0
|
||||
// SPDX-FileCopyrightText: Copyright The LanceDB Authors
|
||||
|
||||
import * as http from "http";
|
||||
import { Catalog, connectCatalog } from "../lancedb";
|
||||
|
||||
type RecordedRequest = {
|
||||
url: string;
|
||||
headers: http.IncomingHttpHeaders;
|
||||
body: Record<string, unknown>;
|
||||
};
|
||||
|
||||
async function withCatalog(
|
||||
responses: [number, unknown][],
|
||||
callback: (catalog: Catalog, requests: RecordedRequest[]) => Promise<void>,
|
||||
) {
|
||||
const requests: RecordedRequest[] = [];
|
||||
const server = http.createServer(async (req, res) => {
|
||||
const chunks: Buffer[] = [];
|
||||
for await (const chunk of req) chunks.push(Buffer.from(chunk));
|
||||
const body = Buffer.concat(chunks).toString();
|
||||
requests.push({
|
||||
url: req.url ?? "",
|
||||
headers: req.headers,
|
||||
body: body ? JSON.parse(body) : {},
|
||||
});
|
||||
const [status, response] = responses.shift() ?? [
|
||||
500,
|
||||
{ error: "Unexpected request" },
|
||||
];
|
||||
res.writeHead(status, { "content-type": "application/json" });
|
||||
res.end(status === 204 ? undefined : JSON.stringify(response));
|
||||
});
|
||||
await new Promise<void>((resolve) => server.listen(0, "127.0.0.1", resolve));
|
||||
const address = server.address();
|
||||
if (!address || typeof address === "string")
|
||||
throw new Error("Missing server address");
|
||||
try {
|
||||
const catalog = await connectCatalog(`http://127.0.0.1:${address.port}`, {
|
||||
apiKey: "secret",
|
||||
sqlHostOverride: "grpc+tls://sql.example.com:10026",
|
||||
clientConfig: {
|
||||
extraHeaders: {
|
||||
"X-LanceDB-Database": "wrong-static",
|
||||
"X-LanceDB-Database-Prefix": "wrong",
|
||||
},
|
||||
},
|
||||
headerProvider: () => ({
|
||||
"X-LanceDB-Database": "wrong-dynamic",
|
||||
"X-LanceDB-Database-Prefix": "wrong",
|
||||
authorization: "Bearer refreshed",
|
||||
}),
|
||||
});
|
||||
await callback(catalog, requests);
|
||||
expect(responses).toHaveLength(0);
|
||||
} finally {
|
||||
server.closeAllConnections();
|
||||
await new Promise<void>((resolve) => server.close(() => resolve()));
|
||||
}
|
||||
}
|
||||
|
||||
describe("remote catalog", () => {
|
||||
it("uses root namespace routes and preserves independent database scope", async () => {
|
||||
await withCatalog(
|
||||
[
|
||||
[204, null],
|
||||
[200, { tables: [] }],
|
||||
[200, {}],
|
||||
[200, { tables: [] }],
|
||||
[200, { tables: [] }],
|
||||
// biome-ignore lint/style/useNamingConvention: server wire format
|
||||
[200, { namespaces: ["team/search"], page_token: "next" }],
|
||||
[204, null],
|
||||
],
|
||||
async (catalog, requests) => {
|
||||
const first = await catalog.createDatabase("team/search", {
|
||||
existOk: true,
|
||||
});
|
||||
expect(await first.tableNames()).toEqual([]);
|
||||
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" });
|
||||
await catalog.dropDatabase("team/search", { ignoreMissing: true });
|
||||
expect(requests[0].url).toBe("/v1/namespace/team%2Fsearch/create");
|
||||
expect(requests[0].body).toEqual({ mode: "ExistOk" });
|
||||
expect(requests[5].url).toBe(
|
||||
"/v1/namespace/%24/list?limit=1&page_token=a%2Fb",
|
||||
);
|
||||
expect(requests[6].body).toEqual({
|
||||
mode: "Skip",
|
||||
behavior: "Restrict",
|
||||
});
|
||||
for (const [i, request] of requests.entries()) {
|
||||
expect(request.headers["x-lancedb-database"]).toBe(
|
||||
i === 1 || i === 4 ? "team/search" : i === 3 ? "other" : undefined,
|
||||
);
|
||||
expect(request.headers["x-lancedb-database-prefix"]).toBeUndefined();
|
||||
expect(request.headers.authorization).toBe("Bearer refreshed");
|
||||
}
|
||||
},
|
||||
);
|
||||
});
|
||||
|
||||
it("propagates lifecycle errors and sends restricted drops", async () => {
|
||||
await withCatalog(
|
||||
[
|
||||
[404, {}],
|
||||
[409, {}],
|
||||
[400, {}],
|
||||
[404, {}],
|
||||
],
|
||||
async (catalog, requests) => {
|
||||
await expect(catalog.connectDatabase("missing")).rejects.toThrow(
|
||||
"missing",
|
||||
);
|
||||
await expect(catalog.createDatabase("exists")).rejects.toThrow(
|
||||
"exists",
|
||||
);
|
||||
await expect(catalog.dropDatabase("full")).rejects.toThrow();
|
||||
await catalog.dropDatabase("missing", { ignoreMissing: true });
|
||||
expect(requests[2].body).toEqual({
|
||||
mode: "Fail",
|
||||
behavior: "Restrict",
|
||||
});
|
||||
},
|
||||
);
|
||||
});
|
||||
|
||||
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");
|
||||
}
|
||||
await expect(catalog.connectDatabase("a$b")).rejects.toThrow(
|
||||
"Invalid database name",
|
||||
);
|
||||
});
|
||||
});
|
||||
@@ -0,0 +1,103 @@
|
||||
// SPDX-License-Identifier: Apache-2.0
|
||||
// SPDX-FileCopyrightText: Copyright The LanceDB Authors
|
||||
|
||||
import { Connection, LocalConnection } from "./connection";
|
||||
import { HeaderProvider } from "./header";
|
||||
import {
|
||||
JsHeaderProvider,
|
||||
ListDatabasesResponse,
|
||||
Catalog as NativeCatalog,
|
||||
CatalogOptions as NativeCatalogOptions,
|
||||
} from "./native.js";
|
||||
import { OAuthConfig } from "./oauth";
|
||||
|
||||
/** Options shared by a catalog and the database connections it returns. */
|
||||
export interface CatalogOptions
|
||||
extends Omit<NativeCatalogOptions, "oauthConfig"> {
|
||||
oauthConfig?: OAuthConfig;
|
||||
/** Called for each request to supply authentication headers. */
|
||||
headerProvider?:
|
||||
| HeaderProvider
|
||||
| (() => Record<string, string> | Promise<Record<string, string>>);
|
||||
}
|
||||
|
||||
export type { ListDatabasesResponse } from "./native.js";
|
||||
|
||||
/** A remote catalog manages databases through the server's root namespace. */
|
||||
export class Catalog {
|
||||
/** @hidden */
|
||||
constructor(private readonly inner: NativeCatalog) {}
|
||||
|
||||
/** The root namespace endpoint. */
|
||||
get uri(): string {
|
||||
return this.inner.uri;
|
||||
}
|
||||
|
||||
/** Create a database, or open an existing database when existOk is true. */
|
||||
async createDatabase(
|
||||
name: string,
|
||||
options: { existOk?: boolean } = {},
|
||||
): Promise<Connection> {
|
||||
return new LocalConnection(
|
||||
await this.inner.createDatabase(name, options.existOk),
|
||||
);
|
||||
}
|
||||
|
||||
/** Connect to an existing database by its logical name. */
|
||||
async connectDatabase(name: string): Promise<Connection> {
|
||||
return new LocalConnection(await this.inner.connectDatabase(name));
|
||||
}
|
||||
|
||||
/** Drop an empty database. The server rejects nonempty databases. */
|
||||
async dropDatabase(
|
||||
name: string,
|
||||
options: { ignoreMissing?: boolean } = {},
|
||||
): Promise<void> {
|
||||
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> {
|
||||
if (
|
||||
options.limit !== undefined &&
|
||||
(!Number.isInteger(options.limit) ||
|
||||
options.limit <= 0 ||
|
||||
options.limit > 2147483647)
|
||||
) {
|
||||
throw new Error(
|
||||
"Database list limit must be an integer between 1 and 2147483647",
|
||||
);
|
||||
}
|
||||
return this.inner.listDatabases(options.limit, options.pageToken);
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Connect to an HTTP(S) catalog endpoint. Catalog requests omit database-selection
|
||||
* headers; opened database connections inherit authentication and client options.
|
||||
*
|
||||
* @example
|
||||
* ```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 });
|
||||
* ```
|
||||
*/
|
||||
export async function connectCatalog(
|
||||
endpoint: string,
|
||||
options: CatalogOptions = {},
|
||||
): Promise<Catalog> {
|
||||
const { headerProvider, ...nativeOptions } = options;
|
||||
const provider = headerProvider
|
||||
? new JsHeaderProvider(async () =>
|
||||
typeof headerProvider === "function"
|
||||
? headerProvider()
|
||||
: headerProvider.getHeaders(),
|
||||
)
|
||||
: undefined;
|
||||
return new Catalog(
|
||||
await NativeCatalog.new(endpoint, nativeOptions, provider),
|
||||
);
|
||||
}
|
||||
@@ -629,3 +629,10 @@ export async function connectNamespace(
|
||||
);
|
||||
return new LocalConnection(nativeConn);
|
||||
}
|
||||
|
||||
export {
|
||||
Catalog,
|
||||
CatalogOptions,
|
||||
ListDatabasesResponse,
|
||||
connectCatalog,
|
||||
} from "./catalog";
|
||||
|
||||
@@ -0,0 +1,123 @@
|
||||
// SPDX-License-Identifier: Apache-2.0
|
||||
// SPDX-FileCopyrightText: Copyright The LanceDB Authors
|
||||
|
||||
use std::sync::Arc;
|
||||
use std::time::Duration;
|
||||
|
||||
use lancedb::catalog::{
|
||||
CatalogConnection, CreateDatabaseRequest, DropDatabaseRequest, ListDatabasesRequest,
|
||||
};
|
||||
use napi::bindgen_prelude::*;
|
||||
use napi_derive::napi;
|
||||
|
||||
use crate::connection::Connection;
|
||||
use crate::error::NapiErrorExt;
|
||||
use crate::header::JsHeaderProvider;
|
||||
use crate::remote::{ClientConfig, OAuthConfig};
|
||||
|
||||
#[napi(object)]
|
||||
pub struct CatalogOptions {
|
||||
pub api_key: Option<String>,
|
||||
pub client_config: Option<ClientConfig>,
|
||||
/// SQL service endpoint inherited by database connections.
|
||||
pub sql_host_override: Option<String>,
|
||||
pub read_consistency_interval: Option<f64>,
|
||||
pub oauth_config: Option<OAuthConfig>,
|
||||
}
|
||||
|
||||
#[napi(object)]
|
||||
pub struct ListDatabasesResponse {
|
||||
pub databases: Vec<String>,
|
||||
pub page_token: Option<String>,
|
||||
}
|
||||
|
||||
#[napi]
|
||||
pub struct Catalog {
|
||||
inner: CatalogConnection,
|
||||
}
|
||||
|
||||
#[napi]
|
||||
impl Catalog {
|
||||
#[napi(factory)]
|
||||
pub async fn new(
|
||||
endpoint: String,
|
||||
options: CatalogOptions,
|
||||
header_provider: Option<&JsHeaderProvider>,
|
||||
) -> Result<Self> {
|
||||
let mut builder = lancedb::connect_catalog(endpoint);
|
||||
if let Some(key) = options.api_key {
|
||||
builder = builder.api_key(key);
|
||||
}
|
||||
let mut config: lancedb::remote::ClientConfig =
|
||||
options.client_config.unwrap_or_default().into();
|
||||
if let Some(provider) = header_provider {
|
||||
config.header_provider = Some(Arc::new(provider.clone()));
|
||||
}
|
||||
builder = builder.client_config(config);
|
||||
if let Some(endpoint) = options.sql_host_override {
|
||||
builder = builder.sql_host_override(endpoint);
|
||||
}
|
||||
if let Some(interval) = options.read_consistency_interval {
|
||||
let interval = Duration::try_from_secs_f64(interval).map_err(|err| {
|
||||
Error::from_reason(format!("Invalid read consistency interval: {err}"))
|
||||
})?;
|
||||
builder = builder.read_consistency_interval(interval);
|
||||
}
|
||||
if let Some(oauth) = options.oauth_config {
|
||||
builder = builder.oauth_config(oauth.try_into().default_error()?);
|
||||
}
|
||||
Ok(Self {
|
||||
inner: builder.execute().await.default_error()?,
|
||||
})
|
||||
}
|
||||
|
||||
#[napi(getter)]
|
||||
pub fn uri(&self) -> String {
|
||||
self.inner.uri().to_string()
|
||||
}
|
||||
|
||||
#[napi]
|
||||
pub async fn create_database(
|
||||
&self,
|
||||
name: String,
|
||||
exist_ok: Option<bool>,
|
||||
) -> Result<Connection> {
|
||||
self.inner
|
||||
.create_database(CreateDatabaseRequest::new(name).exist_ok(exist_ok.unwrap_or(false)))
|
||||
.await
|
||||
.map(Connection::inner_new)
|
||||
.default_error()
|
||||
}
|
||||
#[napi]
|
||||
pub async fn connect_database(&self, name: String) -> Result<Connection> {
|
||||
self.inner
|
||||
.connect_database(name)
|
||||
.await
|
||||
.map(Connection::inner_new)
|
||||
.default_error()
|
||||
}
|
||||
#[napi]
|
||||
pub async fn drop_database(&self, name: String, ignore_missing: Option<bool>) -> Result<()> {
|
||||
self.inner
|
||||
.drop_database(
|
||||
DropDatabaseRequest::new(name).ignore_missing(ignore_missing.unwrap_or(false)),
|
||||
)
|
||||
.await
|
||||
.default_error()
|
||||
}
|
||||
#[napi]
|
||||
pub async 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,
|
||||
})
|
||||
}
|
||||
}
|
||||
@@ -11,6 +11,7 @@ use env_logger::Env;
|
||||
use napi_derive::*;
|
||||
|
||||
mod blob;
|
||||
mod catalog;
|
||||
mod connection;
|
||||
mod error;
|
||||
mod header;
|
||||
|
||||
@@ -52,6 +52,14 @@ from .namespace import (
|
||||
AsyncLanceNamespaceDBConnection,
|
||||
)
|
||||
|
||||
from .catalog import (
|
||||
AsyncCatalog,
|
||||
Catalog,
|
||||
ListDatabasesResponse,
|
||||
connect_catalog,
|
||||
connect_catalog_async,
|
||||
)
|
||||
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from lance.blob import BlobType as BlobType
|
||||
@@ -561,6 +569,11 @@ async def connect_async(
|
||||
|
||||
|
||||
__all__ = [
|
||||
"Catalog",
|
||||
"AsyncCatalog",
|
||||
"ListDatabasesResponse",
|
||||
"connect_catalog",
|
||||
"connect_catalog_async",
|
||||
"AsyncMaterializedView",
|
||||
"MaterializedView",
|
||||
"MaterializedViewDefinition",
|
||||
|
||||
@@ -856,3 +856,27 @@ def fts_query_to_json(query: Any) -> str: ...
|
||||
|
||||
class PermutationReader:
|
||||
def __init__(self, base_table: Table, permutation_table: Table): ...
|
||||
|
||||
class Catalog:
|
||||
@property
|
||||
def uri(self) -> str: ...
|
||||
async def create_database(
|
||||
self, name: str, *, exist_ok: bool = False
|
||||
) -> Connection: ...
|
||||
async def connect_database(self, name: str) -> Connection: ...
|
||||
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]]: ...
|
||||
|
||||
async def connect_catalog(
|
||||
endpoint: str,
|
||||
*,
|
||||
api_key: Optional[str] = None,
|
||||
client_config: Optional[Any] = None,
|
||||
sql_host_override: Optional[str] = None,
|
||||
read_consistency_interval: Optional[float] = None,
|
||||
oauth_config: Optional[Any] = None,
|
||||
) -> Catalog: ...
|
||||
|
||||
@@ -0,0 +1,196 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright The LanceDB Authors
|
||||
|
||||
"""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 . import _lancedb
|
||||
from .background_loop import LOOP
|
||||
from .db import AsyncConnection, DBConnection
|
||||
from .remote import ClientConfig, OAuthConfig
|
||||
from .remote.db import RemoteDBConnection
|
||||
|
||||
|
||||
@dataclass
|
||||
class ListDatabasesResponse:
|
||||
"""A page of database names and an optional continuation token."""
|
||||
|
||||
databases: list[str]
|
||||
page_token: Optional[str] = None
|
||||
|
||||
|
||||
class AsyncCatalog:
|
||||
"""An asynchronous remote catalog returned by
|
||||
[connect_catalog_async][lancedb.connect_catalog_async].
|
||||
|
||||
Create/connect return ordinary [AsyncConnection][lancedb.db.AsyncConnection]
|
||||
instances. Drop uses restricted behavior: remove the database's tables first.
|
||||
"""
|
||||
|
||||
def __init__(self, inner: _lancedb.Catalog):
|
||||
self._inner = inner
|
||||
|
||||
@property
|
||||
def uri(self) -> str:
|
||||
"""The catalog's root namespace endpoint."""
|
||||
return self._inner.uri
|
||||
|
||||
async def create_database(
|
||||
self, name: str, *, exist_ok: bool = False
|
||||
) -> AsyncConnection:
|
||||
"""Create a database, or open an existing one when ``exist_ok=True``."""
|
||||
return AsyncConnection(
|
||||
await self._inner.create_database(name, exist_ok=exist_ok)
|
||||
)
|
||||
|
||||
async def connect_database(self, name: str) -> AsyncConnection:
|
||||
"""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
|
||||
)
|
||||
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."""
|
||||
await self._inner.drop_database(name, ignore_missing=ignore_missing)
|
||||
|
||||
|
||||
class Catalog:
|
||||
"""A synchronous remote catalog returned by
|
||||
[connect_catalog][lancedb.connect_catalog].
|
||||
|
||||
Examples
|
||||
--------
|
||||
```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)
|
||||
```
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
inner: AsyncCatalog,
|
||||
*,
|
||||
api_key=None,
|
||||
client_config=None,
|
||||
sql_host_override: Optional[str] = None,
|
||||
oauth_config: Optional[OAuthConfig] = None,
|
||||
):
|
||||
self._inner = inner
|
||||
self._api_key = api_key
|
||||
self._client_config = client_config
|
||||
self._sql_host_override = sql_host_override
|
||||
self._oauth_config = oauth_config
|
||||
|
||||
@property
|
||||
def uri(self) -> str:
|
||||
"""The catalog's root namespace endpoint."""
|
||||
return self._inner.uri
|
||||
|
||||
def create_database(self, name: str, *, exist_ok: bool = False) -> DBConnection:
|
||||
"""Create a database, or open an existing one when ``exist_ok=True``."""
|
||||
inner = LOOP.run(self._inner.create_database(name, exist_ok=exist_ok))
|
||||
return self._wrap_database(name, inner)
|
||||
|
||||
def connect_database(self, name: str) -> DBConnection:
|
||||
"""Connect to an existing database by its logical name."""
|
||||
return self._wrap_database(name, LOOP.run(self._inner.connect_database(name)))
|
||||
|
||||
def _wrap_database(self, name: str, inner: AsyncConnection) -> DBConnection:
|
||||
return RemoteDBConnection._from_catalog(
|
||||
inner,
|
||||
name,
|
||||
self.uri,
|
||||
self._api_key,
|
||||
self._client_config,
|
||||
self._oauth_config,
|
||||
self._sql_host_override,
|
||||
)
|
||||
|
||||
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))
|
||||
|
||||
def drop_database(self, name: str, *, ignore_missing: bool = False) -> None:
|
||||
"""Drop an empty database. A nonempty database is an error."""
|
||||
LOOP.run(self._inner.drop_database(name, ignore_missing=ignore_missing))
|
||||
|
||||
|
||||
async def connect_catalog_async(
|
||||
endpoint: str,
|
||||
*,
|
||||
api_key: Optional[str] = None,
|
||||
client_config: Optional[Union[ClientConfig, dict[str, Any]]] = None,
|
||||
sql_host_override: Optional[str] = None,
|
||||
read_consistency_interval: Optional[timedelta] = None,
|
||||
oauth_config: Optional[OAuthConfig] = None,
|
||||
) -> AsyncCatalog:
|
||||
"""Connect to an HTTP(S) server's root catalog.
|
||||
|
||||
Root requests omit database-selection headers. API key, client configuration,
|
||||
OAuth, and table read consistency settings are inherited by opened databases.
|
||||
Database names containing slashes remain single logical names.
|
||||
Set ``sql_host_override`` to the SQL service endpoint to execute SQL through
|
||||
returned connections when the catalog endpoint uses HTTPS.
|
||||
"""
|
||||
if isinstance(client_config, dict):
|
||||
client_config = ClientConfig(**client_config)
|
||||
if client_config is None:
|
||||
client_config = ClientConfig()
|
||||
inner = await _lancedb.connect_catalog(
|
||||
endpoint,
|
||||
api_key=api_key,
|
||||
client_config=client_config,
|
||||
sql_host_override=sql_host_override,
|
||||
read_consistency_interval=(
|
||||
read_consistency_interval.total_seconds()
|
||||
if read_consistency_interval is not None
|
||||
else None
|
||||
),
|
||||
oauth_config=oauth_config,
|
||||
)
|
||||
return AsyncCatalog(inner)
|
||||
|
||||
|
||||
def connect_catalog(
|
||||
endpoint: str,
|
||||
*,
|
||||
api_key: Optional[str] = None,
|
||||
client_config: Optional[Union[ClientConfig, dict[str, Any]]] = None,
|
||||
sql_host_override: Optional[str] = None,
|
||||
read_consistency_interval: Optional[timedelta] = None,
|
||||
oauth_config: Optional[OAuthConfig] = None,
|
||||
) -> Catalog:
|
||||
"""Connect synchronously to an HTTP(S) server's root catalog.
|
||||
|
||||
See [connect_catalog_async][lancedb.connect_catalog_async] for options.
|
||||
Local filesystem and object-store catalogs are not supported.
|
||||
"""
|
||||
return Catalog(
|
||||
LOOP.run(
|
||||
connect_catalog_async(
|
||||
endpoint,
|
||||
api_key=api_key,
|
||||
client_config=client_config,
|
||||
sql_host_override=sql_host_override,
|
||||
read_consistency_interval=read_consistency_interval,
|
||||
oauth_config=oauth_config,
|
||||
)
|
||||
),
|
||||
api_key=api_key,
|
||||
client_config=client_config,
|
||||
sql_host_override=sql_host_override,
|
||||
oauth_config=oauth_config,
|
||||
)
|
||||
@@ -2,6 +2,7 @@
|
||||
# SPDX-FileCopyrightText: Copyright The LanceDB Authors
|
||||
|
||||
|
||||
from dataclasses import replace
|
||||
from datetime import timedelta
|
||||
import json
|
||||
import logging
|
||||
@@ -187,11 +188,51 @@ class RemoteDBConnection(DBConnection):
|
||||
)
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def _from_catalog(
|
||||
cls,
|
||||
inner,
|
||||
name,
|
||||
endpoint,
|
||||
api_key,
|
||||
client_config,
|
||||
oauth_config,
|
||||
sql_host_override,
|
||||
):
|
||||
config = (
|
||||
ClientConfig(**client_config)
|
||||
if isinstance(client_config, dict)
|
||||
else (client_config or ClientConfig())
|
||||
)
|
||||
headers = {
|
||||
key: value
|
||||
for key, value in (config.extra_headers or {}).items()
|
||||
if key.lower() not in ("x-lancedb-database", "x-lancedb-database-prefix")
|
||||
}
|
||||
headers["x-lancedb-database"] = name
|
||||
result = cls.__new__(cls)
|
||||
result.db_url = inner.uri
|
||||
result.db_name = name
|
||||
result.api_key = api_key or ""
|
||||
result.region = "us-east-1"
|
||||
result.host_override = endpoint
|
||||
result.sql_host_override = sql_host_override
|
||||
result.storage_options = None
|
||||
result.client_config = replace(config, extra_headers=headers)
|
||||
result._catalog_oauth = oauth_config is not None
|
||||
result._conn = inner
|
||||
return result
|
||||
|
||||
def __repr__(self) -> str:
|
||||
return f"RemoteConnect(name={self.db_name})"
|
||||
|
||||
@override
|
||||
def serialize(self) -> str:
|
||||
if getattr(self, "_catalog_oauth", False):
|
||||
raise ValueError(
|
||||
"Cannot serialize a catalog connection using OAuth; "
|
||||
"provide a worker-side connection factory"
|
||||
)
|
||||
return json.dumps(
|
||||
{
|
||||
"connection_type": "remote",
|
||||
|
||||
@@ -0,0 +1,140 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright The LanceDB Authors
|
||||
|
||||
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
|
||||
import json
|
||||
from threading import Thread
|
||||
|
||||
import pytest
|
||||
|
||||
import lancedb
|
||||
from lancedb.db import AsyncConnection, DBConnection
|
||||
from lancedb.remote.errors import HttpError
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def catalog_server():
|
||||
requests = []
|
||||
responses = []
|
||||
|
||||
class Handler(BaseHTTPRequestHandler):
|
||||
def handle_request(self):
|
||||
body = self.rfile.read(int(self.headers.get("Content-Length", "0")))
|
||||
requests.append(
|
||||
(
|
||||
self.path,
|
||||
dict(self.headers.items()),
|
||||
json.loads(body) if body else None,
|
||||
)
|
||||
)
|
||||
status, response = responses.pop(0)
|
||||
self.send_response(status)
|
||||
self.send_header("Content-Type", "application/json")
|
||||
self.end_headers()
|
||||
if status != 204:
|
||||
self.wfile.write(json.dumps(response).encode())
|
||||
|
||||
do_GET = handle_request
|
||||
do_POST = handle_request
|
||||
|
||||
with ThreadingHTTPServer(("127.0.0.1", 0), Handler) as server:
|
||||
thread = Thread(target=server.serve_forever, daemon=True)
|
||||
thread.start()
|
||||
try:
|
||||
yield f"http://127.0.0.1:{server.server_port}", requests, responses
|
||||
finally:
|
||||
server.shutdown()
|
||||
thread.join()
|
||||
|
||||
|
||||
def test_catalog_sync_scope_and_serialization(catalog_server):
|
||||
endpoint, requests, responses = catalog_server
|
||||
responses.extend(
|
||||
[
|
||||
(204, None),
|
||||
(200, {}),
|
||||
(200, {"tables": []}),
|
||||
(200, {"tables": []}),
|
||||
(200, {"namespaces": ["team/search"], "page_token": "next"}),
|
||||
(204, None),
|
||||
]
|
||||
)
|
||||
catalog = lancedb.connect_catalog(
|
||||
endpoint,
|
||||
api_key="secret",
|
||||
sql_host_override="invalid://localhost",
|
||||
client_config={
|
||||
"extra_headers": {
|
||||
"X-LanceDB-Database": "wrong",
|
||||
"X-LanceDB-Database-Prefix": "wrong",
|
||||
}
|
||||
},
|
||||
)
|
||||
assert isinstance(catalog, lancedb.Catalog)
|
||||
assert catalog.uri == endpoint
|
||||
db = catalog.create_database("team/search", exist_ok=True)
|
||||
assert isinstance(db, DBConnection)
|
||||
with pytest.raises(ValueError, match="sql_host_override must use"):
|
||||
db.execute_query_async("SELECT 1")
|
||||
db = catalog.connect_database("team/search")
|
||||
assert isinstance(db, DBConnection)
|
||||
assert db.table_names() == []
|
||||
restored = lancedb.deserialize_conn(db.serialize())
|
||||
assert restored.sql_host_override == "invalid://localhost"
|
||||
for connection in (db, restored):
|
||||
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")
|
||||
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/%24/list?limit=1&page_token=a%2Fb"
|
||||
assert requests[5][2] == {"mode": "Skip", "behavior": "Restrict"}
|
||||
for i, (_, headers, _) in enumerate(requests):
|
||||
headers = {key.lower(): value for key, value in headers.items()}
|
||||
assert headers.get("x-lancedb-database") == (
|
||||
"team/search" if i in (2, 3) else None
|
||||
)
|
||||
assert "x-lancedb-database-prefix" not in headers
|
||||
assert headers["x-api-key"] == "secret"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_catalog_async_and_errors(catalog_server):
|
||||
endpoint, requests, responses = catalog_server
|
||||
responses.extend(
|
||||
[
|
||||
(200, {}),
|
||||
(200, {"tables": []}),
|
||||
(404, {"error": "missing"}),
|
||||
(409, {"error": "exists"}),
|
||||
(400, {"error": "not empty"}),
|
||||
(404, {"error": "missing"}),
|
||||
]
|
||||
)
|
||||
catalog = await lancedb.connect_catalog_async(
|
||||
endpoint, sql_host_override="invalid://localhost"
|
||||
)
|
||||
assert isinstance(catalog, lancedb.AsyncCatalog)
|
||||
db = await catalog.connect_database("analytics")
|
||||
assert isinstance(db, AsyncConnection)
|
||||
with pytest.raises(ValueError, match="sql_host_override must use"):
|
||||
await db.execute_query_async("SELECT 1")
|
||||
assert await db.table_names() == []
|
||||
with pytest.raises(ValueError, match="missing"):
|
||||
await catalog.connect_database("missing")
|
||||
with pytest.raises(ValueError, match="exists"):
|
||||
await catalog.create_database("exists")
|
||||
with pytest.raises(HttpError):
|
||||
await catalog.drop_database("full")
|
||||
await catalog.drop_database("missing", ignore_missing=True)
|
||||
assert requests[4][2] == {"mode": "Fail", "behavior": "Restrict"}
|
||||
|
||||
|
||||
@pytest.mark.parametrize("endpoint", ["/tmp/catalog", "s3://bucket", "db://db"])
|
||||
def test_catalog_requires_remote_endpoint(endpoint):
|
||||
with pytest.raises(ValueError, match="endpoint"):
|
||||
lancedb.connect_catalog(endpoint)
|
||||
@@ -0,0 +1,125 @@
|
||||
// SPDX-License-Identifier: Apache-2.0
|
||||
// SPDX-FileCopyrightText: Copyright The LanceDB Authors
|
||||
|
||||
use std::time::Duration;
|
||||
|
||||
use lancedb::catalog::{
|
||||
CatalogConnection, CreateDatabaseRequest, DropDatabaseRequest, ListDatabasesRequest,
|
||||
};
|
||||
use pyo3::exceptions::PyValueError;
|
||||
use pyo3::{Bound, PyAny, PyRef, PyResult, Python, pyclass, pyfunction, pymethods};
|
||||
|
||||
use crate::connection::{Connection, PyClientConfig};
|
||||
use crate::error::PythonErrorExt;
|
||||
use crate::runtime::future_into_py;
|
||||
|
||||
#[pyclass]
|
||||
pub struct Catalog {
|
||||
inner: CatalogConnection,
|
||||
}
|
||||
|
||||
#[pymethods]
|
||||
impl Catalog {
|
||||
#[getter]
|
||||
fn uri(&self) -> &str {
|
||||
self.inner.uri()
|
||||
}
|
||||
|
||||
#[pyo3(signature = (name, *, exist_ok=false))]
|
||||
fn create_database<'py>(
|
||||
self_: PyRef<'py, Self>,
|
||||
name: String,
|
||||
exist_ok: bool,
|
||||
) -> PyResult<Bound<'py, PyAny>> {
|
||||
let inner = self_.inner.clone();
|
||||
future_into_py(self_.py(), async move {
|
||||
inner
|
||||
.create_database(CreateDatabaseRequest::new(name).exist_ok(exist_ok))
|
||||
.await
|
||||
.map(Connection::new)
|
||||
.infer_error()
|
||||
})
|
||||
}
|
||||
|
||||
fn connect_database<'py>(self_: PyRef<'py, Self>, name: String) -> PyResult<Bound<'py, PyAny>> {
|
||||
let inner = self_.inner.clone();
|
||||
future_into_py(self_.py(), async move {
|
||||
inner
|
||||
.connect_database(name)
|
||||
.await
|
||||
.map(Connection::new)
|
||||
.infer_error()
|
||||
})
|
||||
}
|
||||
|
||||
#[pyo3(signature = (name, *, ignore_missing=false))]
|
||||
fn drop_database<'py>(
|
||||
self_: PyRef<'py, Self>,
|
||||
name: String,
|
||||
ignore_missing: bool,
|
||||
) -> PyResult<Bound<'py, PyAny>> {
|
||||
let inner = self_.inner.clone();
|
||||
future_into_py(self_.py(), async move {
|
||||
inner
|
||||
.drop_database(DropDatabaseRequest::new(name).ignore_missing(ignore_missing))
|
||||
.await
|
||||
.infer_error()
|
||||
})
|
||||
}
|
||||
|
||||
#[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>> {
|
||||
let inner = self_.inner.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))
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
#[pyfunction]
|
||||
#[pyo3(signature = (endpoint, *, api_key=None, client_config=None, sql_host_override=None, read_consistency_interval=None, oauth_config=None))]
|
||||
pub fn connect_catalog(
|
||||
py: Python<'_>,
|
||||
endpoint: String,
|
||||
api_key: Option<String>,
|
||||
client_config: Option<PyClientConfig>,
|
||||
sql_host_override: Option<String>,
|
||||
read_consistency_interval: Option<f64>,
|
||||
oauth_config: Option<crate::oauth::PyOAuthConfig>,
|
||||
) -> PyResult<Bound<'_, PyAny>> {
|
||||
let interval = read_consistency_interval
|
||||
.map(Duration::try_from_secs_f64)
|
||||
.transpose()
|
||||
.map_err(|err| {
|
||||
PyValueError::new_err(format!("Invalid read consistency interval: {err}"))
|
||||
})?;
|
||||
future_into_py(py, async move {
|
||||
let mut builder = lancedb::connect_catalog(endpoint);
|
||||
if let Some(api_key) = api_key {
|
||||
builder = builder.api_key(api_key);
|
||||
}
|
||||
if let Some(config) = client_config {
|
||||
builder = builder.client_config(config.into());
|
||||
}
|
||||
if let Some(endpoint) = sql_host_override {
|
||||
builder = builder.sql_host_override(endpoint);
|
||||
}
|
||||
if let Some(interval) = interval {
|
||||
builder = builder.read_consistency_interval(interval);
|
||||
}
|
||||
if let Some(config) = oauth_config {
|
||||
builder = builder.oauth_config(config.try_into().infer_error()?);
|
||||
}
|
||||
Ok(Catalog {
|
||||
inner: builder.execute().await.infer_error()?,
|
||||
})
|
||||
})
|
||||
}
|
||||
@@ -31,6 +31,8 @@ impl<T> PythonErrorExt<T> for std::result::Result<T, LanceError> {
|
||||
| LanceError::TableNotFound { .. }
|
||||
| LanceError::NotAMaterializedView { .. }
|
||||
| LanceError::Schema { .. }
|
||||
| LanceError::DatabaseNotFound { .. }
|
||||
| LanceError::DatabaseAlreadyExists { .. }
|
||||
| LanceError::TableAlreadyExists { .. } => self.value_error(),
|
||||
LanceError::CreateDir { .. } => self.os_error(),
|
||||
LanceError::ObjectStore { .. } => Err(PyIOError::new_err(err.to_string())),
|
||||
|
||||
@@ -21,6 +21,7 @@ use table::{
|
||||
};
|
||||
|
||||
pub mod arrow;
|
||||
pub mod catalog;
|
||||
pub mod connection;
|
||||
pub mod error;
|
||||
pub mod expr;
|
||||
@@ -61,6 +62,8 @@ pub fn _lancedb(py: Python, m: &Bound<'_, PyModule>) -> PyResult<()> {
|
||||
.write_style("LANCEDB_LOG_STYLE");
|
||||
env_logger::init_from_env(env);
|
||||
m.add_class::<Connection>()?;
|
||||
m.add_class::<catalog::Catalog>()?;
|
||||
m.add_function(wrap_pyfunction!(catalog::connect_catalog, m)?)?;
|
||||
m.add_class::<Session>()?;
|
||||
m.add_class::<Table>()?;
|
||||
m.add_class::<crate::oauth::PyOAuthSession>()?;
|
||||
|
||||
@@ -0,0 +1,259 @@
|
||||
// SPDX-License-Identifier: Apache-2.0
|
||||
// SPDX-FileCopyrightText: Copyright The LanceDB Authors
|
||||
|
||||
//! Catalogs manage databases. A remote catalog is the root namespace of a server.
|
||||
//!
|
||||
//! ```
|
||||
//! # #[cfg(feature = "remote")]
|
||||
//! # async fn example() -> lancedb::Result<()> {
|
||||
//! let catalog = lancedb::connect_catalog("https://my-server.example")
|
||||
//! .api_key("my-api-key")
|
||||
//! .execute().await?;
|
||||
//! let database = catalog.create_database("analytics").await?;
|
||||
//! # Ok(())
|
||||
//! # }
|
||||
//! ```
|
||||
|
||||
use std::fmt;
|
||||
use std::sync::Arc;
|
||||
|
||||
use crate::Result;
|
||||
use crate::connection::Connection;
|
||||
use crate::database::Database;
|
||||
use crate::embeddings::{EmbeddingRegistry, MemoryRegistry};
|
||||
|
||||
/// Options for creating a database. By default an existing name is an error.
|
||||
#[derive(Clone, Debug)]
|
||||
#[non_exhaustive]
|
||||
pub struct CreateDatabaseRequest {
|
||||
/// Logical database name, including any literal slashes.
|
||||
pub name: String,
|
||||
/// Open the existing database if it is already registered.
|
||||
pub exist_ok: bool,
|
||||
}
|
||||
|
||||
impl CreateDatabaseRequest {
|
||||
/// Initialize a request with the default behavior.
|
||||
pub fn new(name: impl Into<String>) -> Self {
|
||||
Self {
|
||||
name: name.into(),
|
||||
exist_ok: false,
|
||||
}
|
||||
}
|
||||
|
||||
/// Open an existing database instead of failing if its name already exists.
|
||||
pub fn exist_ok(mut self, value: bool) -> Self {
|
||||
self.exist_ok = value;
|
||||
self
|
||||
}
|
||||
}
|
||||
|
||||
impl<T: Into<String>> From<T> for CreateDatabaseRequest {
|
||||
fn from(name: T) -> Self {
|
||||
Self::new(name)
|
||||
}
|
||||
}
|
||||
|
||||
/// Options for restricted database deletion. Tables must be removed first.
|
||||
#[derive(Clone, Debug)]
|
||||
#[non_exhaustive]
|
||||
pub struct DropDatabaseRequest {
|
||||
/// Logical database name, including any literal slashes.
|
||||
pub name: String,
|
||||
/// Succeed if the database is absent.
|
||||
pub ignore_missing: bool,
|
||||
}
|
||||
|
||||
impl DropDatabaseRequest {
|
||||
/// Initialize a request with the default behavior.
|
||||
pub fn new(name: impl Into<String>) -> Self {
|
||||
Self {
|
||||
name: name.into(),
|
||||
ignore_missing: false,
|
||||
}
|
||||
}
|
||||
|
||||
/// Succeed when the database does not exist.
|
||||
pub fn ignore_missing(mut self, value: bool) -> Self {
|
||||
self.ignore_missing = value;
|
||||
self
|
||||
}
|
||||
}
|
||||
|
||||
impl<T: Into<String>> From<T> for DropDatabaseRequest {
|
||||
fn from(name: T) -> Self {
|
||||
Self::new(name)
|
||||
}
|
||||
}
|
||||
|
||||
/// Pagination options for listing databases.
|
||||
#[derive(Clone, Debug, Default)]
|
||||
#[non_exhaustive]
|
||||
pub struct ListDatabasesRequest {
|
||||
/// Maximum number of names to return. None uses the server default.
|
||||
pub limit: Option<u32>,
|
||||
/// Opaque continuation token from a previous response. None starts a listing.
|
||||
pub page_token: Option<String>,
|
||||
}
|
||||
|
||||
impl ListDatabasesRequest {
|
||||
/// Set the maximum page size (1 through 2147483647).
|
||||
pub fn limit(mut self, limit: u32) -> Self {
|
||||
self.limit = Some(limit);
|
||||
self
|
||||
}
|
||||
/// Resume a listing from an opaque continuation token.
|
||||
pub fn page_token(mut self, token: impl Into<String>) -> Self {
|
||||
self.page_token = Some(token.into());
|
||||
self
|
||||
}
|
||||
}
|
||||
|
||||
/// One page of database names, relative to the catalog.
|
||||
#[derive(Clone, Debug, Default)]
|
||||
pub struct ListDatabasesResponse {
|
||||
/// Logical database names on this page.
|
||||
pub databases: Vec<String>,
|
||||
/// None indicates the end of the listing.
|
||||
pub page_token: Option<String>,
|
||||
}
|
||||
|
||||
/// A backend that manages databases. Implementations own database lifecycle semantics.
|
||||
#[async_trait::async_trait]
|
||||
pub trait Catalog: Send + Sync + std::fmt::Debug + 'static {
|
||||
/// Catalog endpoint or location.
|
||||
fn uri(&self) -> &str;
|
||||
/// Create a database, or open it when `exist_ok` permits.
|
||||
async fn create_database(&self, request: CreateDatabaseRequest) -> Result<Arc<dyn Database>>;
|
||||
/// Drop an empty database. This must not cascade to tables.
|
||||
async fn drop_database(&self, request: DropDatabaseRequest) -> Result<()>;
|
||||
/// List a page of database names.
|
||||
async fn list_databases(&self, request: ListDatabasesRequest) -> Result<ListDatabasesResponse>;
|
||||
/// Connect to an existing database by its logical name.
|
||||
async fn connect_database(&self, name: &str) -> Result<Arc<dyn Database>>;
|
||||
}
|
||||
|
||||
/// A catalog connection that returns ordinary LanceDB database connections.
|
||||
#[derive(Clone)]
|
||||
pub struct CatalogConnection {
|
||||
catalog: Arc<dyn Catalog>,
|
||||
embedding_registry: Arc<dyn EmbeddingRegistry>,
|
||||
}
|
||||
|
||||
impl fmt::Debug for CatalogConnection {
|
||||
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
|
||||
f.debug_struct("CatalogConnection")
|
||||
.field("uri", &self.uri())
|
||||
.finish_non_exhaustive()
|
||||
}
|
||||
}
|
||||
|
||||
impl CatalogConnection {
|
||||
/// Wrap a catalog implementation using the default in-memory embedding registry.
|
||||
pub fn new(catalog: Arc<dyn Catalog>) -> Self {
|
||||
Self {
|
||||
catalog,
|
||||
embedding_registry: Arc::new(MemoryRegistry::new()),
|
||||
}
|
||||
}
|
||||
|
||||
/// Provide the registry used by databases opened through this connection.
|
||||
pub fn with_embedding_registry(mut self, registry: Arc<dyn EmbeddingRegistry>) -> Self {
|
||||
self.embedding_registry = registry;
|
||||
self
|
||||
}
|
||||
|
||||
/// The catalog endpoint or location.
|
||||
pub fn uri(&self) -> &str {
|
||||
self.catalog.uri()
|
||||
}
|
||||
/// Access the underlying backend.
|
||||
pub fn catalog(&self) -> &Arc<dyn Catalog> {
|
||||
&self.catalog
|
||||
}
|
||||
|
||||
/// Create a database; pass a name or [`CreateDatabaseRequest`] for additional options.
|
||||
pub async fn create_database(
|
||||
&self,
|
||||
request: impl Into<CreateDatabaseRequest>,
|
||||
) -> Result<Connection> {
|
||||
Ok(Connection::new(
|
||||
self.catalog.create_database(request.into()).await?,
|
||||
self.embedding_registry.clone(),
|
||||
))
|
||||
}
|
||||
|
||||
/// Drop an empty database; pass a name or [`DropDatabaseRequest`] for additional options.
|
||||
pub async fn drop_database(&self, request: impl Into<DropDatabaseRequest>) -> Result<()> {
|
||||
self.catalog.drop_database(request.into()).await
|
||||
}
|
||||
|
||||
/// List a page of databases. Pass [`ListDatabasesRequest::default`] for the first page.
|
||||
pub async fn list_databases(
|
||||
&self,
|
||||
request: ListDatabasesRequest,
|
||||
) -> Result<ListDatabasesResponse> {
|
||||
self.catalog.list_databases(request).await
|
||||
}
|
||||
|
||||
/// Connect to a database without creating it if it is missing.
|
||||
pub async fn connect_database(&self, name: impl AsRef<str>) -> Result<Connection> {
|
||||
Ok(Connection::new(
|
||||
self.catalog.connect_database(name.as_ref()).await?,
|
||||
self.embedding_registry.clone(),
|
||||
))
|
||||
}
|
||||
}
|
||||
|
||||
/// Configure a connection to a remote catalog.
|
||||
#[cfg(feature = "remote")]
|
||||
#[derive(Debug)]
|
||||
pub struct ConnectCatalogBuilder {
|
||||
endpoint: String,
|
||||
options: crate::remote::RemoteCatalogOptions,
|
||||
}
|
||||
|
||||
#[cfg(feature = "remote")]
|
||||
impl ConnectCatalogBuilder {
|
||||
/// Start configuring a remote HTTP(S) catalog connection.
|
||||
pub fn new(endpoint: impl Into<String>) -> Self {
|
||||
Self {
|
||||
endpoint: endpoint.into(),
|
||||
options: Default::default(),
|
||||
}
|
||||
}
|
||||
/// Authenticate with an API key.
|
||||
pub fn api_key(mut self, key: impl Into<String>) -> Self {
|
||||
self.options.api_key = Some(key.into());
|
||||
self
|
||||
}
|
||||
/// Configure headers, TLS, timeouts, and other shared client settings.
|
||||
pub fn client_config(mut self, config: crate::remote::ClientConfig) -> Self {
|
||||
self.options.client_config = config;
|
||||
self
|
||||
}
|
||||
/// Set the SQL service endpoint inherited by database connections.
|
||||
///
|
||||
/// Required to execute SQL when the catalog endpoint uses HTTPS. The SQL
|
||||
/// connection is initialized lazily, using the ordinary remote SQL client.
|
||||
pub fn sql_host_override(mut self, endpoint: impl Into<String>) -> Self {
|
||||
self.options.sql_host_override = Some(endpoint.into());
|
||||
self
|
||||
}
|
||||
/// Configure table read consistency for opened databases.
|
||||
pub fn read_consistency_interval(mut self, interval: std::time::Duration) -> Self {
|
||||
self.options.read_consistency_interval = Some(interval);
|
||||
self
|
||||
}
|
||||
/// Authenticate using OAuth; mutually exclusive with API keys and header providers.
|
||||
pub fn oauth_config(mut self, config: crate::remote::OAuthConfig) -> Self {
|
||||
self.options.oauth_config = Some(config);
|
||||
self
|
||||
}
|
||||
/// Connect to the server's root namespace. Database-scoped headers are omitted.
|
||||
pub async fn execute(self) -> Result<CatalogConnection> {
|
||||
Ok(CatalogConnection::new(Arc::new(
|
||||
crate::remote::RemoteCatalog::try_new(&self.endpoint, self.options)?,
|
||||
)))
|
||||
}
|
||||
}
|
||||
@@ -174,6 +174,7 @@
|
||||
|
||||
pub mod arrow;
|
||||
pub mod blob;
|
||||
pub mod catalog;
|
||||
pub mod connection;
|
||||
pub mod data;
|
||||
pub mod database;
|
||||
@@ -383,3 +384,9 @@ pub use lance_io::object_store::ObjectStoreRegistry;
|
||||
/// declaring their own (potentially mismatched) direct `datafusion` dependency.
|
||||
/// See <https://github.com/lancedb/lancedb/issues/3575>.
|
||||
pub use datafusion;
|
||||
|
||||
/// Connect to a remote catalog through its HTTP(S) root namespace endpoint.
|
||||
#[cfg(feature = "remote")]
|
||||
pub fn connect_catalog(endpoint: impl Into<String>) -> catalog::ConnectCatalogBuilder {
|
||||
catalog::ConnectCatalogBuilder::new(endpoint)
|
||||
}
|
||||
|
||||
@@ -6,6 +6,7 @@
|
||||
//! building client/server applications with LanceDB or as a client for some
|
||||
//! other custom LanceDB service.
|
||||
|
||||
pub mod catalog;
|
||||
pub(crate) mod client;
|
||||
pub(crate) mod db;
|
||||
pub(crate) mod job;
|
||||
@@ -36,3 +37,5 @@ pub use oauth::{
|
||||
AuthorizationCodeOptions, ClientAuthMethod, OAuthConfig, OAuthFlow, OAuthHeaderProvider,
|
||||
};
|
||||
pub use token_cache::{OAuthSession, SessionLogout, SessionStatus, TokenCacheOptions};
|
||||
|
||||
pub use catalog::{RemoteCatalog, RemoteCatalogOptions};
|
||||
|
||||
@@ -0,0 +1,598 @@
|
||||
// SPDX-License-Identifier: Apache-2.0
|
||||
// SPDX-FileCopyrightText: Copyright The LanceDB Authors
|
||||
|
||||
use std::collections::HashMap;
|
||||
use std::fmt;
|
||||
use std::sync::Arc;
|
||||
use std::time::Duration;
|
||||
|
||||
use async_trait::async_trait;
|
||||
use http::StatusCode;
|
||||
use lance_namespace::models::{
|
||||
CreateNamespaceRequest, DescribeNamespaceRequest, DropNamespaceRequest, ListNamespacesRequest,
|
||||
};
|
||||
|
||||
use super::db::RemoteDatabase;
|
||||
use super::{ClientConfig, HeaderProvider, OAuthConfig, OAuthHeaderProvider};
|
||||
use crate::catalog::{
|
||||
Catalog, CreateDatabaseRequest, DropDatabaseRequest, ListDatabasesRequest,
|
||||
ListDatabasesResponse,
|
||||
};
|
||||
use crate::database::Database;
|
||||
use crate::{Error, Result};
|
||||
|
||||
/// Authentication and client settings shared by a catalog and its databases.
|
||||
#[derive(Clone, Default)]
|
||||
#[non_exhaustive]
|
||||
pub struct RemoteCatalogOptions {
|
||||
/// Optional API key for catalog and database requests.
|
||||
pub api_key: Option<String>,
|
||||
/// Shared transport and authentication settings.
|
||||
pub client_config: ClientConfig,
|
||||
/// SQL service endpoint used by returned database connections.
|
||||
/// Required for SQL when the catalog endpoint uses HTTPS.
|
||||
pub sql_host_override: Option<String>,
|
||||
/// Read consistency interval for tables opened in returned databases.
|
||||
pub read_consistency_interval: Option<Duration>,
|
||||
/// OAuth authentication, mutually exclusive with an API key or header provider.
|
||||
pub oauth_config: Option<OAuthConfig>,
|
||||
}
|
||||
|
||||
impl std::fmt::Debug for RemoteCatalogOptions {
|
||||
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
|
||||
f.debug_struct("RemoteCatalogOptions")
|
||||
.field("read_consistency_interval", &self.read_consistency_interval)
|
||||
.finish_non_exhaustive()
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Clone)]
|
||||
pub(crate) struct ScopedHeaderProvider {
|
||||
pub provider: Option<Arc<dyn HeaderProvider>>,
|
||||
pub database: Option<String>,
|
||||
}
|
||||
|
||||
impl std::fmt::Debug for ScopedHeaderProvider {
|
||||
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
|
||||
f.debug_struct("ScopedHeaderProvider")
|
||||
.field("database", &self.database)
|
||||
.finish_non_exhaustive()
|
||||
}
|
||||
}
|
||||
|
||||
impl ScopedHeaderProvider {
|
||||
pub(crate) fn apply(&self, headers: &mut HashMap<String, String>) {
|
||||
headers.retain(|name, _| {
|
||||
!name.eq_ignore_ascii_case("x-lancedb-database")
|
||||
&& !name.eq_ignore_ascii_case("x-lancedb-database-prefix")
|
||||
});
|
||||
if let Some(database) = &self.database {
|
||||
headers.insert("x-lancedb-database".into(), database.clone());
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl HeaderProvider for ScopedHeaderProvider {
|
||||
async fn get_headers(&self) -> Result<HashMap<String, String>> {
|
||||
let mut headers = match &self.provider {
|
||||
Some(provider) => provider.get_headers().await?,
|
||||
None => HashMap::new(),
|
||||
};
|
||||
self.apply(&mut headers);
|
||||
Ok(headers)
|
||||
}
|
||||
}
|
||||
|
||||
/// A catalog backed by the server's root namespace APIs.
|
||||
///
|
||||
/// Database management requests omit database-selection headers. Opened database
|
||||
/// connections retain their own scope and authentication independently.
|
||||
pub struct RemoteCatalog {
|
||||
endpoint: String,
|
||||
root: RemoteDatabase,
|
||||
options: RemoteCatalogOptions,
|
||||
}
|
||||
|
||||
impl fmt::Debug for RemoteCatalog {
|
||||
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
|
||||
f.debug_struct("RemoteCatalog")
|
||||
.field("uri", &self.endpoint)
|
||||
.finish_non_exhaustive()
|
||||
}
|
||||
}
|
||||
|
||||
impl RemoteCatalog {
|
||||
/// Connect to an HTTP(S) root namespace endpoint.
|
||||
///
|
||||
/// ```
|
||||
/// # use lancedb::remote::{RemoteCatalog, RemoteCatalogOptions};
|
||||
/// # fn example() -> lancedb::Result<()> {
|
||||
/// let catalog = RemoteCatalog::try_new("https://my-server.example", RemoteCatalogOptions::default())?;
|
||||
/// # Ok(())
|
||||
/// # }
|
||||
/// ```
|
||||
pub fn try_new(endpoint: impl AsRef<str>, mut options: RemoteCatalogOptions) -> Result<Self> {
|
||||
let url = url::Url::parse(endpoint.as_ref()).map_err(|err| Error::InvalidInput {
|
||||
message: format!("Invalid catalog endpoint: {err}"),
|
||||
})?;
|
||||
if !matches!(url.scheme(), "http" | "https")
|
||||
|| url.host_str().is_none()
|
||||
|| !url.username().is_empty()
|
||||
|| url.password().is_some()
|
||||
|| url.query().is_some()
|
||||
|| url.fragment().is_some()
|
||||
{
|
||||
return Err(Error::InvalidInput { message: "Catalog endpoint must be an HTTP(S) URL without credentials, query, or fragment".into() });
|
||||
}
|
||||
if options
|
||||
.client_config
|
||||
.id_delimiter
|
||||
.as_ref()
|
||||
.is_some_and(|d| d.is_empty())
|
||||
{
|
||||
return Err(Error::InvalidInput {
|
||||
message: "Catalog identifier delimiter cannot be empty".into(),
|
||||
});
|
||||
}
|
||||
if let Some(oauth) = options.oauth_config.take() {
|
||||
if options.api_key.is_some() || options.client_config.header_provider.is_some() {
|
||||
return Err(Error::InvalidInput {
|
||||
message: "oauth_config cannot be combined with api_key or header_provider"
|
||||
.into(),
|
||||
});
|
||||
}
|
||||
options.client_config.header_provider =
|
||||
Some(Arc::new(OAuthHeaderProvider::new(oauth)?));
|
||||
}
|
||||
let endpoint = url.to_string().trim_end_matches('/').to_string();
|
||||
let root = RemoteDatabase::for_catalog(&endpoint, None, &options)?;
|
||||
Ok(Self {
|
||||
endpoint,
|
||||
root,
|
||||
options,
|
||||
})
|
||||
}
|
||||
|
||||
fn validate_name(&self, name: &str) -> Result<()> {
|
||||
let delimiter = self
|
||||
.options
|
||||
.client_config
|
||||
.id_delimiter
|
||||
.as_deref()
|
||||
.unwrap_or("$");
|
||||
if name.is_empty()
|
||||
|| name.trim() != name
|
||||
|| !name.is_ascii()
|
||||
|| name.chars().any(char::is_control)
|
||||
|| name.contains(delimiter)
|
||||
|| matches!(name, "." | "..")
|
||||
{
|
||||
return Err(Error::InvalidInput {
|
||||
message: format!(
|
||||
"Invalid database name '{name}': expected a nonempty ASCII name without surrounding whitespace, control characters, or namespace delimiter '{delimiter}'"
|
||||
),
|
||||
});
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn database(&self, name: &str) -> Result<Arc<dyn Database>> {
|
||||
Ok(Arc::new(RemoteDatabase::for_catalog(
|
||||
&self.endpoint,
|
||||
Some(name),
|
||||
&self.options,
|
||||
)?))
|
||||
}
|
||||
|
||||
fn map_missing(name: &str, err: Error) -> Error {
|
||||
match err {
|
||||
Error::Http {
|
||||
status_code: Some(StatusCode::NOT_FOUND),
|
||||
..
|
||||
} => Error::DatabaseNotFound { name: name.into() },
|
||||
err => err,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl Catalog for RemoteCatalog {
|
||||
fn uri(&self) -> &str {
|
||||
&self.endpoint
|
||||
}
|
||||
|
||||
async fn create_database(&self, request: CreateDatabaseRequest) -> Result<Arc<dyn Database>> {
|
||||
self.validate_name(&request.name)?;
|
||||
self.root
|
||||
.create_namespace(CreateNamespaceRequest {
|
||||
id: Some(vec![request.name.clone()]),
|
||||
mode: Some(
|
||||
if request.exist_ok {
|
||||
"ExistOk"
|
||||
} else {
|
||||
"Create"
|
||||
}
|
||||
.into(),
|
||||
),
|
||||
..Default::default()
|
||||
})
|
||||
.await
|
||||
.map_err(|err| match err {
|
||||
Error::Http {
|
||||
status_code: Some(StatusCode::CONFLICT),
|
||||
..
|
||||
} => Error::DatabaseAlreadyExists {
|
||||
name: request.name.clone(),
|
||||
},
|
||||
err => err,
|
||||
})?;
|
||||
self.database(&request.name)
|
||||
}
|
||||
|
||||
async fn drop_database(&self, request: DropDatabaseRequest) -> Result<()> {
|
||||
self.validate_name(&request.name)?;
|
||||
let result = self
|
||||
.root
|
||||
.drop_namespace(DropNamespaceRequest {
|
||||
id: Some(vec![request.name.clone()]),
|
||||
mode: Some(
|
||||
if request.ignore_missing {
|
||||
"Skip"
|
||||
} else {
|
||||
"Fail"
|
||||
}
|
||||
.into(),
|
||||
),
|
||||
behavior: Some("Restrict".into()),
|
||||
..Default::default()
|
||||
})
|
||||
.await;
|
||||
match result {
|
||||
Ok(_) => Ok(()),
|
||||
Err(Error::Http {
|
||||
status_code: Some(StatusCode::NOT_FOUND),
|
||||
..
|
||||
}) if request.ignore_missing => Ok(()),
|
||||
Err(err) => Err(Self::map_missing(&request.name, err)),
|
||||
}
|
||||
}
|
||||
|
||||
async fn list_databases(&self, request: ListDatabasesRequest) -> Result<ListDatabasesResponse> {
|
||||
let limit = request
|
||||
.limit
|
||||
.map(|limit| {
|
||||
i32::try_from(limit)
|
||||
.ok()
|
||||
.filter(|limit| *limit > 0)
|
||||
.ok_or_else(|| Error::InvalidInput {
|
||||
message: "Database list limit must be between 1 and 2147483647".into(),
|
||||
})
|
||||
})
|
||||
.transpose()?;
|
||||
let response = self
|
||||
.root
|
||||
.list_namespaces(ListNamespacesRequest {
|
||||
id: Some(vec![]),
|
||||
limit,
|
||||
page_token: request.page_token,
|
||||
..Default::default()
|
||||
})
|
||||
.await?;
|
||||
Ok(ListDatabasesResponse {
|
||||
databases: response.namespaces,
|
||||
page_token: response.page_token.filter(|token| !token.is_empty()),
|
||||
})
|
||||
}
|
||||
|
||||
async fn connect_database(&self, name: &str) -> Result<Arc<dyn Database>> {
|
||||
self.validate_name(name)?;
|
||||
self.root
|
||||
.describe_namespace(DescribeNamespaceRequest {
|
||||
id: Some(vec![name.into()]),
|
||||
..Default::default()
|
||||
})
|
||||
.await
|
||||
.map_err(|err| Self::map_missing(name, err))?;
|
||||
self.database(name)
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use crate::catalog::CatalogConnection;
|
||||
use serde_json::{Value, json};
|
||||
use tokio::io::{AsyncReadExt, AsyncWriteExt};
|
||||
use tokio::net::TcpListener;
|
||||
use tokio::task::JoinHandle;
|
||||
|
||||
#[derive(Debug)]
|
||||
struct Request {
|
||||
line: String,
|
||||
headers: HashMap<String, String>,
|
||||
body: Value,
|
||||
}
|
||||
|
||||
async fn server(responses: Vec<(u16, Value)>) -> (String, JoinHandle<Vec<Request>>) {
|
||||
let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
|
||||
let endpoint = format!("http://{}", listener.local_addr().unwrap());
|
||||
let task = tokio::spawn(async move {
|
||||
let mut requests = Vec::new();
|
||||
for (status, body) in responses {
|
||||
let (mut socket, _) = listener.accept().await.unwrap();
|
||||
let mut bytes = Vec::new();
|
||||
let header_end = loop {
|
||||
let mut buf = [0; 4096];
|
||||
let n = socket.read(&mut buf).await.unwrap();
|
||||
assert!(n > 0);
|
||||
bytes.extend_from_slice(&buf[..n]);
|
||||
if let Some(pos) = bytes.windows(4).position(|b| b == b"\r\n\r\n") {
|
||||
break pos + 4;
|
||||
}
|
||||
};
|
||||
let header = String::from_utf8(bytes[..header_end].to_vec()).unwrap();
|
||||
let mut lines = header.lines();
|
||||
let line = lines.next().unwrap().to_string();
|
||||
let headers: HashMap<_, _> = lines
|
||||
.filter_map(|line| line.split_once(':'))
|
||||
.map(|(name, value)| (name.to_ascii_lowercase(), value.trim().to_string()))
|
||||
.collect();
|
||||
let length: usize = headers
|
||||
.get("content-length")
|
||||
.map(|s| s.parse().unwrap())
|
||||
.unwrap_or(0);
|
||||
while bytes.len() < header_end + length {
|
||||
let mut buf = [0; 4096];
|
||||
let n = socket.read(&mut buf).await.unwrap();
|
||||
assert!(n > 0);
|
||||
bytes.extend_from_slice(&buf[..n]);
|
||||
}
|
||||
let request_body = if length == 0 {
|
||||
Value::Null
|
||||
} else {
|
||||
serde_json::from_slice(&bytes[header_end..header_end + length]).unwrap()
|
||||
};
|
||||
requests.push(Request {
|
||||
line,
|
||||
headers,
|
||||
body: request_body,
|
||||
});
|
||||
let body = if status == 204 {
|
||||
String::new()
|
||||
} else {
|
||||
body.to_string()
|
||||
};
|
||||
socket.write_all(format!("HTTP/1.1 {status} Response\r\nContent-Type: application/json\r\nContent-Length: {}\r\nConnection: close\r\n\r\n{body}", body.len()).as_bytes()).await.unwrap();
|
||||
}
|
||||
requests
|
||||
});
|
||||
(endpoint, task)
|
||||
}
|
||||
|
||||
#[derive(Debug)]
|
||||
struct AuthProvider;
|
||||
|
||||
#[async_trait]
|
||||
impl HeaderProvider for AuthProvider {
|
||||
async fn get_headers(&self) -> Result<HashMap<String, String>> {
|
||||
Ok(HashMap::from([
|
||||
("Authorization".into(), "Bearer refreshed".into()),
|
||||
("X-LanceDB-Database".into(), "wrong-dynamic".into()),
|
||||
("X-LanceDB-Database-Prefix".into(), "wrong-prefix".into()),
|
||||
]))
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn catalog_routes_root_and_independent_database_scopes() {
|
||||
let (endpoint, task) = server(vec![
|
||||
(
|
||||
200,
|
||||
json!({"namespaces": ["team/search"], "page_token": "next"}),
|
||||
),
|
||||
(204, Value::Null),
|
||||
(200, json!({"namespaces": []})),
|
||||
(200, json!({})),
|
||||
(200, json!({"namespaces": []})),
|
||||
(200, json!({"namespaces": []})),
|
||||
(200, json!({"namespaces": [], "page_token": ""})),
|
||||
(204, Value::Null),
|
||||
])
|
||||
.await;
|
||||
let mut options = RemoteCatalogOptions {
|
||||
api_key: Some("test-key".into()),
|
||||
..Default::default()
|
||||
};
|
||||
options.client_config.extra_headers = HashMap::from([
|
||||
("x-lancedb-database".into(), "wrong-static".into()),
|
||||
(
|
||||
"x-lancedb-database-prefix".into(),
|
||||
"wrong-static-prefix".into(),
|
||||
),
|
||||
]);
|
||||
options.client_config.header_provider = Some(Arc::new(AuthProvider));
|
||||
let catalog = CatalogConnection::new(Arc::new(
|
||||
RemoteCatalog::try_new(&endpoint, options).unwrap(),
|
||||
));
|
||||
let page = catalog
|
||||
.list_databases(ListDatabasesRequest::default().limit(1).page_token("a/b"))
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(page.databases, ["team/search"]);
|
||||
assert_eq!(page.page_token.as_deref(), Some("next"));
|
||||
let first = catalog
|
||||
.create_database(CreateDatabaseRequest::new("team/search").exist_ok(true))
|
||||
.await
|
||||
.unwrap();
|
||||
first
|
||||
.database()
|
||||
.list_namespaces(ListNamespacesRequest::default())
|
||||
.await
|
||||
.unwrap();
|
||||
let second = catalog.connect_database("other").await.unwrap();
|
||||
second
|
||||
.database()
|
||||
.list_namespaces(ListNamespacesRequest::default())
|
||||
.await
|
||||
.unwrap();
|
||||
first
|
||||
.database()
|
||||
.list_namespaces(ListNamespacesRequest::default())
|
||||
.await
|
||||
.unwrap();
|
||||
assert!(
|
||||
catalog
|
||||
.list_databases(ListDatabasesRequest::default())
|
||||
.await
|
||||
.unwrap()
|
||||
.page_token
|
||||
.is_none()
|
||||
);
|
||||
catalog
|
||||
.drop_database(DropDatabaseRequest::new("team/search").ignore_missing(true))
|
||||
.await
|
||||
.unwrap();
|
||||
let requests = task.await.unwrap();
|
||||
for (i, request) in requests.iter().enumerate() {
|
||||
let database = match i {
|
||||
2 | 5 => Some("team/search"),
|
||||
4 => Some("other"),
|
||||
_ => None,
|
||||
};
|
||||
assert_eq!(
|
||||
request
|
||||
.headers
|
||||
.get("x-lancedb-database")
|
||||
.map(String::as_str),
|
||||
database
|
||||
);
|
||||
assert!(!request.headers.contains_key("x-lancedb-database-prefix"));
|
||||
assert_eq!(request.headers["authorization"], "Bearer refreshed");
|
||||
}
|
||||
assert_eq!(
|
||||
requests[0].line,
|
||||
"GET /v1/namespace/%24/list?limit=1&page_token=a%2Fb HTTP/1.1"
|
||||
);
|
||||
assert_eq!(
|
||||
requests[1].line,
|
||||
"POST /v1/namespace/team%2Fsearch/create HTTP/1.1"
|
||||
);
|
||||
assert_eq!(requests[1].body, json!({"mode": "ExistOk"}));
|
||||
assert_eq!(
|
||||
requests[3].line,
|
||||
"POST /v1/namespace/other/describe HTTP/1.1"
|
||||
);
|
||||
assert_eq!(
|
||||
requests[7].body,
|
||||
json!({"mode": "Skip", "behavior": "Restrict"})
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn catalog_preserves_errors_and_never_cascades() {
|
||||
let (endpoint, task) = server(vec![
|
||||
(404, json!({"error": "missing"})),
|
||||
(409, json!({"error": "exists"})),
|
||||
(400, json!({"error": "not empty"})),
|
||||
(404, json!({"error": "missing"})),
|
||||
(404, json!({"error": "missing"})),
|
||||
(401, json!({"error": "unauthorized"})),
|
||||
])
|
||||
.await;
|
||||
let catalog = RemoteCatalog::try_new(endpoint, RemoteCatalogOptions::default()).unwrap();
|
||||
assert!(matches!(
|
||||
catalog.connect_database("missing").await,
|
||||
Err(Error::DatabaseNotFound { .. })
|
||||
));
|
||||
assert!(matches!(
|
||||
catalog.create_database("exists".into()).await,
|
||||
Err(Error::DatabaseAlreadyExists { .. })
|
||||
));
|
||||
assert!(matches!(
|
||||
catalog.drop_database("full".into()).await,
|
||||
Err(Error::Http {
|
||||
status_code: Some(StatusCode::BAD_REQUEST),
|
||||
..
|
||||
})
|
||||
));
|
||||
assert!(matches!(
|
||||
catalog.drop_database("missing".into()).await,
|
||||
Err(Error::DatabaseNotFound { .. })
|
||||
));
|
||||
catalog
|
||||
.drop_database(DropDatabaseRequest::new("missing").ignore_missing(true))
|
||||
.await
|
||||
.unwrap();
|
||||
assert!(
|
||||
catalog
|
||||
.list_databases(ListDatabasesRequest::default())
|
||||
.await
|
||||
.is_err()
|
||||
);
|
||||
let requests = task.await.unwrap();
|
||||
assert_eq!(requests[1].body, json!({"mode": "Create"}));
|
||||
assert_eq!(
|
||||
requests[2].body,
|
||||
json!({"mode": "Fail", "behavior": "Restrict"})
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn catalog_validates_before_sending_requests() {
|
||||
for endpoint in [
|
||||
"/tmp/catalog",
|
||||
"s3://bucket",
|
||||
"db://database",
|
||||
"https://user:pass@host",
|
||||
"https://host?q=1",
|
||||
"https://host#fragment",
|
||||
] {
|
||||
assert!(RemoteCatalog::try_new(endpoint, RemoteCatalogOptions::default()).is_err());
|
||||
}
|
||||
let catalog =
|
||||
RemoteCatalog::try_new("http://127.0.0.1:1", RemoteCatalogOptions::default()).unwrap();
|
||||
for name in ["", "a$b", "\r\ninjected", "..", "café", " padded "] {
|
||||
assert!(matches!(
|
||||
catalog.create_database(name.into()).await,
|
||||
Err(Error::InvalidInput { .. })
|
||||
));
|
||||
assert!(matches!(
|
||||
catalog.connect_database(name).await,
|
||||
Err(Error::InvalidInput { .. })
|
||||
));
|
||||
assert!(matches!(
|
||||
catalog.drop_database(name.into()).await,
|
||||
Err(Error::InvalidInput { .. })
|
||||
));
|
||||
}
|
||||
for limit in [0, u32::MAX] {
|
||||
assert!(matches!(
|
||||
catalog
|
||||
.list_databases(ListDatabasesRequest::default().limit(limit))
|
||||
.await,
|
||||
Err(Error::InvalidInput { .. })
|
||||
));
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn catalog_debug_redacts_credentials() {
|
||||
let mut options = RemoteCatalogOptions {
|
||||
api_key: Some("catalog-secret-key".into()),
|
||||
..Default::default()
|
||||
};
|
||||
options
|
||||
.client_config
|
||||
.extra_headers
|
||||
.insert("authorization".into(), "Bearer catalog-secret-token".into());
|
||||
let catalog = RemoteCatalog::try_new("https://catalog.example", options).unwrap();
|
||||
let catalog_debug = format!("{catalog:?}");
|
||||
let connection = CatalogConnection::new(Arc::new(catalog));
|
||||
for debug in [catalog_debug, format!("{connection:?}")] {
|
||||
assert!(debug.contains("https://catalog.example"));
|
||||
assert!(!debug.contains("catalog-secret-key"));
|
||||
assert!(!debug.contains("catalog-secret-token"));
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -379,7 +379,11 @@ pub fn parse_db_url(db_url: &str) -> Result<ParsedDbUrl> {
|
||||
message: format!("Invalid database URL (missing host) '{}'", db_url),
|
||||
});
|
||||
}
|
||||
let db_name = parsed_url.host_str().unwrap().to_string();
|
||||
let db_name = urlencoding::decode(parsed_url.host_str().unwrap())
|
||||
.map_err(|err| Error::InvalidInput {
|
||||
message: format!("Invalid encoded database name: {err}"),
|
||||
})?
|
||||
.into_owned();
|
||||
let db_prefix = {
|
||||
let prefix = parsed_url.path().trim_start_matches('/');
|
||||
if prefix.is_empty() {
|
||||
@@ -1098,6 +1102,16 @@ mod tests {
|
||||
ENV_MUTEX.lock().unwrap_or_else(|e| e.into_inner())
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_parse_catalog_database_uri() {
|
||||
let parsed = parse_db_url("db://team%2Fsearch").unwrap();
|
||||
assert_eq!(parsed.db_name, "team/search");
|
||||
assert!(parsed.db_prefix.is_none());
|
||||
let parsed = parse_db_url("db://db/prefix").unwrap();
|
||||
assert_eq!(parsed.db_name, "db");
|
||||
assert_eq!(parsed.db_prefix.as_deref(), Some("prefix"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_timeout_config_default() {
|
||||
let config = TimeoutConfig::default();
|
||||
|
||||
@@ -289,6 +289,66 @@ impl RemoteDatabase {
|
||||
read_consistency_interval: Option<std::time::Duration>,
|
||||
) -> Result<Self> {
|
||||
let parsed = super::client::parse_db_url(uri)?;
|
||||
Self::try_new_with_identity(
|
||||
uri,
|
||||
api_key,
|
||||
region,
|
||||
host_overrides,
|
||||
client_config,
|
||||
options,
|
||||
read_consistency_interval,
|
||||
parsed,
|
||||
)
|
||||
}
|
||||
|
||||
pub(crate) fn for_catalog(
|
||||
endpoint: &str,
|
||||
name: Option<&str>,
|
||||
options: &super::catalog::RemoteCatalogOptions,
|
||||
) -> Result<Self> {
|
||||
let scope = super::catalog::ScopedHeaderProvider {
|
||||
provider: options.client_config.header_provider.clone(),
|
||||
database: name.map(str::to_string),
|
||||
};
|
||||
let mut config = options.client_config.clone();
|
||||
scope.apply(&mut config.extra_headers);
|
||||
config.header_provider = Some(Arc::new(scope));
|
||||
let uri = name
|
||||
.map(|name| format!("db://{}", urlencoding::encode(name)))
|
||||
.unwrap_or_else(|| endpoint.to_string());
|
||||
let mut db = Self::try_new_with_identity(
|
||||
&uri,
|
||||
options.api_key.as_deref().unwrap_or(""),
|
||||
"us-east-1",
|
||||
RemoteHostOverrides {
|
||||
rest: Some(endpoint.to_string()),
|
||||
sql: options.sql_host_override.clone(),
|
||||
},
|
||||
config,
|
||||
RemoteOptions::default(),
|
||||
options.read_consistency_interval,
|
||||
super::client::ParsedDbUrl {
|
||||
db_name: name.unwrap_or("").to_string(),
|
||||
db_prefix: None,
|
||||
},
|
||||
)?;
|
||||
if name.is_none() {
|
||||
db.sql_client = None;
|
||||
}
|
||||
Ok(db)
|
||||
}
|
||||
|
||||
#[allow(clippy::too_many_arguments)]
|
||||
fn try_new_with_identity(
|
||||
uri: &str,
|
||||
api_key: &str,
|
||||
region: &str,
|
||||
host_overrides: RemoteHostOverrides,
|
||||
client_config: ClientConfig,
|
||||
options: RemoteOptions,
|
||||
read_consistency_interval: Option<std::time::Duration>,
|
||||
parsed: super::client::ParsedDbUrl,
|
||||
) -> Result<Self> {
|
||||
let sql_client = SqlClient::new(
|
||||
parsed.db_name.clone(),
|
||||
parsed.db_prefix.clone(),
|
||||
@@ -301,7 +361,7 @@ impl RemoteDatabase {
|
||||
api_key,
|
||||
region,
|
||||
&parsed.db_name,
|
||||
host_overrides.rest.is_some(),
|
||||
host_overrides.rest.is_some() && !parsed.db_name.is_empty(),
|
||||
&options,
|
||||
parsed.db_prefix.as_deref(),
|
||||
&client_config,
|
||||
@@ -1220,6 +1280,7 @@ impl<S: HttpSend> Database for RemoteDatabase<S> {
|
||||
) -> Result<ListNamespacesResponse> {
|
||||
let namespace_parts = request.id.as_deref().unwrap_or(&[]);
|
||||
let namespace_id = build_namespace_identifier(namespace_parts, &self.client.id_delimiter);
|
||||
let namespace_id = urlencoding::encode(&namespace_id);
|
||||
let mut req = self
|
||||
.client
|
||||
.get(&format!("/v1/namespace/{}/list", namespace_id));
|
||||
@@ -1242,6 +1303,7 @@ impl<S: HttpSend> Database for RemoteDatabase<S> {
|
||||
) -> Result<CreateNamespaceResponse> {
|
||||
let namespace_parts = request.id.as_deref().unwrap_or(&[]);
|
||||
let namespace_id = build_namespace_identifier(namespace_parts, &self.client.id_delimiter);
|
||||
let namespace_id = urlencoding::encode(&namespace_id);
|
||||
let mut req = self
|
||||
.client
|
||||
.post(&format!("/v1/namespace/{}/create", namespace_id));
|
||||
@@ -1256,7 +1318,7 @@ impl<S: HttpSend> Database for RemoteDatabase<S> {
|
||||
}
|
||||
|
||||
let body = CreateNamespaceRequestBody {
|
||||
mode: request.mode.as_ref().map(|m| format!("{:?}", m)),
|
||||
mode: request.mode,
|
||||
properties: request.properties,
|
||||
};
|
||||
|
||||
@@ -1264,12 +1326,16 @@ impl<S: HttpSend> Database for RemoteDatabase<S> {
|
||||
let (request_id, resp) = self.client.send(req).await?;
|
||||
let resp = self.client.check_response(&request_id, resp).await?;
|
||||
|
||||
if resp.status() == StatusCode::NO_CONTENT {
|
||||
return Ok(CreateNamespaceResponse::default());
|
||||
}
|
||||
resp.json().await.err_to_http(request_id)
|
||||
}
|
||||
|
||||
async fn drop_namespace(&self, request: DropNamespaceRequest) -> Result<DropNamespaceResponse> {
|
||||
let namespace_parts = request.id.as_deref().unwrap_or(&[]);
|
||||
let namespace_id = build_namespace_identifier(namespace_parts, &self.client.id_delimiter);
|
||||
let namespace_id = urlencoding::encode(&namespace_id);
|
||||
let mut req = self
|
||||
.client
|
||||
.post(&format!("/v1/namespace/{}/drop", namespace_id));
|
||||
@@ -1284,14 +1350,17 @@ impl<S: HttpSend> Database for RemoteDatabase<S> {
|
||||
}
|
||||
|
||||
let body = DropNamespaceRequestBody {
|
||||
mode: request.mode.as_ref().map(|m| format!("{:?}", m)),
|
||||
behavior: request.behavior.as_ref().map(|b| format!("{:?}", b)),
|
||||
mode: request.mode,
|
||||
behavior: request.behavior,
|
||||
};
|
||||
|
||||
req = req.json(&body);
|
||||
let (request_id, resp) = self.client.send(req).await?;
|
||||
let resp = self.client.check_response(&request_id, resp).await?;
|
||||
|
||||
if resp.status() == StatusCode::NO_CONTENT {
|
||||
return Ok(DropNamespaceResponse::default());
|
||||
}
|
||||
resp.json().await.err_to_http(request_id)
|
||||
}
|
||||
|
||||
@@ -1301,6 +1370,7 @@ impl<S: HttpSend> Database for RemoteDatabase<S> {
|
||||
) -> Result<DescribeNamespaceResponse> {
|
||||
let namespace_parts = request.id.as_deref().unwrap_or(&[]);
|
||||
let namespace_id = build_namespace_identifier(namespace_parts, &self.client.id_delimiter);
|
||||
let namespace_id = urlencoding::encode(&namespace_id);
|
||||
let req = self
|
||||
.client
|
||||
.post(&format!("/v1/namespace/{}/describe", namespace_id))
|
||||
|
||||
@@ -18,7 +18,10 @@ use futures::{StreamExt, TryStreamExt};
|
||||
use tonic::{Request, Response, Status, Streaming};
|
||||
|
||||
use super::*;
|
||||
use crate::database::Database;
|
||||
use crate::remote::RemoteCatalogOptions;
|
||||
use crate::remote::client::HeaderProvider;
|
||||
use crate::remote::db::RemoteDatabase;
|
||||
|
||||
#[derive(Debug, Default)]
|
||||
struct DelayedHeaderProvider {
|
||||
@@ -173,7 +176,10 @@ impl FlightService for TestSqlService {
|
||||
namespace_path: header("namespace-path"),
|
||||
request_id: header("x-request-id"),
|
||||
api_key: header("x-api-key"),
|
||||
database_prefix: header("x-lancedb-database-prefix"),
|
||||
database_prefix: metadata
|
||||
.get("x-lancedb-database-prefix")
|
||||
.map(|value| value.to_str().unwrap().to_string())
|
||||
.unwrap_or_default(),
|
||||
});
|
||||
|
||||
let command = Any::decode(request.get_ref().cmd.as_ref())
|
||||
@@ -386,6 +392,60 @@ impl FlightService for TestSqlService {
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn catalog_connections_use_explicit_sql_endpoint_and_database_scope() {
|
||||
let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
|
||||
let address = listener.local_addr().unwrap();
|
||||
let incoming = futures::stream::try_unfold(listener, |listener| async {
|
||||
let (socket, _) = listener.accept().await?;
|
||||
Ok::<_, std::io::Error>(Some((socket, listener)))
|
||||
});
|
||||
let service = TestSqlService::default();
|
||||
let headers = service.headers.clone();
|
||||
let expected = service.result.clone();
|
||||
let (shutdown_tx, shutdown_rx) = tokio::sync::oneshot::channel();
|
||||
let server = tokio::spawn(
|
||||
tonic::transport::Server::builder()
|
||||
.add_service(FlightServiceServer::new(service))
|
||||
.serve_with_incoming_shutdown(incoming, async {
|
||||
let _ = shutdown_rx.await;
|
||||
}),
|
||||
);
|
||||
let options = RemoteCatalogOptions {
|
||||
api_key: Some("catalog-key".into()),
|
||||
sql_host_override: Some(format!("grpc://{address}")),
|
||||
..Default::default()
|
||||
};
|
||||
for name in ["analytics", "team/search"] {
|
||||
let database =
|
||||
RemoteDatabase::for_catalog("https://catalog.example", Some(name), &options).unwrap();
|
||||
let query = database
|
||||
.execute_query_async("SELECT 42", &[])
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(
|
||||
collect_result(&query).await.unwrap(),
|
||||
vec![expected.clone()]
|
||||
);
|
||||
}
|
||||
{
|
||||
let headers = headers.lock().unwrap();
|
||||
assert!(headers.iter().any(|header| header.database == "analytics"));
|
||||
assert!(
|
||||
headers
|
||||
.iter()
|
||||
.any(|header| header.database == "team/search")
|
||||
);
|
||||
for header in headers.iter() {
|
||||
assert_eq!(header.api_key, "catalog-key");
|
||||
assert_eq!(header.namespace_path, "public");
|
||||
assert!(header.database_prefix.is_empty());
|
||||
}
|
||||
}
|
||||
shutdown_tx.send(()).unwrap();
|
||||
server.await.unwrap().unwrap();
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn submits_polls_fetches_cancels_and_reuses_client() {
|
||||
let listener = std::net::TcpListener::bind("127.0.0.1:0").unwrap();
|
||||
|
||||
Reference in New Issue
Block a user