From 7b2dfb9e361bddfbf4b3e16f2d950f377d5b5e33 Mon Sep 17 00:00:00 2001 From: Jack Ye Date: Thu, 12 Nov 2026 01:19:05 -0800 Subject: [PATCH] fix(python): align catalog and database connection inputs (#4983) Align Python catalog connection inputs with database connections, warn when region is ignored, and preserve remote storage settings in opened databases. --- docs/web/basics/catalog.mdx | 18 +++ python/python/lancedb/_lancedb.pyi | 1 + python/python/lancedb/catalog.py | 154 ++++++++++++++++++++++-- python/python/lancedb/remote/db.py | 3 +- python/python/tests/test_catalog.py | 176 +++++++++++++++++++++++++++- python/python/type_tests/connect.py | 20 +++- python/src/catalog.rs | 8 +- rust/lancedb/src/catalog.rs | 12 ++ rust/lancedb/src/remote/catalog.rs | 10 ++ rust/lancedb/src/remote/db.rs | 2 +- 10 files changed, 389 insertions(+), 15 deletions(-) diff --git a/docs/web/basics/catalog.mdx b/docs/web/basics/catalog.mdx index 1e95aeca4..519d86067 100644 --- a/docs/web/basics/catalog.mdx +++ b/docs/web/basics/catalog.mdx @@ -158,6 +158,24 @@ list, create and drop databases across the whole deployment. Its `connect_database(name)` and `create_database(name)` then return the same kind of database connection `lancedb.connect()` does. +In Python, `connect_catalog` and `connect_catalog_async` accept the same +connection settings as their database counterparts. Use `host_override` to +share settings that include a `db://` URI; the catalog connection targets +the root of that deployment and ignores the database name. An explicitly +supplied `region` is ignored with a warning, since the HTTP(S) endpoint +selects the deployment. Embedded-only session and namespace settings are +unsupported. + +```python +settings = { + "uri": "db://analytics", + "api_key": "secret", + "host_override": "https://my-server.example", +} +db = lancedb.connect(**settings) +catalog = lancedb.connect_catalog(**settings) +``` + This connection is for database administration, not table or query work: list, create and drop databases through it, then switch to a **database** connection for everything else -- tables, queries, and SQL all run within a diff --git a/python/python/lancedb/_lancedb.pyi b/python/python/lancedb/_lancedb.pyi index 521cc46bb..11798a520 100644 --- a/python/python/lancedb/_lancedb.pyi +++ b/python/python/lancedb/_lancedb.pyi @@ -1227,6 +1227,7 @@ async def connect_catalog( sql_host_override: Optional[str] = None, read_consistency_interval: Optional[float] = None, oauth_config: Optional[Any] = None, + storage_options: Optional[dict[str, str]] = None, ) -> Catalog: ... class NameListing: diff --git a/python/python/lancedb/catalog.py b/python/python/lancedb/catalog.py index 338be98b8..5b3b0509c 100644 --- a/python/python/lancedb/catalog.py +++ b/python/python/lancedb/catalog.py @@ -3,13 +3,17 @@ """Remote catalogs manage databases.""" -from functools import cached_property +import os +import warnings +from concurrent.futures import ThreadPoolExecutor from datetime import timedelta +from functools import cached_property from typing import Any, AsyncIterator, Iterator, Optional, Union from . import _lancedb from .background_loop import LOOP from .authz import AsyncAuthorization, Authorization +from .common import URI, sanitize_uri from .db import AsyncConnection, DBConnection from .remote import ClientConfig, OAuthConfig from .remote.db import RemoteDBConnection @@ -163,12 +167,14 @@ class Catalog: client_config=None, sql_host_override: Optional[str] = None, oauth_config: Optional[OAuthConfig] = None, + storage_options: Optional[dict[str, str]] = 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 + self._storage_options = storage_options @cached_property def authz(self) -> Authorization: @@ -198,6 +204,7 @@ class Catalog: self._client_config, self._oauth_config, self._sql_host_override, + self._storage_options, ) def list_databases( @@ -227,28 +234,71 @@ class Catalog: async def connect_catalog_async( - endpoint: str, + uri: Optional[URI] = None, *, api_key: Optional[str] = None, - client_config: Optional[Union[ClientConfig, dict[str, Any]]] = None, + region: Optional[str] = None, + host_override: Optional[str] = None, sql_host_override: Optional[str] = None, read_consistency_interval: Optional[timedelta] = None, + client_config: Optional[Union[ClientConfig, dict[str, Any]]] = None, + storage_options: Optional[dict[str, str]] = None, + session: Optional[_lancedb.Session] = None, + manifest_enabled: bool = False, + namespace_client_properties: Optional[dict[str, str]] = None, oauth_config: Optional[OAuthConfig] = None, ) -> AsyncCatalog: - """Connect to an HTTP(S) server's root catalog. + """Connect to a server's root catalog using database connection settings. + + ``host_override`` selects the HTTP(S) catalog endpoint when provided; otherwise + ``uri`` must be an HTTP(S) endpoint. A ``db://`` URI can be shared with + [connect_async][lancedb.connect_async] when ``host_override`` is set: the URI's + database is ignored and requests target the root catalog. + + ``region`` is ignored with a warning when supplied. Catalog connections use + the explicit endpoint rather than a region-derived Cloud (legacy) endpoint. + ``api_key`` defaults to the ``LANCEDB_API_KEY`` environment variable. + Embedded-only ``session``, ``manifest_enabled`` and + ``namespace_client_properties`` are not supported. Root requests omit database-selection headers. API key, client configuration, - OAuth, and table read consistency settings are inherited by opened databases. + OAuth, storage options, 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 region is not None: + warnings.warn( + "region is ignored by connect_catalog_async; the catalog endpoint " + "is selected by host_override or uri.", + UserWarning, + stacklevel=2, + ) + for name, value in ( + ("session", session), + ("manifest_enabled", manifest_enabled or None), + ("namespace_client_properties", namespace_client_properties), + ): + if value is not None: + raise NotImplementedError( + f"{name} is not supported for catalog connections" + ) + endpoint = host_override if host_override is not None else uri + if endpoint is None: + raise ValueError("uri or host_override is required for catalog connections") if isinstance(client_config, dict): client_config = ClientConfig(**client_config) if client_config is None: client_config = ClientConfig() + if ( + api_key is None + and oauth_config is None + and client_config.header_provider is None + ): + api_key = os.environ.get("LANCEDB_API_KEY") inner = await _lancedb.connect_catalog( - endpoint, + sanitize_uri(endpoint), api_key=api_key, client_config=client_config, sql_host_override=sql_host_override, @@ -258,32 +308,115 @@ async def connect_catalog_async( else None ), oauth_config=oauth_config, + storage_options=storage_options, ) return AsyncCatalog(inner) def connect_catalog( - endpoint: str, + uri: Optional[URI] = None, *, api_key: Optional[str] = None, - client_config: Optional[Union[ClientConfig, dict[str, Any]]] = None, + region: Optional[str] = None, + host_override: Optional[str] = None, sql_host_override: Optional[str] = None, read_consistency_interval: Optional[timedelta] = None, + request_thread_pool: Optional[Union[int, ThreadPoolExecutor]] = None, + client_config: Optional[Union[ClientConfig, dict[str, Any]]] = None, + storage_options: Optional[dict[str, str]] = None, + session: Optional[_lancedb.Session] = None, + manifest_enabled: bool = False, + namespace_client_impl: Optional[str] = None, + namespace_client_properties: Optional[dict[str, str]] = None, + namespace_client_pushdown_operations: Optional[list[str]] = None, oauth_config: Optional[OAuthConfig] = None, + **kwargs: Any, ) -> Catalog: - """Connect synchronously to an HTTP(S) server's root catalog. + """Connect synchronously to a root catalog using database connection settings. See [connect_catalog_async][lancedb.connect_catalog_async] for options. + The inputs match [connect][lancedb.connect]; ``region`` is ignored with a + warning when supplied. Namespace client options and embedded session settings + are not supported. ``request_thread_pool`` is deprecated and unused, as in + ``connect``. Local filesystem and object-store catalogs are not supported. + + Examples + -------- + ```python + settings = {"api_key": "secret", "host_override": "https://my-server.example"} + catalog = lancedb.connect_catalog("db://analytics", **settings) + db = lancedb.connect("db://analytics", **settings) + ``` """ + unknown_kwargs = { + key: value + for key, value in kwargs.items() + if key not in ("connection_timeout", "read_timeout") + } + if unknown_kwargs: + raise ValueError(f"Unknown keyword arguments: {unknown_kwargs}") + for name, value in ( + ("namespace_client_impl", namespace_client_impl), + ("namespace_client_pushdown_operations", namespace_client_pushdown_operations), + ): + if value is not None: + raise NotImplementedError( + f"{name} is not supported for catalog connections" + ) + if region is not None: + warnings.warn( + "region is ignored by connect_catalog; the catalog endpoint " + "is selected by host_override or uri.", + UserWarning, + stacklevel=2, + ) + if request_thread_pool is not None: + warnings.warn( + "request_thread_pool is no longer used and will be removed in " + "a future release.", + DeprecationWarning, + stacklevel=2, + ) + if isinstance(client_config, dict): + client_config = ClientConfig(**client_config) + if client_config is None: + client_config = ClientConfig() + for name, field in ( + ("connection_timeout", "connect_timeout"), + ("read_timeout", "read_timeout"), + ): + if kwargs.get(name) is not None: + warnings.warn( + f"{name} is deprecated and will be removed in a future release. " + f"Please use client_config.timeout_config.{field} instead.", + DeprecationWarning, + stacklevel=2, + ) + setattr( + client_config.timeout_config, + field, + timedelta(seconds=kwargs[name]), + ) + if ( + api_key is None + and oauth_config is None + and client_config.header_provider is None + ): + api_key = os.environ.get("LANCEDB_API_KEY") return Catalog( LOOP.run( connect_catalog_async( - endpoint, + uri, api_key=api_key, + host_override=host_override, client_config=client_config, sql_host_override=sql_host_override, read_consistency_interval=read_consistency_interval, + storage_options=storage_options, + session=session, + manifest_enabled=manifest_enabled, + namespace_client_properties=namespace_client_properties, oauth_config=oauth_config, ) ), @@ -291,4 +424,5 @@ def connect_catalog( client_config=client_config, sql_host_override=sql_host_override, oauth_config=oauth_config, + storage_options=storage_options, ) diff --git a/python/python/lancedb/remote/db.py b/python/python/lancedb/remote/db.py index 21ad37b36..4219ca40b 100644 --- a/python/python/lancedb/remote/db.py +++ b/python/python/lancedb/remote/db.py @@ -214,6 +214,7 @@ class RemoteDBConnection(DBConnection): client_config, oauth_config, sql_host_override, + storage_options, ): config = ( ClientConfig(**client_config) @@ -233,7 +234,7 @@ class RemoteDBConnection(DBConnection): result.region = "us-east-1" result.host_override = endpoint result.sql_host_override = sql_host_override - result.storage_options = None + result.storage_options = storage_options result.client_config = replace(config, extra_headers=headers) result._catalog_oauth = oauth_config is not None result._conn = inner diff --git a/python/python/tests/test_catalog.py b/python/python/tests/test_catalog.py index 95b03c6e7..edac21bc1 100644 --- a/python/python/tests/test_catalog.py +++ b/python/python/tests/test_catalog.py @@ -2,7 +2,9 @@ # SPDX-FileCopyrightText: Copyright The LanceDB Authors from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer +from datetime import timedelta import json +import warnings import pytest @@ -62,6 +64,7 @@ def test_catalog_sync_scope_and_serialization(catalog_server): endpoint, api_key="secret", sql_host_override="invalid://localhost", + storage_options={"azure_storage_account_name": "catalog-account"}, client_config={ "extra_headers": { "X-LanceDB-Database": "wrong", @@ -80,6 +83,7 @@ def test_catalog_sync_scope_and_serialization(catalog_server): assert db.table_names() == [] restored = lancedb.deserialize_conn(db.serialize()) assert restored.sql_host_override == "invalid://localhost" + assert restored.storage_options == {"azure_storage_account_name": "catalog-account"} for connection in (db, restored): with pytest.raises(ValueError, match="sql_host_override must use"): connection.execute_query_async("SELECT 1") @@ -98,6 +102,7 @@ def test_catalog_sync_scope_and_serialization(catalog_server): ) assert "x-lancedb-database-prefix" not in headers assert headers["x-api-key"] == "secret" + assert headers["x-azure-storage-account-name"] == "catalog-account" @pytest.mark.asyncio @@ -114,7 +119,9 @@ async def test_catalog_async_and_errors(catalog_server): ] ) catalog = await lancedb.connect_catalog_async( - endpoint, sql_host_override="invalid://localhost" + endpoint, + sql_host_override="invalid://localhost", + storage_options={"account_name": "async-account"}, ) assert isinstance(catalog, lancedb.AsyncCatalog) db = await catalog.connect_database("analytics") @@ -130,6 +137,9 @@ async def test_catalog_async_and_errors(catalog_server): await catalog.drop_database("full") await catalog.drop_database("missing", ignore_missing=True) assert requests[4][2] == {"mode": "Fail", "behavior": "Restrict"} + for _, headers, _ in requests: + headers = {key.lower(): value for key, value in headers.items()} + assert headers["x-azure-storage-account-name"] == "async-account" @pytest.mark.parametrize("endpoint", ["/tmp/catalog", "s3://bucket", "db://db"]) @@ -138,6 +148,170 @@ def test_catalog_requires_remote_endpoint(endpoint): lancedb.connect_catalog(endpoint) +@pytest.mark.parametrize("asynchronous", [False, True]) +@pytest.mark.parametrize("region", ["us-east-1", "not-a-region"]) +@pytest.mark.asyncio +async def test_catalog_accepts_database_connection_settings( + catalog_server, asynchronous, region +): + endpoint, requests, responses = catalog_server + responses.extend([(200, {"namespaces": []}), (200, {"tables": []})]) + settings = { + "uri": "db://analytics", + "api_key": "shared-secret", + "region": region, + "host_override": endpoint, + "sql_host_override": "invalid://localhost", + "read_consistency_interval": timedelta(seconds=5), + "client_config": {"extra_headers": {"x-shared-setting": "shared"}}, + "storage_options": {"azure_storage_account_name": "shared-account"}, + "session": None, + "manifest_enabled": False, + "namespace_client_properties": None, + } + if not asynchronous: + settings.update( + request_thread_pool=None, + namespace_client_impl=None, + namespace_client_pushdown_operations=None, + ) + with pytest.warns( + UserWarning, match="region is ignored by connect_catalog" + ) as caught: + catalog = ( + await lancedb.connect_catalog_async(**settings) + if asynchronous + else lancedb.connect_catalog(**settings) + ) + assert len(caught) == 1 + assert catalog.uri == endpoint + assert ( + [name async for name in catalog.list_databases()] + if asynchronous + else list(catalog.list_databases()) + ) == [] + db = ( + await lancedb.connect_async(**settings) + if asynchronous + else lancedb.connect(**settings) + ) + assert (await db.table_names() if asynchronous else db.table_names()) == [] + root_headers = {key.lower(): value for key, value in requests[0][1].items()} + db_headers = {key.lower(): value for key, value in requests[1][1].items()} + for headers in (root_headers, db_headers): + assert headers["x-api-key"] == "shared-secret" + assert headers["x-shared-setting"] == "shared" + assert headers["x-azure-storage-account-name"] == "shared-account" + assert "x-lancedb-database" not in root_headers + assert db_headers["x-lancedb-database"] == "analytics" + + +@pytest.mark.parametrize("asynchronous", [False, True]) +@pytest.mark.parametrize("host_only", [False, True]) +@pytest.mark.asyncio +async def test_catalog_default_region_is_silent_and_api_key_uses_environment( + catalog_server, monkeypatch, asynchronous, host_only +): + endpoint, requests, responses = catalog_server + monkeypatch.setenv("LANCEDB_API_KEY", "environment-secret") + responses.extend([(200, {}), (200, {"tables": []})]) + settings = {"host_override": endpoint} if host_only else {"uri": endpoint} + with warnings.catch_warnings(record=True) as caught: + warnings.simplefilter("always") + catalog = ( + await lancedb.connect_catalog_async(**settings) + if asynchronous + else lancedb.connect_catalog(**settings) + ) + assert caught == [] + assert catalog.uri == endpoint + db = ( + await catalog.connect_database("analytics") + if asynchronous + else catalog.connect_database("analytics") + ) + if not asynchronous: + db = lancedb.deserialize_conn(db.serialize()) + assert (await db.table_names() if asynchronous else db.table_names()) == [] + for _, headers, _ in requests: + headers = {key.lower(): value for key, value in headers.items()} + assert headers["x-api-key"] == "environment-secret" + + +@pytest.mark.parametrize("asynchronous", [False, True]) +@pytest.mark.asyncio +async def test_catalog_host_override_selects_endpoint(catalog_server, asynchronous): + endpoint, _, _ = catalog_server + settings = {"uri": "https://unused.example", "host_override": endpoint} + catalog = ( + await lancedb.connect_catalog_async(**settings) + if asynchronous + else lancedb.connect_catalog(**settings) + ) + assert catalog.uri == endpoint + + +@pytest.mark.parametrize("asynchronous", [False, True]) +@pytest.mark.parametrize( + "options, message", + [ + ({}, "uri or host_override is required"), + ({"session": object()}, "session"), + ({"manifest_enabled": True}, "manifest_enabled"), + ({"namespace_client_properties": {}}, "namespace_client_properties"), + ], +) +@pytest.mark.asyncio +async def test_catalog_rejects_unsupported_connection_settings( + catalog_server, asynchronous, options, message +): + _, requests, _ = catalog_server + error = NotImplementedError if options else ValueError + with pytest.raises(error, match=message): + if asynchronous: + await lancedb.connect_catalog_async(**options) + else: + lancedb.connect_catalog(**options) + assert requests == [] + + +@pytest.mark.parametrize( + "options, error, message", + [ + ( + {"namespace_client_impl": "dir"}, + NotImplementedError, + "namespace_client_impl", + ), + ( + {"namespace_client_pushdown_operations": []}, + NotImplementedError, + "namespace_client_pushdown_operations", + ), + ({"unknown_option": True}, ValueError, "Unknown keyword arguments"), + ], +) +def test_catalog_rejects_unsupported_sync_settings(options, error, message): + with pytest.raises(error, match=message): + lancedb.connect_catalog(**options) + + +def test_catalog_accepts_deprecated_sync_settings(catalog_server): + endpoint, _, responses = catalog_server + responses.append((200, {})) + with pytest.warns(DeprecationWarning) as caught: + catalog = lancedb.connect_catalog( + endpoint, + request_thread_pool=2, + connection_timeout=5, + read_timeout=10, + ) + assert len(caught) == 3 + db = catalog.connect_database("analytics") + assert db.client_config.timeout_config.connect_timeout == timedelta(seconds=5) + assert db.client_config.timeout_config.read_timeout == timedelta(seconds=10) + + @pytest.mark.parametrize("asynchronous", [False, True]) @pytest.mark.asyncio @pytest.mark.parametrize("page_token,page_limit", [(None, None), ("start/token", 2)]) diff --git a/python/python/type_tests/connect.py b/python/python/type_tests/connect.py index eb2cba37c..6da32d2ec 100644 --- a/python/python/type_tests/connect.py +++ b/python/python/type_tests/connect.py @@ -4,12 +4,30 @@ from typing import assert_type import lancedb -from lancedb import AsyncConnection, DBConnection +from lancedb import AsyncCatalog, AsyncConnection, Catalog, DBConnection def check_connect_type() -> None: assert_type(lancedb.connect("memory://"), DBConnection) + assert_type( + lancedb.connect_catalog( + uri="db://analytics", + api_key="secret", + region="us-west-2", + host_override="https://my-server.example", + ), + Catalog, + ) async def check_connect_async_type() -> None: assert_type(await lancedb.connect_async("memory://"), AsyncConnection) + assert_type( + await lancedb.connect_catalog_async( + uri="db://analytics", + api_key="secret", + region="us-west-2", + host_override="https://my-server.example", + ), + AsyncCatalog, + ) diff --git a/python/src/catalog.rs b/python/src/catalog.rs index fb2e7e8a0..456d62a3c 100644 --- a/python/src/catalog.rs +++ b/python/src/catalog.rs @@ -1,6 +1,7 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright The LanceDB Authors +use std::collections::HashMap; use std::sync::{Arc, Mutex}; use std::time::Duration; @@ -127,7 +128,8 @@ impl DatabaseNames { } #[pyfunction] -#[pyo3(signature = (endpoint, *, api_key=None, client_config=None, sql_host_override=None, read_consistency_interval=None, oauth_config=None))] +#[allow(clippy::too_many_arguments)] +#[pyo3(signature = (endpoint, *, api_key=None, client_config=None, sql_host_override=None, read_consistency_interval=None, oauth_config=None, storage_options=None))] pub fn connect_catalog( py: Python<'_>, endpoint: String, @@ -136,6 +138,7 @@ pub fn connect_catalog( sql_host_override: Option, read_consistency_interval: Option, oauth_config: Option, + storage_options: Option>, ) -> PyResult> { let interval = read_consistency_interval .map(Duration::try_from_secs_f64) @@ -151,6 +154,9 @@ pub fn connect_catalog( if let Some(config) = client_config { builder = builder.client_config(config.into()); } + if let Some(options) = storage_options { + builder = builder.storage_options(options); + } if let Some(endpoint) = sql_host_override { builder = builder.sql_host_override(endpoint); } diff --git a/rust/lancedb/src/catalog.rs b/rust/lancedb/src/catalog.rs index df80b6de6..dcea8fd90 100644 --- a/rust/lancedb/src/catalog.rs +++ b/rust/lancedb/src/catalog.rs @@ -349,6 +349,18 @@ impl ConnectCatalogBuilder { self.options.client_config = config; self } + /// Set storage options inherited by the catalog's remote database connections. + pub fn storage_options( + mut self, + options: impl IntoIterator, impl Into)>, + ) -> Self { + self.options.storage_options.extend( + options + .into_iter() + .map(|(key, value)| (key.into(), value.into())), + ); + self + } /// Set the SQL service endpoint inherited by database connections. /// /// Required to execute SQL when the catalog endpoint uses HTTPS. The SQL diff --git a/rust/lancedb/src/remote/catalog.rs b/rust/lancedb/src/remote/catalog.rs index 95facd04b..c620cf272 100644 --- a/rust/lancedb/src/remote/catalog.rs +++ b/rust/lancedb/src/remote/catalog.rs @@ -30,6 +30,8 @@ pub struct RemoteCatalogOptions { pub api_key: Option, /// Shared transport and authentication settings. pub client_config: ClientConfig, + /// Storage options supported by remote database connections. + pub storage_options: HashMap, /// SQL service endpoint used by returned database connections. /// Required for SQL when the catalog endpoint uses HTTPS. pub sql_host_override: Option, @@ -391,6 +393,10 @@ mod tests { .await; let mut options = RemoteCatalogOptions { api_key: Some("test-key".into()), + storage_options: HashMap::from([( + "azure_storage_account_name".into(), + "catalog-account".into(), + )]), ..Default::default() }; options.client_config.extra_headers = HashMap::from([ @@ -460,6 +466,10 @@ mod tests { ); assert!(!request.headers.contains_key("x-lancedb-database-prefix")); assert_eq!(request.headers["authorization"], "Bearer refreshed"); + assert_eq!( + request.headers["x-azure-storage-account-name"], + "catalog-account" + ); } assert_eq!( requests[0].line, diff --git a/rust/lancedb/src/remote/db.rs b/rust/lancedb/src/remote/db.rs index a9d55b511..549baf345 100644 --- a/rust/lancedb/src/remote/db.rs +++ b/rust/lancedb/src/remote/db.rs @@ -385,7 +385,7 @@ impl RemoteDatabase { sql: options.sql_host_override.clone(), }, config, - RemoteOptions::default(), + StorageOptions(options.storage_options.clone()).into(), options.read_consistency_interval, super::client::ParsedDbUrl { db_name: name.unwrap_or("").to_string(),