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.

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:

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,
     )

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

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)])

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,
+    )

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<String>,
     read_consistency_interval: Option<f64>,
     oauth_config: Option<crate::oauth::PyOAuthConfig>,
+    storage_options: Option<HashMap<String, String>>,
 ) -> PyResult<Bound<'_, PyAny>> {
     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);
         }