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:
@@ -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>()?;
|
||||
|
||||
Reference in New Issue
Block a user