diff --git a/.github/workflows/nodejs.yml b/.github/workflows/nodejs.yml index 55c050a7d..98acccf02 100644 --- a/.github/workflows/nodejs.yml +++ b/.github/workflows/nodejs.yml @@ -207,4 +207,4 @@ jobs: pnpm tsc - name: Test run: | - pnpm test + pnpm test --runInBand diff --git a/docs/src/js/classes/Catalog.md b/docs/src/js/classes/Catalog.md new file mode 100644 index 000000000..deac5e3ce --- /dev/null +++ b/docs/src/js/classes/Catalog.md @@ -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 +``` + +Connect to an existing database by its logical name. + +#### Parameters + +* **name**: `string` + +#### Returns + +`Promise`<[`Connection`](Connection.md)> + +*** + +### createDatabase() + +```ts +createDatabase(name, options): Promise +``` + +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 +``` + +Drop an empty database. The server rejects nonempty databases. + +#### Parameters + +* **name**: `string` + +* **options** = `{}` + +* **options.ignoreMissing?**: `boolean` + +#### Returns + +`Promise`<`void`> + +*** + +### listDatabases() + +```ts +listDatabases(options): Promise +``` + +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)> diff --git a/docs/src/js/functions/connectCatalog.md b/docs/src/js/functions/connectCatalog.md new file mode 100644 index 000000000..2bdb3182f --- /dev/null +++ b/docs/src/js/functions/connectCatalog.md @@ -0,0 +1,32 @@ +[**@lancedb/lancedb**](../README.md) • **Docs** + +*** + +[@lancedb/lancedb](../globals.md) / connectCatalog + +# Function: connectCatalog() + +```ts +function connectCatalog(endpoint, options): Promise +``` + +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 }); +``` diff --git a/docs/src/js/globals.md b/docs/src/js/globals.md index 4e09853f6..92e7940f0 100644 --- a/docs/src/js/globals.md +++ b/docs/src/js/globals.md @@ -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) diff --git a/docs/src/js/interfaces/CatalogOptions.md b/docs/src/js/interfaces/CatalogOptions.md new file mode 100644 index 000000000..c66823f6c --- /dev/null +++ b/docs/src/js/interfaces/CatalogOptions.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 | Promise>; +``` + +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` diff --git a/docs/src/js/interfaces/ListDatabasesResponse.md b/docs/src/js/interfaces/ListDatabasesResponse.md new file mode 100644 index 000000000..a1a247746 --- /dev/null +++ b/docs/src/js/interfaces/ListDatabasesResponse.md @@ -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; +``` diff --git a/docs/src/python/python.md b/docs/src/python/python.md index d14e5cf94..96161fbac 100644 --- a/docs/src/python/python.md +++ b/docs/src/python/python.md @@ -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. diff --git a/nodejs/__test__/catalog.test.ts b/nodejs/__test__/catalog.test.ts new file mode 100644 index 000000000..d44103ab9 --- /dev/null +++ b/nodejs/__test__/catalog.test.ts @@ -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; +}; + +async function withCatalog( + responses: [number, unknown][], + callback: (catalog: Catalog, requests: RecordedRequest[]) => Promise, +) { + 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((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((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", + ); + }); +}); diff --git a/nodejs/lancedb/catalog.ts b/nodejs/lancedb/catalog.ts new file mode 100644 index 000000000..5cfe2c5ea --- /dev/null +++ b/nodejs/lancedb/catalog.ts @@ -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 { + oauthConfig?: OAuthConfig; + /** Called for each request to supply authentication headers. */ + headerProvider?: + | HeaderProvider + | (() => Record | Promise>); +} + +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 { + return new LocalConnection( + await this.inner.createDatabase(name, options.existOk), + ); + } + + /** Connect to an existing database by its logical name. */ + async connectDatabase(name: string): Promise { + 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 { + 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 { + 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 { + 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), + ); +} diff --git a/nodejs/lancedb/index.ts b/nodejs/lancedb/index.ts index 55078bba3..39ceec917 100644 --- a/nodejs/lancedb/index.ts +++ b/nodejs/lancedb/index.ts @@ -629,3 +629,10 @@ export async function connectNamespace( ); return new LocalConnection(nativeConn); } + +export { + Catalog, + CatalogOptions, + ListDatabasesResponse, + connectCatalog, +} from "./catalog"; diff --git a/nodejs/src/catalog.rs b/nodejs/src/catalog.rs new file mode 100644 index 000000000..c2288bc4a --- /dev/null +++ b/nodejs/src/catalog.rs @@ -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, + pub client_config: Option, + /// SQL service endpoint inherited by database connections. + pub sql_host_override: Option, + pub read_consistency_interval: Option, + pub oauth_config: Option, +} + +#[napi(object)] +pub struct ListDatabasesResponse { + pub databases: Vec, + pub page_token: Option, +} + +#[napi] +pub struct Catalog { + inner: CatalogConnection, +} + +#[napi] +impl Catalog { + #[napi(factory)] + pub async fn new( + endpoint: String, + options: CatalogOptions, + header_provider: Option<&JsHeaderProvider>, + ) -> Result { + 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, + ) -> Result { + 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 { + self.inner + .connect_database(name) + .await + .map(Connection::inner_new) + .default_error() + } + #[napi] + pub async fn drop_database(&self, name: String, ignore_missing: Option) -> 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, + page_token: Option, + ) -> Result { + 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, + }) + } +} diff --git a/nodejs/src/lib.rs b/nodejs/src/lib.rs index 288a2b925..1bb5f761d 100644 --- a/nodejs/src/lib.rs +++ b/nodejs/src/lib.rs @@ -11,6 +11,7 @@ use env_logger::Env; use napi_derive::*; mod blob; +mod catalog; mod connection; mod error; mod header; diff --git a/python/python/lancedb/__init__.py b/python/python/lancedb/__init__.py index cb9b57be3..19114e1a1 100644 --- a/python/python/lancedb/__init__.py +++ b/python/python/lancedb/__init__.py @@ -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", diff --git a/python/python/lancedb/_lancedb.pyi b/python/python/lancedb/_lancedb.pyi index e18657eb7..f267b3a55 100644 --- a/python/python/lancedb/_lancedb.pyi +++ b/python/python/lancedb/_lancedb.pyi @@ -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: ... diff --git a/python/python/lancedb/catalog.py b/python/python/lancedb/catalog.py new file mode 100644 index 000000000..4fd941ddd --- /dev/null +++ b/python/python/lancedb/catalog.py @@ -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, + ) diff --git a/python/python/lancedb/remote/db.py b/python/python/lancedb/remote/db.py index fb4e30fdf..aceebf987 100644 --- a/python/python/lancedb/remote/db.py +++ b/python/python/lancedb/remote/db.py @@ -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", diff --git a/python/python/tests/test_catalog.py b/python/python/tests/test_catalog.py new file mode 100644 index 000000000..63386a998 --- /dev/null +++ b/python/python/tests/test_catalog.py @@ -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) diff --git a/python/src/catalog.rs b/python/src/catalog.rs new file mode 100644 index 000000000..f61b409ff --- /dev/null +++ b/python/src/catalog.rs @@ -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> { + 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> { + 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> { + 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, + page_token: Option, + ) -> PyResult> { + 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, + client_config: Option, + sql_host_override: Option, + read_consistency_interval: Option, + oauth_config: Option, +) -> PyResult> { + 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()?, + }) + }) +} diff --git a/python/src/error.rs b/python/src/error.rs index b46fa6c83..869a5a247 100644 --- a/python/src/error.rs +++ b/python/src/error.rs @@ -31,6 +31,8 @@ impl PythonErrorExt for std::result::Result { | 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())), diff --git a/python/src/lib.rs b/python/src/lib.rs index 0bd7d5d6d..fd9a23798 100644 --- a/python/src/lib.rs +++ b/python/src/lib.rs @@ -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::()?; + m.add_class::()?; + m.add_function(wrap_pyfunction!(catalog::connect_catalog, m)?)?; m.add_class::()?; m.add_class::()?; m.add_class::()?; diff --git a/rust/lancedb/src/catalog.rs b/rust/lancedb/src/catalog.rs new file mode 100644 index 000000000..5eaa1d49e --- /dev/null +++ b/rust/lancedb/src/catalog.rs @@ -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) -> 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> From 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) -> 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> From 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, + /// Opaque continuation token from a previous response. None starts a listing. + pub page_token: Option, +} + +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) -> 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, + /// None indicates the end of the listing. + pub page_token: Option, +} + +/// 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>; + /// 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; + /// Connect to an existing database by its logical name. + async fn connect_database(&self, name: &str) -> Result>; +} + +/// A catalog connection that returns ordinary LanceDB database connections. +#[derive(Clone)] +pub struct CatalogConnection { + catalog: Arc, + embedding_registry: Arc, +} + +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) -> 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) -> 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 { + &self.catalog + } + + /// Create a database; pass a name or [`CreateDatabaseRequest`] for additional options. + pub async fn create_database( + &self, + request: impl Into, + ) -> Result { + 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) -> 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 { + 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) -> Result { + 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) -> Self { + Self { + endpoint: endpoint.into(), + options: Default::default(), + } + } + /// Authenticate with an API key. + pub fn api_key(mut self, key: impl Into) -> 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) -> 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 { + Ok(CatalogConnection::new(Arc::new( + crate::remote::RemoteCatalog::try_new(&self.endpoint, self.options)?, + ))) + } +} diff --git a/rust/lancedb/src/lib.rs b/rust/lancedb/src/lib.rs index 44c5dd616..b173a63cb 100644 --- a/rust/lancedb/src/lib.rs +++ b/rust/lancedb/src/lib.rs @@ -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 . 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) -> catalog::ConnectCatalogBuilder { + catalog::ConnectCatalogBuilder::new(endpoint) +} diff --git a/rust/lancedb/src/remote.rs b/rust/lancedb/src/remote.rs index 4441e01ee..bc534c0f4 100644 --- a/rust/lancedb/src/remote.rs +++ b/rust/lancedb/src/remote.rs @@ -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}; diff --git a/rust/lancedb/src/remote/catalog.rs b/rust/lancedb/src/remote/catalog.rs new file mode 100644 index 000000000..098a696a7 --- /dev/null +++ b/rust/lancedb/src/remote/catalog.rs @@ -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, + /// 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, + /// Read consistency interval for tables opened in returned databases. + pub read_consistency_interval: Option, + /// OAuth authentication, mutually exclusive with an API key or header provider. + pub oauth_config: Option, +} + +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>, + pub database: Option, +} + +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) { + 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> { + 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, mut options: RemoteCatalogOptions) -> Result { + 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> { + 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> { + 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 { + 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> { + 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, + body: Value, + } + + async fn server(responses: Vec<(u16, Value)>) -> (String, JoinHandle>) { + 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> { + 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")); + } + } +} diff --git a/rust/lancedb/src/remote/client.rs b/rust/lancedb/src/remote/client.rs index 6e4764ad6..55d272b51 100644 --- a/rust/lancedb/src/remote/client.rs +++ b/rust/lancedb/src/remote/client.rs @@ -379,7 +379,11 @@ pub fn parse_db_url(db_url: &str) -> Result { 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(); diff --git a/rust/lancedb/src/remote/db.rs b/rust/lancedb/src/remote/db.rs index c24ede4e1..e7ae78108 100644 --- a/rust/lancedb/src/remote/db.rs +++ b/rust/lancedb/src/remote/db.rs @@ -289,6 +289,66 @@ impl RemoteDatabase { read_consistency_interval: Option, ) -> Result { 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 { + 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, + parsed: super::client::ParsedDbUrl, + ) -> Result { 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 Database for RemoteDatabase { ) -> Result { 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 Database for RemoteDatabase { ) -> Result { 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 Database for RemoteDatabase { } let body = CreateNamespaceRequestBody { - mode: request.mode.as_ref().map(|m| format!("{:?}", m)), + mode: request.mode, properties: request.properties, }; @@ -1264,12 +1326,16 @@ impl Database for RemoteDatabase { 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 { 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 Database for RemoteDatabase { } 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 Database for RemoteDatabase { ) -> Result { 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)) diff --git a/rust/lancedb/src/remote/sql_test.rs b/rust/lancedb/src/remote/sql_test.rs index a05a11259..d20ba495f 100644 --- a/rust/lancedb/src/remote/sql_test.rs +++ b/rust/lancedb/src/remote/sql_test.rs @@ -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();