mirror of
https://github.com/lancedb/lancedb.git
synced 2026-11-03 01:06:52 +00:00
Merge remote-tracking branch 'origin/main' into xuanwo/oss-2012-json-fts
# Conflicts: # rust/lancedb/src/remote/table.rs # rust/lancedb/src/table/create_index.rs
This commit is contained in:
+2
-1
@@ -1,6 +1,6 @@
|
||||
[package]
|
||||
name = "lancedb-python"
|
||||
version = "0.39.0-beta.7"
|
||||
version = "0.40.0-beta.3"
|
||||
publish = false
|
||||
edition.workspace = true
|
||||
description = "Python bindings for LanceDB"
|
||||
@@ -15,6 +15,7 @@ name = "_lancedb"
|
||||
crate-type = ["cdylib"]
|
||||
|
||||
[dependencies]
|
||||
arc-swap = "1.9"
|
||||
arrow = { workspace = true, features = ["pyarrow"] }
|
||||
async-trait.workspace = true
|
||||
bytes.workspace = true
|
||||
|
||||
@@ -63,7 +63,7 @@ tests = [
|
||||
"polars>=0.19, <=1.32.3",
|
||||
"pyarrow<25",
|
||||
"pyarrow-stubs>=16.0",
|
||||
"pylance==9.0.0rc1",
|
||||
"pylance==9.0.0",
|
||||
"requests>=2.31.0",
|
||||
"datafusion>=54,<55",
|
||||
"opentelemetry-sdk>=1.30.0",
|
||||
@@ -126,6 +126,7 @@ include = [
|
||||
"python/lancedb/rerankers/util.py",
|
||||
"python/lancedb/rerankers/__init__.py",
|
||||
"python/lancedb/rerankers/voyageai.py",
|
||||
"python/lancedb/rerankers/typesafe.py",
|
||||
"python/lancedb/rerankers/jinaai.py",
|
||||
"python/lancedb/rerankers/openai.py",
|
||||
"python/lancedb/rerankers/cross_encoder.py",
|
||||
|
||||
@@ -12,6 +12,9 @@ __version__ = importlib.metadata.version("lancedb")
|
||||
|
||||
from ._lancedb import connect as lancedb_connect
|
||||
from ._lancedb import FtsToken
|
||||
from ._lancedb import FunctionErrorFragment as FunctionErrorFragment
|
||||
from ._lancedb import FunctionErrorRecord as FunctionErrorRecord
|
||||
from ._lancedb import FunctionErrors as FunctionErrors
|
||||
from ._lancedb import LsmWriteSpec
|
||||
from ._lancedb import tokenize as _tokenize
|
||||
from .common import URI, sanitize_uri
|
||||
@@ -37,6 +40,8 @@ from .functions import (
|
||||
UdfDefinition as UdfDefinition,
|
||||
udf as udf,
|
||||
)
|
||||
from .secrets import EnvVarSecret as EnvVarSecret
|
||||
from .secrets import SecretInfo as SecretInfo
|
||||
from .materialized_view import (
|
||||
AsyncMaterializedView,
|
||||
MaterializedView,
|
||||
@@ -52,6 +57,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 +574,11 @@ async def connect_async(
|
||||
|
||||
|
||||
__all__ = [
|
||||
"Catalog",
|
||||
"AsyncCatalog",
|
||||
"ListDatabasesResponse",
|
||||
"connect_catalog",
|
||||
"connect_catalog_async",
|
||||
"AsyncMaterializedView",
|
||||
"MaterializedView",
|
||||
"MaterializedViewDefinition",
|
||||
|
||||
@@ -153,6 +153,21 @@ class Connection(object):
|
||||
async def get_function(self, name: str, version: str) -> str: ...
|
||||
async def list_functions(self) -> List[str]: ...
|
||||
async def drop_function(self, name: str, version: str) -> bool: ...
|
||||
async def create_secret(
|
||||
self, name: str, value: str, namespace_path: Optional[List[str]] = None
|
||||
) -> None: ...
|
||||
async def alter_secret(
|
||||
self, name: str, value: str, namespace_path: Optional[List[str]] = None
|
||||
) -> None: ...
|
||||
async def list_secrets(
|
||||
self, namespace_path: Optional[List[str]] = None
|
||||
) -> List[str]: ...
|
||||
async def drop_secret(
|
||||
self, name: str, namespace_path: Optional[List[str]] = None
|
||||
) -> None: ...
|
||||
async def describe_secret(
|
||||
self, name: str, namespace_path: Optional[List[str]] = None
|
||||
) -> Tuple[str, int, int]: ...
|
||||
async def list_jobs(self) -> List[JobInfo]: ...
|
||||
async def cancel_job(self, job_id: str) -> bool: ...
|
||||
async def execute_query_async(
|
||||
@@ -211,8 +226,24 @@ class Connection(object):
|
||||
projections: Optional[List[Tuple[str, str]]] = None,
|
||||
filter: Optional[str] = None,
|
||||
limit: Optional[int] = None,
|
||||
with_no_data: bool = False,
|
||||
) -> Table: ...
|
||||
async def create_materialized_view_async(
|
||||
self,
|
||||
name: str,
|
||||
source: str,
|
||||
projections: Optional[List[Tuple[str, str]]] = None,
|
||||
filter: Optional[str] = None,
|
||||
limit: Optional[int] = None,
|
||||
with_no_data: bool = False,
|
||||
) -> Job: ...
|
||||
async def list_materialized_views(self) -> List[str]: ...
|
||||
async def drop_materialized_view(
|
||||
self, name: str, namespace_path: Optional[List[str]] = None
|
||||
) -> None: ...
|
||||
async def drop_materialized_view_async(
|
||||
self, name: str, namespace_path: Optional[List[str]] = None
|
||||
) -> Job: ...
|
||||
async def drop_table(
|
||||
self, name: str, namespace_path: Optional[List[str]] = None
|
||||
) -> None: ...
|
||||
@@ -267,6 +298,37 @@ class JobInfo:
|
||||
@property
|
||||
def created_at_millis(self) -> int: ...
|
||||
|
||||
class SessionStatus:
|
||||
@property
|
||||
def resource(self) -> Optional[str]: ...
|
||||
@property
|
||||
def audience(self) -> Optional[str]: ...
|
||||
@property
|
||||
def refreshable(self) -> bool: ...
|
||||
@property
|
||||
def issuer_url(self) -> str: ...
|
||||
@property
|
||||
def client_id(self) -> str: ...
|
||||
@property
|
||||
def scopes(self) -> List[str]: ...
|
||||
@property
|
||||
def flow(self) -> str: ...
|
||||
@property
|
||||
def obtained_at(self) -> Optional[int]: ...
|
||||
def __repr__(self) -> str: ...
|
||||
|
||||
class SessionLogout:
|
||||
@property
|
||||
def removed(self) -> bool: ...
|
||||
def __repr__(self) -> str: ...
|
||||
|
||||
class OAuthSession:
|
||||
def __init__(self, config: Any) -> None: ...
|
||||
async def login(self) -> SessionStatus: ...
|
||||
async def status(self) -> SessionStatus: ...
|
||||
async def logout(self) -> SessionLogout: ...
|
||||
def __repr__(self) -> str: ...
|
||||
|
||||
class JobFailureInfo:
|
||||
@property
|
||||
def phase(self) -> Optional[str]: ...
|
||||
@@ -399,9 +461,19 @@ class Table:
|
||||
) -> AddColumnsResult: ...
|
||||
async def refresh_column(self, column: str) -> RefreshColumnResult: ...
|
||||
async def refresh_column_async(self, column: str) -> Job: ...
|
||||
async def function_errors(
|
||||
self,
|
||||
job_id: Optional[str] = None,
|
||||
column: Optional[str] = None,
|
||||
limit: Optional[int] = None,
|
||||
) -> FunctionErrors: ...
|
||||
async def refresh_materialized_view(
|
||||
self, full: bool = False, source_version: Optional[int] = None
|
||||
) -> RefreshMaterializedViewResult: ...
|
||||
async def refresh_materialized_view_async(
|
||||
self, full: bool = False, source_version: Optional[int] = None
|
||||
) -> Job: ...
|
||||
async def materialized_view_definition(self) -> str: ...
|
||||
async def add_columns_with_schema(self, schema: pa.Schema) -> AddColumnsResult: ...
|
||||
async def alter_columns(
|
||||
self, columns: list[dict[str, Any]]
|
||||
@@ -754,7 +826,32 @@ class RefreshColumnResult:
|
||||
rows_filled: int
|
||||
version: int
|
||||
|
||||
class FunctionErrorRecord:
|
||||
job_id: str
|
||||
fragment_id: int
|
||||
row_offset: Optional[int]
|
||||
column: str
|
||||
function: str
|
||||
function_version: str
|
||||
table_version: int
|
||||
error_type: str
|
||||
error_message: str
|
||||
created_at_millis: int
|
||||
|
||||
class FunctionErrorFragment:
|
||||
job_id: str
|
||||
fragment_id: int
|
||||
rows_skipped: int
|
||||
rows_recorded: int
|
||||
|
||||
class FunctionErrors:
|
||||
records: list[FunctionErrorRecord]
|
||||
fragments: list[FunctionErrorFragment]
|
||||
truncated: bool
|
||||
|
||||
class RefreshMaterializedViewResult:
|
||||
@staticmethod
|
||||
def from_json(value: str) -> RefreshMaterializedViewResult: ...
|
||||
mode: str
|
||||
rows_written: int
|
||||
source_version: int
|
||||
@@ -803,3 +900,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,
|
||||
)
|
||||
+391
-32
@@ -17,6 +17,7 @@ from typing import (
|
||||
List,
|
||||
Literal,
|
||||
Optional,
|
||||
Sequence,
|
||||
Union,
|
||||
)
|
||||
from uuid import UUID
|
||||
@@ -57,6 +58,12 @@ from .materialized_view import (
|
||||
SelectArg,
|
||||
normalize_select,
|
||||
)
|
||||
from .secrets import (
|
||||
EnvVarSecret,
|
||||
SecretInfo,
|
||||
validate_namespace_path,
|
||||
validate_secret_name,
|
||||
)
|
||||
from .table import (
|
||||
AsyncTable,
|
||||
LanceTable,
|
||||
@@ -529,13 +536,14 @@ class DBConnection(EnforceOverrides):
|
||||
select: SelectArg = None,
|
||||
where: Optional[str] = None,
|
||||
limit: Optional[int] = None,
|
||||
with_no_data: bool = False,
|
||||
) -> MaterializedView:
|
||||
"""Define a materialized view named ``name`` over the table ``source``.
|
||||
|
||||
The view is created empty, with the query recorded in its schema
|
||||
metadata; ``view.refresh()`` computes the rows. The view is a normal
|
||||
table: it can be queried, indexed and searched, and it appears in
|
||||
``table_names``. Local databases only.
|
||||
The view is populated before creation returns. Pass
|
||||
``with_no_data=True`` to create only its definition and empty backing
|
||||
table. The view is a normal table: it can be queried, indexed and
|
||||
searched, and it appears in ``table_names``.
|
||||
|
||||
The source table must have stable row ids (create it with the
|
||||
``new_table_enable_stable_row_ids`` storage option): they keep the
|
||||
@@ -556,6 +564,8 @@ class DBConnection(EnforceOverrides):
|
||||
SQL predicate; only matching source rows appear in the view.
|
||||
limit: int, optional
|
||||
Cap the view at this many rows, in materialization order.
|
||||
with_no_data: bool, default False
|
||||
Skip the initial refresh and leave the backing table empty.
|
||||
|
||||
Returns
|
||||
-------
|
||||
@@ -565,6 +575,27 @@ class DBConnection(EnforceOverrides):
|
||||
"materialized views are not supported on this connection type"
|
||||
)
|
||||
|
||||
def create_materialized_view_async(
|
||||
self,
|
||||
name: str,
|
||||
source: str,
|
||||
*,
|
||||
select: SelectArg = None,
|
||||
where: Optional[str] = None,
|
||||
limit: Optional[int] = None,
|
||||
with_no_data: bool = False,
|
||||
) -> Job[None]:
|
||||
"""Submit materialized-view creation and return its job.
|
||||
|
||||
The job may already be complete for a local database. On LanceDB
|
||||
Cloud and Enterprise, its ``id`` is the server job identifier from
|
||||
the ``202 Accepted`` create response. Wait for the job before opening
|
||||
or querying the view.
|
||||
"""
|
||||
raise NotImplementedError(
|
||||
"materialized views are not supported on this connection type"
|
||||
)
|
||||
|
||||
def open_materialized_view(self, name: str) -> MaterializedView:
|
||||
"""Open the materialized view named ``name``.
|
||||
|
||||
@@ -585,6 +616,32 @@ class DBConnection(EnforceOverrides):
|
||||
"materialized views are not supported on this connection type"
|
||||
)
|
||||
|
||||
def drop_materialized_view(
|
||||
self, name: str, namespace_path: Optional[List[str]] = None
|
||||
) -> None:
|
||||
"""Drop a materialized view.
|
||||
|
||||
The view may become unavailable before physical cleanup finishes. Use
|
||||
:meth:`drop_materialized_view_async` to retain and wait for the cleanup
|
||||
job.
|
||||
"""
|
||||
raise NotImplementedError(
|
||||
"materialized views are not supported on this connection type"
|
||||
)
|
||||
|
||||
def drop_materialized_view_async(
|
||||
self, name: str, namespace_path: Optional[List[str]] = None
|
||||
) -> Job[None]:
|
||||
"""Start dropping a materialized view and return its cleanup job.
|
||||
|
||||
The job may already be complete for a local database. On LanceDB Cloud
|
||||
and Enterprise, its ``id`` is the server job identifier from the
|
||||
``202 Accepted`` drop response.
|
||||
"""
|
||||
raise NotImplementedError(
|
||||
"materialized views are not supported on this connection type"
|
||||
)
|
||||
|
||||
def drop_table(self, name: str, namespace_path: Optional[List[str]] = None):
|
||||
"""Drop a table from the database.
|
||||
|
||||
@@ -692,20 +749,55 @@ class DBConnection(EnforceOverrides):
|
||||
"""
|
||||
raise NotImplementedError("serialize is not supported for this connection type")
|
||||
|
||||
def create_function(self, definition: UdfDefinition) -> FunctionVersion:
|
||||
"""Register a scalar Python UDF and wait for its immutable version.
|
||||
def create_function(
|
||||
self,
|
||||
definition: UdfDefinition,
|
||||
*,
|
||||
secrets: Optional[Sequence[EnvVarSecret]] = None,
|
||||
) -> FunctionVersion:
|
||||
"""Build and register a scalar Python UDF, then return its version.
|
||||
|
||||
The server builds the OCI image and registers the completed artifact.
|
||||
This is the blocking counterpart of :meth:`create_function_async`.
|
||||
Local connections raise ``NotImplementedError``.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
definition : UdfDefinition
|
||||
A callable decorated with [udf][lancedb.udf].
|
||||
secrets : sequence of EnvVarSecret, optional
|
||||
One [EnvVarSecret][lancedb.secrets.EnvVarSecret] per credential the
|
||||
Function needs, each naming a Secret and the environment variable
|
||||
its value arrives in. The Function's source is unchanged by this;
|
||||
it reads the variable the way it already did.
|
||||
|
||||
Examples
|
||||
--------
|
||||
```python
|
||||
db.create_secret("openai-prod", os.environ["OPENAI_API_KEY"])
|
||||
db.create_function(
|
||||
analyze_caption,
|
||||
secrets=[
|
||||
EnvVarSecret(
|
||||
secret_name="openai-prod", env_variable="OPENAI_API_KEY"
|
||||
)
|
||||
],
|
||||
)
|
||||
```
|
||||
"""
|
||||
return self.create_function_async(definition).wait()
|
||||
return self.create_function_async(definition, secrets=secrets).wait()
|
||||
|
||||
def create_function_async(self, definition: UdfDefinition) -> Job[FunctionVersion]:
|
||||
"""Register a scalar Python UDF through the remote Function catalog.
|
||||
def create_function_async(
|
||||
self,
|
||||
definition: UdfDefinition,
|
||||
*,
|
||||
secrets: Optional[Sequence[EnvVarSecret]] = None,
|
||||
) -> Job[FunctionVersion]:
|
||||
"""Submit a scalar Python UDF for building and registration.
|
||||
|
||||
Submission returns a typed job. The immutable Function version becomes
|
||||
available only when :meth:`Job.wait` succeeds. Local connections raise
|
||||
``NotImplementedError``.
|
||||
The server-side job builds the OCI image, then registers the completed
|
||||
artifact. Waiting on the job returns the immutable Function version.
|
||||
Local connections raise ``NotImplementedError``.
|
||||
"""
|
||||
raise NotImplementedError(
|
||||
"Function catalog operations are not supported for this connection type"
|
||||
@@ -736,15 +828,81 @@ class DBConnection(EnforceOverrides):
|
||||
)
|
||||
|
||||
def drop_function(self, name: str, *, version: str) -> bool:
|
||||
"""Drop one exact immutable Function version from the remote catalog.
|
||||
"""Remove the current Function name binding from the remote catalog.
|
||||
|
||||
Returns True when the version changed to Dropped and False for an
|
||||
idempotent replay. Local connections raise NotImplementedError.
|
||||
The requested version must exist in the currently named object.
|
||||
Object history and existing computed-column references are retained.
|
||||
Returns True when the name was removed and False when it was absent.
|
||||
Local connections raise NotImplementedError.
|
||||
"""
|
||||
raise NotImplementedError(
|
||||
"Function catalog operations are not supported for this connection type"
|
||||
)
|
||||
|
||||
def create_secret(
|
||||
self, name: str, value: str, *, namespace_path: Optional[List[str]] = None
|
||||
) -> None:
|
||||
"""Create a named Secret in this database.
|
||||
|
||||
Fails if the name is taken, so a create never silently becomes a
|
||||
rotation. Nothing reads the value back: it is bound to a Function by
|
||||
name and resolved by the service when that Function runs. Local
|
||||
connections raise ``NotImplementedError``.
|
||||
"""
|
||||
raise NotImplementedError(
|
||||
"Secret operations are not supported for this connection type"
|
||||
)
|
||||
|
||||
def alter_secret(
|
||||
self, name: str, value: str, *, namespace_path: Optional[List[str]] = None
|
||||
) -> None:
|
||||
"""Replace the credential behind an existing Secret.
|
||||
|
||||
Fails if it does not exist. Every Function bound to the Secret uses the
|
||||
new value from its next job, and no new Function version is created --
|
||||
which is how a rotation reaches columns pinned to a version registered
|
||||
before it. Local connections raise ``NotImplementedError``.
|
||||
"""
|
||||
raise NotImplementedError(
|
||||
"Secret operations are not supported for this connection type"
|
||||
)
|
||||
|
||||
def list_secrets(self, *, namespace_path: Optional[List[str]] = None) -> List[str]:
|
||||
"""The names of every Secret in this database.
|
||||
|
||||
Names only. No method returns a stored credential, by construction
|
||||
rather than by policy. Local connections raise ``NotImplementedError``.
|
||||
"""
|
||||
raise NotImplementedError(
|
||||
"Secret operations are not supported for this connection type"
|
||||
)
|
||||
|
||||
def drop_secret(
|
||||
self, name: str, *, namespace_path: Optional[List[str]] = None
|
||||
) -> None:
|
||||
"""Drop a Secret.
|
||||
|
||||
Functions bound to it fail at their next job, naming the Secret; that
|
||||
is the revocation path. The name becomes free to reuse, and a new
|
||||
Secret under it is picked up by everything still bound to that name.
|
||||
Local connections raise ``NotImplementedError``.
|
||||
"""
|
||||
raise NotImplementedError(
|
||||
"Secret operations are not supported for this connection type"
|
||||
)
|
||||
|
||||
def describe_secret(
|
||||
self, name: str, *, namespace_path: Optional[List[str]] = None
|
||||
) -> SecretInfo:
|
||||
"""What this database records about a Secret: name and timestamps.
|
||||
|
||||
Never the value -- there is no code path that could return one. Local
|
||||
connections raise ``NotImplementedError``.
|
||||
"""
|
||||
raise NotImplementedError(
|
||||
"Secret operations are not supported for this connection type"
|
||||
)
|
||||
|
||||
def open_job(self, job_id: str) -> Job:
|
||||
"""Open a server-side job by id, returning a handle with its record
|
||||
already populated.
|
||||
@@ -1270,6 +1428,7 @@ class LanceDBConnection(DBConnection):
|
||||
select: SelectArg = None,
|
||||
where: Optional[str] = None,
|
||||
limit: Optional[int] = None,
|
||||
with_no_data: bool = False,
|
||||
) -> MaterializedView:
|
||||
"""Define a materialized view named ``name`` over the table ``source``.
|
||||
See
|
||||
@@ -1290,17 +1449,44 @@ class LanceDBConnection(DBConnection):
|
||||
... select=["name", ("shout", "upper(name)")],
|
||||
... where="age >= 18",
|
||||
... )
|
||||
>>> result = view.refresh()
|
||||
>>> result.rows_written
|
||||
>>> view.table.count_rows()
|
||||
1
|
||||
"""
|
||||
LOOP.run(
|
||||
self._conn.create_materialized_view(
|
||||
name, source, select=select, where=where, limit=limit
|
||||
name,
|
||||
source,
|
||||
select=select,
|
||||
where=where,
|
||||
limit=limit,
|
||||
with_no_data=with_no_data,
|
||||
)
|
||||
)
|
||||
return MaterializedView(self.open_table(name))
|
||||
|
||||
@override
|
||||
def create_materialized_view_async(
|
||||
self,
|
||||
name: str,
|
||||
source: str,
|
||||
*,
|
||||
select: SelectArg = None,
|
||||
where: Optional[str] = None,
|
||||
limit: Optional[int] = None,
|
||||
with_no_data: bool = False,
|
||||
) -> Job[None]:
|
||||
job = LOOP.run(
|
||||
self._conn.create_materialized_view_async(
|
||||
name,
|
||||
source,
|
||||
select=select,
|
||||
where=where,
|
||||
limit=limit,
|
||||
with_no_data=with_no_data,
|
||||
)
|
||||
)
|
||||
return Job(job)
|
||||
|
||||
@override
|
||||
def open_materialized_view(self, name: str) -> MaterializedView:
|
||||
"""Open the materialized view named ``name``."""
|
||||
@@ -1313,6 +1499,25 @@ class LanceDBConnection(DBConnection):
|
||||
"""The names of the materialized views in this database."""
|
||||
return LOOP.run(self._conn.list_materialized_views())
|
||||
|
||||
@override
|
||||
def drop_materialized_view(
|
||||
self, name: str, namespace_path: Optional[List[str]] = None
|
||||
) -> None:
|
||||
if namespace_path is None:
|
||||
namespace_path = []
|
||||
LOOP.run(self._conn.drop_materialized_view(name, namespace_path=namespace_path))
|
||||
|
||||
@override
|
||||
def drop_materialized_view_async(
|
||||
self, name: str, namespace_path: Optional[List[str]] = None
|
||||
) -> Job[None]:
|
||||
if namespace_path is None:
|
||||
namespace_path = []
|
||||
job = LOOP.run(
|
||||
self._conn.drop_materialized_view_async(name, namespace_path=namespace_path)
|
||||
)
|
||||
return Job(job)
|
||||
|
||||
def clone_table(
|
||||
self,
|
||||
target_table_name: str,
|
||||
@@ -1457,8 +1662,13 @@ class LanceDBConnection(DBConnection):
|
||||
return Job(LOOP.run(self._conn.open_job(job_id)))
|
||||
|
||||
@override
|
||||
def create_function_async(self, definition: UdfDefinition) -> Job[FunctionVersion]:
|
||||
job = LOOP.run(self._conn.create_function_async(definition))
|
||||
def create_function_async(
|
||||
self,
|
||||
definition: UdfDefinition,
|
||||
*,
|
||||
secrets: Optional[Sequence[EnvVarSecret]] = None,
|
||||
) -> Job[FunctionVersion]:
|
||||
job = LOOP.run(self._conn.create_function_async(definition, secrets=secrets))
|
||||
return Job(job)
|
||||
|
||||
@override
|
||||
@@ -1473,6 +1683,34 @@ class LanceDBConnection(DBConnection):
|
||||
def drop_function(self, name: str, *, version: str) -> bool:
|
||||
return LOOP.run(self._conn.drop_function(name, version=version))
|
||||
|
||||
@override
|
||||
def create_secret(
|
||||
self, name: str, value: str, *, namespace_path: Optional[List[str]] = None
|
||||
) -> None:
|
||||
LOOP.run(self._conn.create_secret(name, value, namespace_path=namespace_path))
|
||||
|
||||
@override
|
||||
def alter_secret(
|
||||
self, name: str, value: str, *, namespace_path: Optional[List[str]] = None
|
||||
) -> None:
|
||||
LOOP.run(self._conn.alter_secret(name, value, namespace_path=namespace_path))
|
||||
|
||||
@override
|
||||
def list_secrets(self, *, namespace_path: Optional[List[str]] = None) -> List[str]:
|
||||
return LOOP.run(self._conn.list_secrets(namespace_path=namespace_path))
|
||||
|
||||
@override
|
||||
def drop_secret(
|
||||
self, name: str, *, namespace_path: Optional[List[str]] = None
|
||||
) -> None:
|
||||
LOOP.run(self._conn.drop_secret(name, namespace_path=namespace_path))
|
||||
|
||||
@override
|
||||
def describe_secret(
|
||||
self, name: str, *, namespace_path: Optional[List[str]] = None
|
||||
) -> SecretInfo:
|
||||
return LOOP.run(self._conn.describe_secret(name, namespace_path=namespace_path))
|
||||
|
||||
@override
|
||||
def list_jobs(self) -> List[JobInfo]:
|
||||
"""List server-side jobs across the database's tables."""
|
||||
@@ -2080,6 +2318,7 @@ class AsyncConnection(object):
|
||||
select: SelectArg = None,
|
||||
where: Optional[str] = None,
|
||||
limit: Optional[int] = None,
|
||||
with_no_data: bool = False,
|
||||
) -> AsyncMaterializedView:
|
||||
"""Define a materialized view named ``name`` over the table ``source``.
|
||||
See
|
||||
@@ -2091,19 +2330,40 @@ class AsyncConnection(object):
|
||||
projections=normalize_select(select),
|
||||
filter=where,
|
||||
limit=limit,
|
||||
with_no_data=with_no_data,
|
||||
)
|
||||
return AsyncMaterializedView(AsyncTable(inner))
|
||||
|
||||
async def create_materialized_view_async(
|
||||
self,
|
||||
name: str,
|
||||
source: str,
|
||||
*,
|
||||
select: SelectArg = None,
|
||||
where: Optional[str] = None,
|
||||
limit: Optional[int] = None,
|
||||
with_no_data: bool = False,
|
||||
) -> AsyncJob[None]:
|
||||
"""Submit materialized-view creation and return its job.
|
||||
|
||||
Wait for the returned job before opening or querying the view.
|
||||
"""
|
||||
inner = await self._inner.create_materialized_view_async(
|
||||
name,
|
||||
source,
|
||||
projections=normalize_select(select),
|
||||
filter=where,
|
||||
limit=limit,
|
||||
with_no_data=with_no_data,
|
||||
)
|
||||
return AsyncJob(inner)
|
||||
|
||||
async def open_materialized_view(self, name: str) -> AsyncMaterializedView:
|
||||
"""Open the materialized view named ``name``.
|
||||
|
||||
Raises ``ValueError`` if the table exists but is not a materialized
|
||||
view.
|
||||
"""
|
||||
if self.uri.startswith("db://"):
|
||||
raise NotImplementedError(
|
||||
"materialized views are supported only on local databases"
|
||||
)
|
||||
view = AsyncMaterializedView(await self.open_table(name))
|
||||
await view.definition()
|
||||
return view
|
||||
@@ -2116,6 +2376,41 @@ class AsyncConnection(object):
|
||||
"""
|
||||
return await self._inner.list_materialized_views()
|
||||
|
||||
async def drop_materialized_view(
|
||||
self,
|
||||
name: str,
|
||||
*,
|
||||
namespace_path: Optional[List[str]] = None,
|
||||
) -> None:
|
||||
"""Drop a materialized view.
|
||||
|
||||
The view may become unavailable before physical cleanup finishes. Use
|
||||
:meth:`drop_materialized_view_async` to retain and wait for the cleanup
|
||||
job.
|
||||
"""
|
||||
if namespace_path is None:
|
||||
namespace_path = []
|
||||
await self._inner.drop_materialized_view(name, namespace_path=namespace_path)
|
||||
|
||||
async def drop_materialized_view_async(
|
||||
self,
|
||||
name: str,
|
||||
*,
|
||||
namespace_path: Optional[List[str]] = None,
|
||||
) -> AsyncJob[None]:
|
||||
"""Start dropping a materialized view and return its cleanup job.
|
||||
|
||||
Await :meth:`AsyncJob.wait` before assuming physical cleanup has
|
||||
finished.
|
||||
"""
|
||||
if namespace_path is None:
|
||||
namespace_path = []
|
||||
return AsyncJob(
|
||||
await self._inner.drop_materialized_view_async(
|
||||
name, namespace_path=namespace_path
|
||||
)
|
||||
)
|
||||
|
||||
async def clone_table(
|
||||
self,
|
||||
target_table_name: str,
|
||||
@@ -2265,18 +2560,24 @@ class AsyncConnection(object):
|
||||
return AsyncJob(await self._inner.open_job(job_id))
|
||||
|
||||
async def create_function_async(
|
||||
self, definition: UdfDefinition
|
||||
self,
|
||||
definition: UdfDefinition,
|
||||
*,
|
||||
secrets: Optional[Sequence[EnvVarSecret]] = None,
|
||||
) -> AsyncJob[FunctionVersion]:
|
||||
"""Register a scalar Python UDF through the remote Function catalog.
|
||||
"""Submit a scalar Python UDF for building and registration.
|
||||
|
||||
The returned typed job resolves to the immutable Function version.
|
||||
Local connections raise ``NotImplementedError``.
|
||||
The server-side job builds the OCI image, then registers the completed
|
||||
artifact. Waiting on the job returns the immutable Function version.
|
||||
``secrets`` is a sequence of
|
||||
[EnvVarSecret][lancedb.secrets.EnvVarSecret], each naming a Secret and
|
||||
the environment variable its value arrives in. Local connections raise
|
||||
``NotImplementedError``.
|
||||
"""
|
||||
if not isinstance(definition, UdfDefinition):
|
||||
raise TypeError("create_function_async requires a @udf definition")
|
||||
inner = await self._inner.create_function_async(
|
||||
definition.registration_request.to_canonical_json()
|
||||
)
|
||||
request = definition.bind_secrets(secrets)
|
||||
inner = await self._inner.create_function_async(request.to_canonical_json())
|
||||
return _typed_job(inner, FunctionVersion.from_json)
|
||||
|
||||
async def get_function(self, name: str, *, version: str) -> FunctionVersion:
|
||||
@@ -2295,9 +2596,67 @@ class AsyncConnection(object):
|
||||
]
|
||||
|
||||
async def drop_function(self, name: str, *, version: str) -> bool:
|
||||
"""Drop one exact immutable Function version from the remote catalog."""
|
||||
"""Remove the current name binding, retaining the object and its history."""
|
||||
return await self._inner.drop_function(name, version)
|
||||
|
||||
async def create_secret(
|
||||
self, name: str, value: str, *, namespace_path: Optional[List[str]] = None
|
||||
) -> None:
|
||||
"""Create a named Secret in this database.
|
||||
|
||||
Fails if the name is taken, so a create never silently becomes a
|
||||
rotation. Nothing reads the value back.
|
||||
"""
|
||||
await self._inner.create_secret(
|
||||
validate_secret_name(name),
|
||||
value,
|
||||
list(validate_namespace_path(namespace_path)),
|
||||
)
|
||||
|
||||
async def alter_secret(
|
||||
self, name: str, value: str, *, namespace_path: Optional[List[str]] = None
|
||||
) -> None:
|
||||
"""Replace the credential behind an existing Secret.
|
||||
|
||||
Fails if it does not exist. Bound Functions use the new value from
|
||||
their next job, with no new Function version.
|
||||
"""
|
||||
await self._inner.alter_secret(
|
||||
validate_secret_name(name),
|
||||
value,
|
||||
list(validate_namespace_path(namespace_path)),
|
||||
)
|
||||
|
||||
async def list_secrets(
|
||||
self, *, namespace_path: Optional[List[str]] = None
|
||||
) -> List[str]:
|
||||
"""The names of every Secret in this database. Names only."""
|
||||
return await self._inner.list_secrets(
|
||||
list(validate_namespace_path(namespace_path))
|
||||
)
|
||||
|
||||
async def drop_secret(
|
||||
self, name: str, *, namespace_path: Optional[List[str]] = None
|
||||
) -> None:
|
||||
"""Drop a Secret. Bound Functions fail at their next job."""
|
||||
await self._inner.drop_secret(
|
||||
validate_secret_name(name), list(validate_namespace_path(namespace_path))
|
||||
)
|
||||
|
||||
async def describe_secret(
|
||||
self, name: str, *, namespace_path: Optional[List[str]] = None
|
||||
) -> SecretInfo:
|
||||
"""What this database records about a Secret. Never the value."""
|
||||
name, created_at_millis, updated_at_millis = await self._inner.describe_secret(
|
||||
validate_secret_name(name),
|
||||
list(validate_namespace_path(namespace_path)),
|
||||
)
|
||||
return SecretInfo(
|
||||
name=name,
|
||||
created_at_millis=created_at_millis,
|
||||
updated_at_millis=updated_at_millis,
|
||||
)
|
||||
|
||||
async def list_jobs(self) -> List[JobInfo]:
|
||||
"""List server-side jobs across the database's tables."""
|
||||
return await self._inner.list_jobs()
|
||||
|
||||
@@ -249,7 +249,7 @@ def retry_with_exponential_backoff(
|
||||
initial_delay (float): Initial delay in seconds (default is 1).
|
||||
exponential_base (float): The base for exponential backoff (default is 2).
|
||||
jitter (bool): Whether to add jitter to the delay (default is True).
|
||||
max_retries (int): Maximum number of retries (default is 10).
|
||||
max_retries (int): Maximum number of retries (default is 7).
|
||||
|
||||
Returns:
|
||||
function: The decorated function.
|
||||
|
||||
@@ -4,7 +4,7 @@
|
||||
"""Canonical Function values exchanged with LanceDB Enterprise services.
|
||||
|
||||
These immutable models contain client/wire state only. Catalog persistence,
|
||||
environment bake, and execution are owned by Sophon.
|
||||
environment bake, secret resolution, and execution are owned by Sophon.
|
||||
``RefreshColumnResult`` is also the backend-neutral result of a local
|
||||
expression-backed refresh job.
|
||||
"""
|
||||
@@ -25,7 +25,7 @@ import re
|
||||
import sys
|
||||
import textwrap
|
||||
import types
|
||||
from collections.abc import Mapping
|
||||
from collections.abc import Mapping, Sequence
|
||||
from datetime import date, datetime
|
||||
from typing import (
|
||||
Annotated,
|
||||
@@ -41,6 +41,7 @@ from typing import (
|
||||
|
||||
import pyarrow as pa
|
||||
from pydantic import (
|
||||
AfterValidator,
|
||||
BaseModel,
|
||||
ConfigDict,
|
||||
Field,
|
||||
@@ -50,6 +51,7 @@ from pydantic import (
|
||||
)
|
||||
|
||||
from .schema import is_blob_v2_field as _is_blob_v2_field
|
||||
from .secrets import EnvVarSecret
|
||||
|
||||
_Int32 = conint(strict=True, ge=-(2**31), le=2**31 - 1)
|
||||
_UInt32 = conint(strict=True, ge=0, le=2**32 - 1)
|
||||
@@ -226,6 +228,33 @@ class FunctionOutput(_OpenRemoteValue):
|
||||
fields: tuple[FunctionResultField, ...] = ()
|
||||
|
||||
|
||||
class SecretReference(_RemoteValue):
|
||||
"""Where a Secret lives, carried as its parts rather than as one string.
|
||||
|
||||
A joined id would need a delimiter, and a delimiter has to be excluded from
|
||||
every name and segment forever, agreed on by both sides, and re-agreed each
|
||||
time either grows a new way to be configured. Naming the parts settles all
|
||||
of that: nothing here is parsed, so nothing can parse two ways.
|
||||
"""
|
||||
|
||||
name: str
|
||||
namespace_path: tuple[str, ...] = ()
|
||||
|
||||
|
||||
class SecretBinding(_RemoteValue):
|
||||
"""How a Secret reaches the Function that binds it.
|
||||
|
||||
One list rather than a field per delivery mode: a binding is the concept,
|
||||
and how it arrives is a property of one. ``kind`` is open, so a binding a
|
||||
newer service introduces decodes here instead of failing the whole
|
||||
FunctionVersion.
|
||||
"""
|
||||
|
||||
kind: str
|
||||
variable: Optional[str] = None
|
||||
secret_ref: Optional[SecretReference] = None
|
||||
|
||||
|
||||
class FunctionSignature(_RemoteValue):
|
||||
inputs: tuple[FunctionParameter, ...]
|
||||
output: FunctionOutput
|
||||
@@ -295,21 +324,40 @@ class PythonRuntimeSpec(_RemoteValue):
|
||||
return self
|
||||
|
||||
|
||||
class FunctionVersion(_RemoteValue):
|
||||
"""An exact immutable Function version returned by Enterprise.
|
||||
class FunctionImage(_RemoteValue):
|
||||
"""A complete OCI Function image identified by its exact manifest digest."""
|
||||
|
||||
The GPU execution requirement is part of this identity. CPU and memory sizing,
|
||||
priority, concurrency, and retry policy belong to the execution platform.
|
||||
"""
|
||||
manifest_digest: str
|
||||
descriptor: Mapping[str, Any]
|
||||
source: bool
|
||||
|
||||
|
||||
def _validate_object_version(value: str) -> str:
|
||||
if int(value) > 2**64 - 1:
|
||||
raise ValueError("Function version exceeds uint64")
|
||||
return value
|
||||
|
||||
|
||||
_ObjectVersion = Annotated[
|
||||
str,
|
||||
Field(strict=True, pattern=r"^[1-9][0-9]*$"),
|
||||
AfterValidator(_validate_object_version),
|
||||
]
|
||||
|
||||
|
||||
class FunctionVersion(_RemoteValue):
|
||||
"""A pinned object revision, independent of its executable image digest."""
|
||||
|
||||
name: str
|
||||
version: str
|
||||
artifact: FunctionArtifact
|
||||
object_id: str
|
||||
location: str
|
||||
version: _ObjectVersion
|
||||
image: FunctionImage
|
||||
signature: FunctionSignature
|
||||
runtime: PythonRuntimeSpec
|
||||
runtime_digest: str
|
||||
environment_digest: str
|
||||
secret_bindings: tuple[SecretBinding, ...] = ()
|
||||
created_at: str
|
||||
metadata: Mapping[str, str]
|
||||
disabled: bool
|
||||
|
||||
def __call__(self, **inputs: Any) -> FunctionApplication:
|
||||
"""Bind this exact version to named table columns.
|
||||
@@ -363,24 +411,39 @@ class FunctionVersion(_RemoteValue):
|
||||
)
|
||||
)
|
||||
return FunctionApplication(
|
||||
function=FunctionVersionRef(name=self.name, version=self.version),
|
||||
function=FunctionVersionRef(
|
||||
name=self.name,
|
||||
object_id=self.object_id,
|
||||
location=self.location,
|
||||
version=self.version,
|
||||
manifest_digest=self.image.manifest_digest,
|
||||
),
|
||||
inputs=tuple(bindings),
|
||||
output=self.signature.output,
|
||||
)
|
||||
|
||||
|
||||
class FunctionRegistrationRequest(_RemoteValue):
|
||||
"""Stable remote registration envelope produced by :func:`udf`."""
|
||||
"""Stable remote registration envelope produced by :func:`udf`.
|
||||
|
||||
Credential values deliberately have no field here. The only secret-shaped
|
||||
thing a client sends is ``secret_bindings``: the name of a Secret the
|
||||
database already holds, which the remote service resolves at execution.
|
||||
"""
|
||||
|
||||
name: str
|
||||
artifact: FunctionArtifactRequest
|
||||
signature: FunctionSignature
|
||||
runtime: PythonRuntimeSpec
|
||||
secret_bindings: tuple[SecretBinding, ...] = ()
|
||||
|
||||
|
||||
class FunctionVersionRef(_OpenRemoteValue):
|
||||
name: str
|
||||
version: str
|
||||
object_id: str
|
||||
location: str
|
||||
version: _ObjectVersion
|
||||
manifest_digest: str
|
||||
|
||||
|
||||
class ApplicationInput(_OpenRemoteValue):
|
||||
@@ -524,6 +587,7 @@ class RefreshColumnResult(_RemoteValue):
|
||||
|
||||
|
||||
_FUNCTION_NAME = re.compile(r"^[A-Za-z_][A-Za-z0-9_.-]*$")
|
||||
|
||||
_FUNCTION_BLOB_V2_TYPE = "blob_v2"
|
||||
_ARROW_EXTENSION_NAME_KEY = "ARROW:extension:name"
|
||||
_BLOB_V2_EXTENSION_NAME = "lance.blob.v2"
|
||||
@@ -1265,9 +1329,73 @@ class UdfDefinition:
|
||||
|
||||
@property
|
||||
def registration_request(self) -> FunctionRegistrationRequest:
|
||||
"""The immutable request sent by ``create_function_async``."""
|
||||
"""The immutable request sent by ``create_function_async``.
|
||||
|
||||
Carries no secret bindings. Binding is a registration-time decision,
|
||||
so a Function bound to Secrets is registered through :meth:`bind_secrets`,
|
||||
which is what ``create_function`` calls.
|
||||
"""
|
||||
return self._request
|
||||
|
||||
def bind_secrets(
|
||||
self, secrets: Optional[Sequence[EnvVarSecret]]
|
||||
) -> FunctionRegistrationRequest:
|
||||
"""The registration request for this definition bound to ``secrets``.
|
||||
|
||||
Binding does not change the Function's source: each
|
||||
[EnvVarSecret][lancedb.secrets.EnvVarSecret] names a Secret and the
|
||||
environment variable its value should arrive in, and the Function reads
|
||||
that variable the way it already did. Whether the named Secrets exist is
|
||||
the server's answer, not this one.
|
||||
"""
|
||||
bindings = () if secrets is None else tuple(secrets)
|
||||
wrong_type = [
|
||||
binding for binding in bindings if not isinstance(binding, EnvVarSecret)
|
||||
]
|
||||
if wrong_type:
|
||||
kinds = sorted({type(binding).__name__ for binding in wrong_type})
|
||||
raise TypeError(
|
||||
f"Function secrets must be EnvVarSecret values, not {kinds!r}; a "
|
||||
"credential value is never sent to this API"
|
||||
)
|
||||
variables = [binding.env_variable for binding in bindings]
|
||||
duplicates = sorted({name for name in variables if variables.count(name) > 1})
|
||||
if duplicates:
|
||||
raise ValueError(
|
||||
"a Function binds each environment variable once; duplicated: "
|
||||
f"{duplicates!r}"
|
||||
)
|
||||
# `env` is ordinary configuration carried in the definition, so a name in
|
||||
# both would have a value visible in the Function's record and a value
|
||||
# that is not. Refuse rather than pick.
|
||||
environment = self._request.runtime.env or {}
|
||||
overlap = sorted(set(environment) & set(variables))
|
||||
if overlap:
|
||||
raise ValueError(
|
||||
f"Function env and secret bindings must be disjoint: {overlap!r}"
|
||||
)
|
||||
if not bindings:
|
||||
return self._request
|
||||
# Sorted, because the list is carried in the FunctionVersion hash and a
|
||||
# caller's argument order is not part of what a Function is.
|
||||
resolved = tuple(
|
||||
sorted(
|
||||
(
|
||||
SecretBinding(
|
||||
kind="env",
|
||||
variable=binding.env_variable,
|
||||
secret_ref=SecretReference(
|
||||
name=binding.secret_name,
|
||||
namespace_path=tuple(binding.secret_namespace_path),
|
||||
),
|
||||
)
|
||||
for binding in bindings
|
||||
),
|
||||
key=lambda binding: (binding.kind, binding.variable or ""),
|
||||
)
|
||||
)
|
||||
return self._request._copy(update={"secret_bindings": resolved})
|
||||
|
||||
def __call__(self, *args, **kwargs):
|
||||
return self._function(*args, **kwargs)
|
||||
|
||||
@@ -1332,7 +1460,9 @@ def udf(
|
||||
conda_channels : sequence of str, optional
|
||||
Conda channels in priority order; requires ``conda``.
|
||||
env : mapping of str to str, optional
|
||||
Environment variables included in the Function definition.
|
||||
Environment variables included in the Function definition. Not for
|
||||
credentials -- these are ordinary configuration, stored with the
|
||||
Function and visible wherever it is.
|
||||
python_version : str, optional
|
||||
Remote Python major/minor version. Defaults to the client version.
|
||||
gpu : bool, default False
|
||||
@@ -1402,6 +1532,7 @@ __all__ = [
|
||||
"FunctionRegistrationRequest",
|
||||
"FunctionResultField",
|
||||
"FunctionSignature",
|
||||
"FunctionImage",
|
||||
"FunctionVersion",
|
||||
"FunctionVersionRef",
|
||||
"InputBinding",
|
||||
|
||||
@@ -11,6 +11,7 @@ from dataclasses import dataclass, field
|
||||
from typing import TYPE_CHECKING, Dict, List, Optional, Sequence, Tuple, Union
|
||||
|
||||
from .background_loop import LOOP
|
||||
from .job import AsyncJob, Job, _typed_job
|
||||
|
||||
if TYPE_CHECKING:
|
||||
import pyarrow as pa
|
||||
@@ -73,6 +74,20 @@ def _definition_from_schema(
|
||||
)
|
||||
|
||||
|
||||
def _definition_from_json(raw: str) -> MaterializedViewDefinition:
|
||||
value = json.loads(raw)
|
||||
return MaterializedViewDefinition(
|
||||
source_table=value["source_table"],
|
||||
projections=[
|
||||
(p["output"], p["expression"]) for p in value.get("projections", [])
|
||||
],
|
||||
filter=value.get("filter"),
|
||||
limit=value.get("limit"),
|
||||
inputs=value.get("inputs", []),
|
||||
source_namespace=value.get("source_namespace", []),
|
||||
)
|
||||
|
||||
|
||||
def _quote_identifier(name: str) -> str:
|
||||
"""Quote a column name as a Lance SQL identifier (backticks)."""
|
||||
escaped = name.replace("`", "``")
|
||||
@@ -126,8 +141,9 @@ class AsyncMaterializedView:
|
||||
return self._table
|
||||
|
||||
async def definition(self) -> MaterializedViewDefinition:
|
||||
"""The query that defines the view, read from its stored schema."""
|
||||
return _definition_from_schema(await self._table.schema(), self.name)
|
||||
"""The query that defines the view."""
|
||||
raw = await self._table._inner.materialized_view_definition()
|
||||
return _definition_from_json(raw)
|
||||
|
||||
async def refresh(
|
||||
self, *, full: bool = False, source_version: Optional[int] = None
|
||||
@@ -148,6 +164,24 @@ class AsyncMaterializedView:
|
||||
full=full, source_version=source_version
|
||||
)
|
||||
|
||||
async def refresh_async(
|
||||
self, *, full: bool = False, source_version: Optional[int] = None
|
||||
) -> "AsyncJob[RefreshMaterializedViewResult]":
|
||||
"""Submit a refresh and return its job without waiting.
|
||||
|
||||
The job may already be complete for a local view. On LanceDB Cloud
|
||||
and Enterprise, its ``id`` is the server job identifier returned by
|
||||
the refresh endpoint.
|
||||
"""
|
||||
from ._lancedb import RefreshMaterializedViewResult
|
||||
|
||||
return _typed_job(
|
||||
await self._table._inner.refresh_materialized_view_async(
|
||||
full=full, source_version=source_version
|
||||
),
|
||||
RefreshMaterializedViewResult.from_json,
|
||||
)
|
||||
|
||||
|
||||
class MaterializedView:
|
||||
"""Synchronous variant of
|
||||
@@ -171,8 +205,8 @@ class MaterializedView:
|
||||
|
||||
@property
|
||||
def definition(self) -> MaterializedViewDefinition:
|
||||
"""The query that defines the view, read from its stored schema."""
|
||||
return _definition_from_schema(self._table.schema, self.name)
|
||||
"""The query that defines the view."""
|
||||
return LOOP.run(self._async.definition())
|
||||
|
||||
def refresh(
|
||||
self, *, full: bool = False, source_version: Optional[int] = None
|
||||
@@ -180,3 +214,17 @@ class MaterializedView:
|
||||
"""Recompute the view from its source. See
|
||||
[AsyncMaterializedView.refresh][lancedb.materialized_view.AsyncMaterializedView.refresh]."""
|
||||
return LOOP.run(self._async.refresh(full=full, source_version=source_version))
|
||||
|
||||
def refresh_async(
|
||||
self, *, full: bool = False, source_version: Optional[int] = None
|
||||
) -> "Job[RefreshMaterializedViewResult]":
|
||||
"""Submit a refresh and return its job without waiting.
|
||||
|
||||
See
|
||||
[AsyncMaterializedView.refresh_async][lancedb.materialized_view.AsyncMaterializedView.refresh_async].
|
||||
"""
|
||||
return Job(
|
||||
LOOP.run(
|
||||
self._async.refresh_async(full=full, source_version=source_version)
|
||||
)
|
||||
)
|
||||
|
||||
@@ -637,6 +637,7 @@ class LanceNamespaceDBConnection(DBConnection):
|
||||
select: "SelectArg" = None,
|
||||
where: Optional[str] = None,
|
||||
limit: Optional[int] = None,
|
||||
with_no_data: bool = False,
|
||||
) -> "MaterializedView":
|
||||
"""Define a materialized view over a table in the root namespace.
|
||||
See
|
||||
@@ -646,12 +647,40 @@ class LanceNamespaceDBConnection(DBConnection):
|
||||
self.open_table(
|
||||
LOOP.run(
|
||||
self._inner.create_materialized_view(
|
||||
name, source, select=select, where=where, limit=limit
|
||||
name,
|
||||
source,
|
||||
select=select,
|
||||
where=where,
|
||||
limit=limit,
|
||||
with_no_data=with_no_data,
|
||||
)
|
||||
).name
|
||||
)
|
||||
)
|
||||
|
||||
@override
|
||||
def create_materialized_view_async(
|
||||
self,
|
||||
name: str,
|
||||
source: str,
|
||||
*,
|
||||
select: "SelectArg" = None,
|
||||
where: Optional[str] = None,
|
||||
limit: Optional[int] = None,
|
||||
with_no_data: bool = False,
|
||||
) -> Job[None]:
|
||||
job = LOOP.run(
|
||||
self._inner.create_materialized_view_async(
|
||||
name,
|
||||
source,
|
||||
select=select,
|
||||
where=where,
|
||||
limit=limit,
|
||||
with_no_data=with_no_data,
|
||||
)
|
||||
)
|
||||
return Job(job)
|
||||
|
||||
@override
|
||||
def open_materialized_view(self, name: str) -> "MaterializedView":
|
||||
"""Open the materialized view named ``name``."""
|
||||
@@ -664,6 +693,30 @@ class LanceNamespaceDBConnection(DBConnection):
|
||||
"""The names of the materialized views in the root namespace."""
|
||||
return LOOP.run(self._inner.list_materialized_views())
|
||||
|
||||
@override
|
||||
def drop_materialized_view(
|
||||
self, name: str, namespace_path: Optional[List[str]] = None
|
||||
) -> None:
|
||||
if namespace_path is None:
|
||||
namespace_path = []
|
||||
LOOP.run(
|
||||
self._inner.drop_materialized_view(name, namespace_path=namespace_path)
|
||||
)
|
||||
|
||||
@override
|
||||
def drop_materialized_view_async(
|
||||
self, name: str, namespace_path: Optional[List[str]] = None
|
||||
) -> Job[None]:
|
||||
if namespace_path is None:
|
||||
namespace_path = []
|
||||
return Job(
|
||||
LOOP.run(
|
||||
self._inner.drop_materialized_view_async(
|
||||
name, namespace_path=namespace_path
|
||||
)
|
||||
)
|
||||
)
|
||||
|
||||
@override
|
||||
def drop_table(self, name: str, namespace_path: Optional[List[str]] = None):
|
||||
if namespace_path is None:
|
||||
@@ -1194,15 +1247,41 @@ class AsyncLanceNamespaceDBConnection:
|
||||
select: "SelectArg" = None,
|
||||
where: Optional[str] = None,
|
||||
limit: Optional[int] = None,
|
||||
with_no_data: bool = False,
|
||||
) -> "AsyncMaterializedView":
|
||||
"""Define a materialized view over a table in the root namespace."""
|
||||
view = await self._inner.create_materialized_view(
|
||||
name, source, select=select, where=where, limit=limit
|
||||
name,
|
||||
source,
|
||||
select=select,
|
||||
where=where,
|
||||
limit=limit,
|
||||
with_no_data=with_no_data,
|
||||
)
|
||||
# Reopen through the namespace so the view's table carries the
|
||||
# namespace client and pushdown configuration a bare inner table lacks.
|
||||
return AsyncMaterializedView(await self.open_table(view.name))
|
||||
|
||||
async def create_materialized_view_async(
|
||||
self,
|
||||
name: str,
|
||||
source: str,
|
||||
*,
|
||||
select: "SelectArg" = None,
|
||||
where: Optional[str] = None,
|
||||
limit: Optional[int] = None,
|
||||
with_no_data: bool = False,
|
||||
) -> AsyncJob[None]:
|
||||
"""Submit materialized-view creation and return its job."""
|
||||
return await self._inner.create_materialized_view_async(
|
||||
name,
|
||||
source,
|
||||
select=select,
|
||||
where=where,
|
||||
limit=limit,
|
||||
with_no_data=with_no_data,
|
||||
)
|
||||
|
||||
async def open_materialized_view(self, name: str) -> "AsyncMaterializedView":
|
||||
"""Open the materialized view named ``name``."""
|
||||
view = AsyncMaterializedView(await self.open_table(name))
|
||||
@@ -1213,6 +1292,24 @@ class AsyncLanceNamespaceDBConnection:
|
||||
"""The names of the materialized views in the root namespace."""
|
||||
return await self._inner.list_materialized_views()
|
||||
|
||||
async def drop_materialized_view(
|
||||
self, name: str, namespace_path: Optional[List[str]] = None
|
||||
) -> None:
|
||||
"""Drop a materialized view from the namespace."""
|
||||
if namespace_path is None:
|
||||
namespace_path = []
|
||||
await self._inner.drop_materialized_view(name, namespace_path=namespace_path)
|
||||
|
||||
async def drop_materialized_view_async(
|
||||
self, name: str, namespace_path: Optional[List[str]] = None
|
||||
) -> AsyncJob[None]:
|
||||
"""Start dropping a materialized view and return its cleanup job."""
|
||||
if namespace_path is None:
|
||||
namespace_path = []
|
||||
return await self._inner.drop_materialized_view_async(
|
||||
name, namespace_path=namespace_path
|
||||
)
|
||||
|
||||
async def drop_table(self, name: str, namespace_path: Optional[List[str]] = None):
|
||||
"""Drop a table from the namespace."""
|
||||
if namespace_path is None:
|
||||
|
||||
@@ -9,7 +9,13 @@ from typing import List, Optional
|
||||
from lancedb import __version__
|
||||
|
||||
from .header import HeaderProvider
|
||||
from .oauth import OAuthConfig, OAuthFlowType
|
||||
from .oauth import (
|
||||
ClientAuthMethod,
|
||||
OAuthConfig,
|
||||
OAuthFlowType,
|
||||
OAuthSession,
|
||||
TokenCacheOptions,
|
||||
)
|
||||
|
||||
# The API reference renders this module with a single mkdocstrings directive,
|
||||
# which only picks up names listed here. New public names must be added to this
|
||||
@@ -22,6 +28,9 @@ __all__ = [
|
||||
"HeaderProvider",
|
||||
"OAuthConfig",
|
||||
"OAuthFlowType",
|
||||
"ClientAuthMethod",
|
||||
"OAuthSession",
|
||||
"TokenCacheOptions",
|
||||
]
|
||||
|
||||
|
||||
@@ -164,7 +173,10 @@ class ClientConfig:
|
||||
extra_headers: Optional[dict]
|
||||
Additional headers to include in requests.
|
||||
id_delimiter: Optional[str]
|
||||
The delimiter to use when constructing object identifiers.
|
||||
The delimiter joining a namespace path and a name into one object
|
||||
identifier. ``"$"`` is the only supported value, and leaving this
|
||||
unset is how to get it; anything else is rejected when the connection
|
||||
is created.
|
||||
tls_config: Optional[TlsConfig]
|
||||
TLS/mTLS configuration for secure connections.
|
||||
header_provider: Optional[HeaderProvider]
|
||||
|
||||
@@ -2,12 +2,22 @@
|
||||
# SPDX-FileCopyrightText: Copyright The LanceDB Authors
|
||||
|
||||
|
||||
from dataclasses import replace
|
||||
from datetime import timedelta
|
||||
import json
|
||||
import logging
|
||||
from concurrent.futures import ThreadPoolExecutor
|
||||
import sys
|
||||
from typing import TYPE_CHECKING, Any, Dict, Iterable, List, Optional, Union
|
||||
from typing import (
|
||||
TYPE_CHECKING,
|
||||
Any,
|
||||
Dict,
|
||||
Iterable,
|
||||
List,
|
||||
Optional,
|
||||
Sequence,
|
||||
Union,
|
||||
)
|
||||
from urllib.parse import urlparse
|
||||
from uuid import UUID
|
||||
import warnings
|
||||
@@ -29,6 +39,7 @@ from ..job import AsyncJob, Job
|
||||
from ..sql import Query as SqlQuery
|
||||
from ..sql import QueryDescription
|
||||
from ..materialized_view import MaterializedView, SelectArg
|
||||
from ..secrets import EnvVarSecret, SecretInfo
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from .._lancedb import JobInfo
|
||||
@@ -187,11 +198,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",
|
||||
@@ -665,22 +716,80 @@ class RemoteDBConnection(DBConnection):
|
||||
select: SelectArg = None,
|
||||
where: Optional[str] = None,
|
||||
limit: Optional[int] = None,
|
||||
with_no_data: bool = False,
|
||||
) -> MaterializedView:
|
||||
raise NotImplementedError(
|
||||
"materialized views are supported only on local databases"
|
||||
from .table import RemoteTable
|
||||
|
||||
view = LOOP.run(
|
||||
self._conn.create_materialized_view(
|
||||
name,
|
||||
source,
|
||||
select=select,
|
||||
where=where,
|
||||
limit=limit,
|
||||
with_no_data=with_no_data,
|
||||
)
|
||||
)
|
||||
return MaterializedView(
|
||||
RemoteTable(
|
||||
view.table,
|
||||
self.db_name,
|
||||
connection_state=self.serialize,
|
||||
namespace_path=[],
|
||||
)
|
||||
)
|
||||
|
||||
@override
|
||||
def create_materialized_view_async(
|
||||
self,
|
||||
name: str,
|
||||
source: str,
|
||||
*,
|
||||
select: SelectArg = None,
|
||||
where: Optional[str] = None,
|
||||
limit: Optional[int] = None,
|
||||
with_no_data: bool = False,
|
||||
) -> Job[None]:
|
||||
job = LOOP.run(
|
||||
self._conn.create_materialized_view_async(
|
||||
name,
|
||||
source,
|
||||
select=select,
|
||||
where=where,
|
||||
limit=limit,
|
||||
with_no_data=with_no_data,
|
||||
)
|
||||
)
|
||||
return Job(job)
|
||||
|
||||
@override
|
||||
def open_materialized_view(self, name: str) -> MaterializedView:
|
||||
raise NotImplementedError(
|
||||
"materialized views are supported only on local databases"
|
||||
)
|
||||
view = MaterializedView(self.open_table(name))
|
||||
view.definition
|
||||
return view
|
||||
|
||||
@override
|
||||
def list_materialized_views(self) -> List[str]:
|
||||
raise NotImplementedError(
|
||||
"materialized views are supported only on local databases"
|
||||
return LOOP.run(self._conn.list_materialized_views())
|
||||
|
||||
@override
|
||||
def drop_materialized_view(
|
||||
self, name: str, namespace_path: Optional[List[str]] = None
|
||||
) -> None:
|
||||
if namespace_path is None:
|
||||
namespace_path = []
|
||||
LOOP.run(self._conn.drop_materialized_view(name, namespace_path=namespace_path))
|
||||
|
||||
@override
|
||||
def drop_materialized_view_async(
|
||||
self, name: str, namespace_path: Optional[List[str]] = None
|
||||
) -> Job[None]:
|
||||
if namespace_path is None:
|
||||
namespace_path = []
|
||||
job = LOOP.run(
|
||||
self._conn.drop_materialized_view_async(name, namespace_path=namespace_path)
|
||||
)
|
||||
return Job(job)
|
||||
|
||||
@override
|
||||
def drop_table(self, name: str, namespace_path: Optional[List[str]] = None):
|
||||
@@ -746,8 +855,14 @@ class RemoteDBConnection(DBConnection):
|
||||
return Job(LOOP.run(self._conn.open_job(job_id)))
|
||||
|
||||
@override
|
||||
def create_function_async(self, definition: UdfDefinition) -> Job[FunctionVersion]:
|
||||
return Job(LOOP.run(self._conn.create_function_async(definition)))
|
||||
def create_function_async(
|
||||
self,
|
||||
definition: UdfDefinition,
|
||||
*,
|
||||
secrets: Optional[Sequence[EnvVarSecret]] = None,
|
||||
) -> Job[FunctionVersion]:
|
||||
job = LOOP.run(self._conn.create_function_async(definition, secrets=secrets))
|
||||
return Job(job)
|
||||
|
||||
@override
|
||||
def get_function(self, name: str, *, version: str) -> FunctionVersion:
|
||||
@@ -761,6 +876,34 @@ class RemoteDBConnection(DBConnection):
|
||||
def drop_function(self, name: str, *, version: str) -> bool:
|
||||
return LOOP.run(self._conn.drop_function(name, version=version))
|
||||
|
||||
@override
|
||||
def create_secret(
|
||||
self, name: str, value: str, *, namespace_path: Optional[List[str]] = None
|
||||
) -> None:
|
||||
LOOP.run(self._conn.create_secret(name, value, namespace_path=namespace_path))
|
||||
|
||||
@override
|
||||
def alter_secret(
|
||||
self, name: str, value: str, *, namespace_path: Optional[List[str]] = None
|
||||
) -> None:
|
||||
LOOP.run(self._conn.alter_secret(name, value, namespace_path=namespace_path))
|
||||
|
||||
@override
|
||||
def describe_secret(
|
||||
self, name: str, *, namespace_path: Optional[List[str]] = None
|
||||
) -> SecretInfo:
|
||||
return LOOP.run(self._conn.describe_secret(name, namespace_path=namespace_path))
|
||||
|
||||
@override
|
||||
def list_secrets(self, *, namespace_path: Optional[List[str]] = None) -> List[str]:
|
||||
return LOOP.run(self._conn.list_secrets(namespace_path=namespace_path))
|
||||
|
||||
@override
|
||||
def drop_secret(
|
||||
self, name: str, *, namespace_path: Optional[List[str]] = None
|
||||
) -> None:
|
||||
LOOP.run(self._conn.drop_secret(name, namespace_path=namespace_path))
|
||||
|
||||
@override
|
||||
def list_jobs(self) -> List["JobInfo"]:
|
||||
"""List server-side jobs across the database's tables."""
|
||||
|
||||
@@ -12,10 +12,72 @@ class OAuthFlowType(str, Enum):
|
||||
CLIENT_CREDENTIALS = "client_credentials"
|
||||
"""Client Credentials grant (service-to-service / M2M)."""
|
||||
|
||||
AUTHORIZATION_CODE = "authorization_code"
|
||||
"""Interactive Authorization Code grant, using PKCE by default."""
|
||||
|
||||
DEVICE_CODE = "device_code"
|
||||
"""Device Authorization grant for CLI and headless environments."""
|
||||
|
||||
AZURE_MANAGED_IDENTITY = "azure_managed_identity"
|
||||
"""Azure Managed Identity via IMDS."""
|
||||
|
||||
|
||||
class ClientAuthMethod(str, Enum):
|
||||
"""How the client authenticates to the OAuth token endpoint.
|
||||
|
||||
The method applies to every OAuth request that carries client
|
||||
authentication: client-credentials, authorization-code exchange,
|
||||
refresh-token, and device-authorization requests. The Azure managed
|
||||
identity flow ignores this option.
|
||||
"""
|
||||
|
||||
NONE = "none"
|
||||
"""No client authentication, for public clients using PKCE or the device
|
||||
flow. Cannot be combined with ``client_secret``."""
|
||||
|
||||
CLIENT_SECRET_BASIC = "client_secret_basic"
|
||||
"""HTTP Basic authentication. This is the RFC 6749 recommended method and
|
||||
the normal default for confidential clients, including default Okta
|
||||
applications. Requires ``client_secret``."""
|
||||
|
||||
CLIENT_SECRET_POST = "client_secret_post"
|
||||
"""Credentials in the request body, for providers configured to require
|
||||
it. Requires ``client_secret``."""
|
||||
|
||||
|
||||
@dataclass
|
||||
class TokenCacheOptions:
|
||||
"""Options for the persistent OAuth token cache.
|
||||
|
||||
The cache is opt-in: it is only used when set as ``token_cache`` on
|
||||
:class:`OAuthConfig`. Only refresh tokens are persisted, in a private
|
||||
directory with owner-only permissions, so short-lived processes can reuse
|
||||
an authenticated session instead of re-prompting on every start.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
cache_dir : Optional[str]
|
||||
Directory that holds cached credentials. Defaults to
|
||||
``$XDG_CACHE_HOME/lancedb/oauth``, ``$HOME/.cache/lancedb/oauth`` on
|
||||
Unix, or ``%LOCALAPPDATA%\\lancedb\\oauth`` on Windows. The directory
|
||||
is created with owner-only permissions (``0700``) when missing.
|
||||
lock_timeout_secs : Optional[int]
|
||||
How long to wait for the cross-process refresh lock before failing
|
||||
(default: 30 seconds).
|
||||
|
||||
Examples
|
||||
--------
|
||||
>>> opts = TokenCacheOptions(cache_dir="/tmp/my-app/oauth-cache")
|
||||
|
||||
Multiple identities (issuer, client, scopes, flow, client
|
||||
authentication) get separate cache entries. Within one identity the most
|
||||
recent login wins.
|
||||
"""
|
||||
|
||||
cache_dir: Optional[str] = None
|
||||
lock_timeout_secs: Optional[int] = None
|
||||
|
||||
|
||||
@dataclass
|
||||
class OAuthConfig:
|
||||
"""OAuth configuration for LanceDB authentication.
|
||||
@@ -38,12 +100,37 @@ class OAuthConfig:
|
||||
Authentication flow to use. Default: CLIENT_CREDENTIALS.
|
||||
client_secret : Optional[str]
|
||||
Client secret (required for CLIENT_CREDENTIALS).
|
||||
client_auth_method : Optional[ClientAuthMethod]
|
||||
How the client authenticates to the token endpoint (default: auto).
|
||||
With a ``client_secret`` the default is
|
||||
``ClientAuthMethod.CLIENT_SECRET_BASIC``, which matches the RFC 6749
|
||||
recommendation and the default configuration of Okta confidential
|
||||
applications; without a secret the client is public and no client
|
||||
authentication is sent.
|
||||
redirect_uri : Optional[str]
|
||||
Loopback redirect URI for AUTHORIZATION_CODE. The default is
|
||||
``http://127.0.0.1:{callback_port}/callback``.
|
||||
callback_port : Optional[int]
|
||||
Port for the AUTHORIZATION_CODE loopback callback server (default: 8400).
|
||||
use_pkce : bool
|
||||
Protect AUTHORIZATION_CODE with S256 PKCE (default: True).
|
||||
managed_identity_client_id : Optional[str]
|
||||
Client ID for user-assigned managed identity (AZURE_MANAGED_IDENTITY).
|
||||
resource : Optional[str]
|
||||
Resource indicator (RFC 8707), forwarded verbatim to authorization and
|
||||
token endpoints, including refresh requests. Must be an absolute URI
|
||||
without a fragment. Not supported for Azure managed identity.
|
||||
audience : Optional[str]
|
||||
Provider-specific audience, forwarded to authorization and token
|
||||
endpoints, including refresh requests. Not supported for Azure managed identity.
|
||||
refresh_buffer_secs : Optional[int]
|
||||
Seconds before expiry to trigger proactive refresh (default: 300).
|
||||
Keep this well below the token TTL; if it is greater than or equal to
|
||||
the TTL, each request refreshes the token.
|
||||
token_cache : Optional[TokenCacheOptions]
|
||||
Opt in to the persistent token cache so short-lived processes reuse
|
||||
one session. Only supported by AUTHORIZATION_CODE and DEVICE_CODE;
|
||||
azure managed identity is rejected. Default: None (memory only).
|
||||
|
||||
Examples
|
||||
--------
|
||||
@@ -56,6 +143,12 @@ class OAuthConfig:
|
||||
... scopes=["api://lancedb-api/.default"],
|
||||
... )
|
||||
|
||||
Providers that require an explicit target can use ``resource`` and/or
|
||||
``audience`` (these are forwarded unchanged):
|
||||
|
||||
>>> config.resource = "https://api.example.com"
|
||||
>>> config.audience = "lancedb-api"
|
||||
|
||||
Azure Managed Identity:
|
||||
|
||||
>>> config = OAuthConfig(
|
||||
@@ -64,6 +157,30 @@ class OAuthConfig:
|
||||
... scopes=["api://lancedb-api/.default"],
|
||||
... flow=OAuthFlowType.AZURE_MANAGED_IDENTITY,
|
||||
... )
|
||||
|
||||
Authorization Code with PKCE:
|
||||
|
||||
The authorization URL is written to standard error before LanceDB tries to
|
||||
open a browser, so it can be copied in headless environments.
|
||||
|
||||
>>> config = OAuthConfig(
|
||||
... issuer_url="https://login.microsoftonline.com/{tenant}/v2.0",
|
||||
... client_id="app-id",
|
||||
... scopes=["openid", "api://lancedb-api/access"],
|
||||
... flow=OAuthFlowType.AUTHORIZATION_CODE,
|
||||
... )
|
||||
|
||||
Device Authorization, with a persistent cache so later processes reuse
|
||||
the session without a new device prompt. The verification URL and user
|
||||
code are written to standard error before polling begins:
|
||||
|
||||
>>> config = OAuthConfig(
|
||||
... issuer_url="https://login.microsoftonline.com/{tenant}/v2.0",
|
||||
... client_id="app-id",
|
||||
... scopes=["openid", "offline_access", "api://lancedb-api/access"],
|
||||
... flow=OAuthFlowType.DEVICE_CODE,
|
||||
... token_cache=TokenCacheOptions(),
|
||||
... )
|
||||
"""
|
||||
|
||||
issuer_url: str
|
||||
@@ -71,5 +188,73 @@ class OAuthConfig:
|
||||
scopes: List[str]
|
||||
flow: OAuthFlowType = OAuthFlowType.CLIENT_CREDENTIALS
|
||||
client_secret: Optional[str] = field(default=None, repr=False)
|
||||
client_auth_method: Optional[ClientAuthMethod] = None
|
||||
redirect_uri: Optional[str] = None
|
||||
callback_port: Optional[int] = None
|
||||
use_pkce: bool = True
|
||||
managed_identity_client_id: Optional[str] = None
|
||||
refresh_buffer_secs: Optional[int] = None
|
||||
token_cache: Optional[TokenCacheOptions] = None
|
||||
resource: Optional[str] = None
|
||||
audience: Optional[str] = None
|
||||
|
||||
|
||||
class OAuthSession:
|
||||
"""Explicit OAuth session lifecycle for the persistent token cache.
|
||||
|
||||
Built from the same :class:`OAuthConfig` used for
|
||||
:func:`lancedb.connect_async` (including its ``token_cache`` options).
|
||||
A connection created with the same configuration shares the cache, so
|
||||
logging in here prepares tokens for later processes without any database
|
||||
request.
|
||||
|
||||
``login`` always runs the configured interactive flow and replaces the
|
||||
cached session (the most recent login wins). ``logout`` removes only the
|
||||
local credential; it does not revoke anything with the provider and does
|
||||
not sign out of a browser SSO session.
|
||||
|
||||
Examples
|
||||
--------
|
||||
>>> config = OAuthConfig(
|
||||
... issuer_url="https://issuer.example.com",
|
||||
... client_id="my-app",
|
||||
... scopes=["openid", "offline_access"],
|
||||
... flow=OAuthFlowType.DEVICE_CODE,
|
||||
... token_cache=TokenCacheOptions(),
|
||||
... )
|
||||
>>> session = OAuthSession(config) # doctest: +SKIP
|
||||
>>> status = await session.login() # doctest: +SKIP
|
||||
>>> status.refreshable # doctest: +SKIP
|
||||
True
|
||||
"""
|
||||
|
||||
def __init__(self, config: OAuthConfig):
|
||||
from lancedb._lancedb import OAuthSession as PyOAuthSession
|
||||
|
||||
self._inner: PyOAuthSession = PyOAuthSession(config)
|
||||
|
||||
async def login(self):
|
||||
"""Eagerly run the configured flow and store the session.
|
||||
|
||||
Returns a :class:`lancedb._lancedb.SessionStatus` describing the
|
||||
cached session. A successful login always replaces any prior cached
|
||||
session for this identity; if the provider does not issue a refresh
|
||||
token (for example without ``offline_access``), the previous record
|
||||
is removed and ``refreshable`` is ``False``.
|
||||
"""
|
||||
return await self._inner.login()
|
||||
|
||||
async def status(self):
|
||||
"""Report whether a cached session exists, with safe metadata.
|
||||
|
||||
Never contacts the identity provider and never exposes token values.
|
||||
"""
|
||||
return await self._inner.status()
|
||||
|
||||
async def logout(self):
|
||||
"""Remove the matching local cached credential.
|
||||
|
||||
Returns a :class:`lancedb._lancedb.SessionLogout` whose ``removed``
|
||||
flag reports whether a credential existed. Logout is idempotent.
|
||||
"""
|
||||
return await self._inner.logout()
|
||||
|
||||
@@ -994,6 +994,16 @@ class RemoteTable(Table):
|
||||
def refresh_column_async(self, column: str) -> Job[RefreshColumnResult]:
|
||||
return Job(LOOP.run(self._table.refresh_column_async(column)))
|
||||
|
||||
def function_errors(
|
||||
self,
|
||||
job_id: Optional[str] = None,
|
||||
column: Optional[str] = None,
|
||||
limit: Optional[int] = None,
|
||||
):
|
||||
return LOOP.run(
|
||||
self._table.function_errors(job_id=job_id, column=column, limit=limit)
|
||||
)
|
||||
|
||||
def alter_columns(
|
||||
self, *alterations: Iterable[Dict[str, str]]
|
||||
) -> AlterColumnsResult:
|
||||
|
||||
@@ -11,6 +11,7 @@ from .jinaai import JinaReranker
|
||||
from .rrf import RRFReranker
|
||||
from .mrr import MRRReranker
|
||||
from .answerdotai import AnswerdotaiRerankers
|
||||
from .typesafe import TypeSafeReranker
|
||||
from .voyageai import VoyageAIReranker
|
||||
from .watsonx import WatsonxReranker
|
||||
|
||||
@@ -27,6 +28,7 @@ __all__ = [
|
||||
"JinaReranker",
|
||||
"RRFReranker",
|
||||
"AnswerdotaiRerankers",
|
||||
"TypeSafeReranker",
|
||||
"VoyageAIReranker",
|
||||
"MRRReranker",
|
||||
"WatsonxReranker",
|
||||
|
||||
@@ -0,0 +1,159 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright The LanceDB Authors
|
||||
|
||||
|
||||
from concurrent.futures import ThreadPoolExecutor
|
||||
from functools import cached_property
|
||||
from typing import Any, Dict, Mapping, Optional
|
||||
|
||||
import pyarrow as pa
|
||||
|
||||
from ..util import attempt_import_or_raise
|
||||
from .base import Reranker
|
||||
|
||||
DEFAULT_INSTRUCTIONS = (
|
||||
"Does the document contain information that answers or directly addresses "
|
||||
"the query?"
|
||||
)
|
||||
DEFAULT_CRITERIA = {
|
||||
"true": "The document answers the query or states facts the query asks about.",
|
||||
"false": "The document is off-topic, or only shares keywords or a general "
|
||||
"subject with the query without addressing it.",
|
||||
}
|
||||
# The key TypeSafe returns the answer under. It is not sent to the model.
|
||||
_QUESTION_ID = "relevance"
|
||||
|
||||
|
||||
class TypeSafeReranker(Reranker):
|
||||
"""
|
||||
Reranks the results using the TypeSafe System One API.
|
||||
https://docs.typesafe.ai/cookbooks/rerank_typesafe
|
||||
|
||||
Each result is scored independently: TypeSafe reads the query together with
|
||||
the result's ``column`` value and answers a yes/no ("noul") question about
|
||||
whether the document is relevant. The probability of "yes", between 0 and 1,
|
||||
becomes the ``_relevance_score``. Because every score is an absolute
|
||||
probability rather than a position in the list, scores are comparable
|
||||
across queries and can be thresholded. They are model estimates, though,
|
||||
and can vary slightly between identical calls, so results with close scores
|
||||
may swap places when the same search is repeated.
|
||||
|
||||
One request is sent per result, up to ``max_concurrency`` at a time.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
model_name : str, default "jev-latest"
|
||||
The TypeSafe model to use.
|
||||
column : str, default "text"
|
||||
The name of the column holding the document text to score.
|
||||
instructions : str, optional
|
||||
The yes/no question asked about each query and document pair. The state
|
||||
TypeSafe reads is ``{"query": <query>, "document": <column value>}``.
|
||||
Defaults to a generic relevance question.
|
||||
criteria : Mapping[str, str], optional
|
||||
What a "yes" and a "no" mean, as a mapping with the keys ``"true"`` and
|
||||
``"false"``. Domain-specific criteria usually rank better than the
|
||||
generic default. Pass an empty mapping to send no criteria.
|
||||
return_score : str, default "relevance"
|
||||
Options are "relevance" or "all". If "all", keeps the vector and FTS
|
||||
scores alongside the relevance score.
|
||||
api_key : str, optional
|
||||
The API key to use. If None, the TypeSafe SDK reads the
|
||||
``TYPESAFE_API_KEY`` environment variable.
|
||||
max_concurrency : int, default 8
|
||||
The maximum number of TypeSafe requests in flight for one rerank call.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
model_name: str = "jev-latest",
|
||||
column: str = "text",
|
||||
instructions: Optional[str] = None,
|
||||
criteria: Optional[Mapping[str, str]] = None,
|
||||
return_score: str = "relevance",
|
||||
api_key: Optional[str] = None,
|
||||
max_concurrency: int = 8,
|
||||
):
|
||||
super().__init__(return_score)
|
||||
if max_concurrency < 1:
|
||||
raise ValueError("max_concurrency must be at least 1")
|
||||
criteria = DEFAULT_CRITERIA if criteria is None else dict(criteria)
|
||||
unknown = set(criteria) - {"true", "false"}
|
||||
if unknown:
|
||||
raise ValueError(
|
||||
f"criteria keys must be 'true' or 'false', got {sorted(unknown)}"
|
||||
)
|
||||
self.model_name = model_name
|
||||
self.column = column
|
||||
self.instructions = instructions or DEFAULT_INSTRUCTIONS
|
||||
self.criteria = criteria
|
||||
self.api_key = api_key
|
||||
self.max_concurrency = max_concurrency
|
||||
|
||||
def __str__(self):
|
||||
return f"TypeSafeReranker(model_name={self.model_name})"
|
||||
|
||||
@cached_property
|
||||
def _client(self):
|
||||
typesafe_sdk = attempt_import_or_raise("typesafe_sdk", "typesafe-sdk")
|
||||
return typesafe_sdk.TypeSafeClient(api_key=self.api_key)
|
||||
|
||||
@cached_property
|
||||
def _question(self) -> Dict[str, Any]:
|
||||
question: Dict[str, Any] = {"type": "noul", "instructions": self.instructions}
|
||||
if self.criteria:
|
||||
question["criteria"] = self.criteria
|
||||
return question
|
||||
|
||||
def _score(self, query: str, document: Optional[str]) -> float:
|
||||
if document is None:
|
||||
return 0.0
|
||||
response = self._client.system_one(
|
||||
state={"query": query, "document": document},
|
||||
questions={_QUESTION_ID: self._question},
|
||||
model=self.model_name,
|
||||
)
|
||||
return response.answers[_QUESTION_ID].noul
|
||||
|
||||
def _rerank(self, result_set: pa.Table, query: str) -> pa.Table:
|
||||
result_set = self._handle_empty_results(result_set)
|
||||
if len(result_set) == 0:
|
||||
return result_set
|
||||
docs = result_set[self.column].to_pylist()
|
||||
# Rerankers are also called synchronously from inside the async query
|
||||
# APIs, so the requests run on threads rather than on an event loop.
|
||||
with ThreadPoolExecutor(
|
||||
max_workers=min(self.max_concurrency, len(docs))
|
||||
) as pool:
|
||||
scores = list(pool.map(lambda doc: self._score(query, doc), docs))
|
||||
result_set = result_set.append_column(
|
||||
"_relevance_score", pa.array(scores, type=pa.float32())
|
||||
)
|
||||
return result_set.sort_by([("_relevance_score", "descending")])
|
||||
|
||||
def rerank_hybrid(
|
||||
self,
|
||||
query: str,
|
||||
vector_results: pa.Table,
|
||||
fts_results: pa.Table,
|
||||
):
|
||||
if self.score == "all":
|
||||
combined_results = self._merge_and_keep_scores(vector_results, fts_results)
|
||||
else:
|
||||
combined_results = self.merge_results(vector_results, fts_results)
|
||||
combined_results = self._rerank(combined_results, query)
|
||||
if self.score == "relevance":
|
||||
combined_results = self._keep_relevance_score(combined_results)
|
||||
return combined_results
|
||||
|
||||
def rerank_vector(self, query: str, vector_results: pa.Table):
|
||||
vector_results = self._rerank(vector_results, query)
|
||||
if self.score == "relevance":
|
||||
vector_results = vector_results.drop_columns(["_distance"])
|
||||
return vector_results
|
||||
|
||||
def rerank_fts(self, query: str, fts_results: pa.Table):
|
||||
fts_results = self._rerank(fts_results, query)
|
||||
if self.score == "relevance":
|
||||
fts_results = fts_results.drop_columns(["_score"])
|
||||
return fts_results
|
||||
@@ -0,0 +1,226 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright The LanceDB Authors
|
||||
|
||||
"""Named Secrets, and the bindings that deliver them to Functions.
|
||||
|
||||
A Secret is a database-scoped named credential. Nothing in this module holds a
|
||||
value: :class:`EnvVarSecret` names one and says which environment variable it
|
||||
should arrive in, and the value is resolved by the remote service when a
|
||||
Function bound to it runs. No API returns a stored credential, by construction
|
||||
rather than by policy -- there is no code path that could.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import re
|
||||
|
||||
# The same characters LanceDB already admits in a namespace or table name, and
|
||||
# no positional rule on top of them: a segment may begin with `_`, `-` or `.`
|
||||
# today, so anything narrower would put Secrets out of reach inside namespaces
|
||||
# that already exist. Matches the service, which admits the same set.
|
||||
_SECRET_NAME = re.compile(r"^[A-Za-z0-9_.-]{1,255}$")
|
||||
_ENV_VARIABLE = re.compile(r"^[A-Za-z_][A-Za-z0-9_]*$")
|
||||
|
||||
|
||||
def validate_secret_name(name: str) -> str:
|
||||
"""Check a Secret name locally and return it unchanged."""
|
||||
if not isinstance(name, str):
|
||||
raise TypeError(f"Secret name must be a string, not {type(name).__name__}")
|
||||
if not _SECRET_NAME.fullmatch(name):
|
||||
raise ValueError(f"invalid Secret name: {name!r}")
|
||||
return name
|
||||
|
||||
|
||||
def validate_namespace_path(namespace_path=None):
|
||||
"""Check a namespace path locally and return it as a tuple.
|
||||
|
||||
``None`` and ``[]`` both mean the root namespace. Segments follow the same
|
||||
rule as Secret names: a binding carries the path and the name as separate
|
||||
fields, so neither is ever parsed out of the other.
|
||||
"""
|
||||
if namespace_path is None:
|
||||
return ()
|
||||
if isinstance(namespace_path, str):
|
||||
raise TypeError(
|
||||
"namespace_path must be a list of segments, not a string; "
|
||||
f"did you mean [{namespace_path!r}]?"
|
||||
)
|
||||
segments = tuple(namespace_path)
|
||||
for segment in segments:
|
||||
if not isinstance(segment, str):
|
||||
raise TypeError(
|
||||
f"namespace path segment must be a string, not {type(segment).__name__}"
|
||||
)
|
||||
if not _SECRET_NAME.fullmatch(segment):
|
||||
raise ValueError(f"invalid namespace path segment: {segment!r}")
|
||||
return segments
|
||||
|
||||
|
||||
def validate_env_variable(name: str) -> str:
|
||||
"""Check an environment variable name locally and return it unchanged."""
|
||||
if not isinstance(name, str):
|
||||
raise TypeError(
|
||||
f"environment variable name must be a string, not {type(name).__name__}"
|
||||
)
|
||||
if not _ENV_VARIABLE.fullmatch(name):
|
||||
raise ValueError(f"invalid environment variable name: {name!r}")
|
||||
return name
|
||||
|
||||
|
||||
class EnvVarSecret:
|
||||
"""A Secret bound to the environment variable a Function's library reads.
|
||||
|
||||
Pass these in the ``secrets`` sequence of
|
||||
[DBConnection.create_function][lancedb.db.DBConnection.create_function]. The
|
||||
Function's source is unchanged by binding: it reads ``OPENAI_API_KEY`` the
|
||||
way it always did, and the binding is what puts a value there.
|
||||
|
||||
This is a local value. Constructing it contacts no server, so it always
|
||||
succeeds and says nothing about whether the Secret exists; that is checked
|
||||
at registration, where a mistyped Secret name surfaces as a clear "does not
|
||||
exist" naming both the Secret and the variable bound to it. A mistyped
|
||||
*variable* name cannot be caught anywhere -- nothing knows which variables a
|
||||
Function reads -- so it surfaces on the first rows instead.
|
||||
|
||||
The type exists so a credential cannot be passed by accident. A bare string
|
||||
in the same position is a plausible-looking mistake with the opposite
|
||||
meaning, and it reads identically in a diff.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
secret_name : str
|
||||
The Secret's database-scoped name.
|
||||
env_variable : str
|
||||
The environment variable the Function reads it from.
|
||||
secret_namespace_path : list of str, optional
|
||||
The namespace the Secret is addressed within. ``None`` and ``[]`` both
|
||||
mean the root namespace. Carried beside the name rather than joined
|
||||
into it, so neither is ever parsed back out of the other.
|
||||
|
||||
Examples
|
||||
--------
|
||||
>>> from lancedb import EnvVarSecret
|
||||
>>> binding = EnvVarSecret(
|
||||
... secret_name="openai-prod", env_variable="OPENAI_API_KEY"
|
||||
... )
|
||||
>>> binding.secret_name, binding.env_variable
|
||||
('openai-prod', 'OPENAI_API_KEY')
|
||||
"""
|
||||
|
||||
__slots__ = ("_secret_name", "_env_variable", "_secret_namespace_path")
|
||||
|
||||
def __init__(
|
||||
self, secret_name: str, env_variable: str, *, secret_namespace_path=None
|
||||
):
|
||||
self._secret_name = validate_secret_name(secret_name)
|
||||
self._env_variable = validate_env_variable(env_variable)
|
||||
self._secret_namespace_path = validate_namespace_path(secret_namespace_path)
|
||||
|
||||
@property
|
||||
def secret_name(self) -> str:
|
||||
"""The Secret's database-scoped name."""
|
||||
return self._secret_name
|
||||
|
||||
@property
|
||||
def env_variable(self) -> str:
|
||||
"""The environment variable the value is delivered in."""
|
||||
return self._env_variable
|
||||
|
||||
@property
|
||||
def secret_namespace_path(self):
|
||||
"""The namespace path the Secret is addressed within, root when empty."""
|
||||
return list(self._secret_namespace_path)
|
||||
|
||||
def __repr__(self) -> str:
|
||||
path = (
|
||||
f", secret_namespace_path={list(self._secret_namespace_path)!r}"
|
||||
if self._secret_namespace_path
|
||||
else ""
|
||||
)
|
||||
return (
|
||||
f"EnvVarSecret(secret_name={self._secret_name!r}, "
|
||||
f"env_variable={self._env_variable!r}{path})"
|
||||
)
|
||||
|
||||
def __eq__(self, other: object) -> bool:
|
||||
return (
|
||||
isinstance(other, EnvVarSecret)
|
||||
and other._secret_name == self._secret_name
|
||||
and other._env_variable == self._env_variable
|
||||
and other._secret_namespace_path == self._secret_namespace_path
|
||||
)
|
||||
|
||||
def __hash__(self) -> int:
|
||||
return hash(
|
||||
(
|
||||
EnvVarSecret,
|
||||
self._secret_name,
|
||||
self._env_variable,
|
||||
self._secret_namespace_path,
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
class SecretInfo:
|
||||
"""What a database records about a Secret. Never its value.
|
||||
|
||||
Returned by
|
||||
[DBConnection.describe_secret][lancedb.db.DBConnection.describe_secret].
|
||||
"""
|
||||
|
||||
__slots__ = ("_name", "_created_at_millis", "_updated_at_millis")
|
||||
|
||||
def __init__(self, name: str, created_at_millis: int, updated_at_millis: int):
|
||||
self._name = name
|
||||
self._created_at_millis = created_at_millis
|
||||
self._updated_at_millis = updated_at_millis
|
||||
|
||||
@property
|
||||
def name(self) -> str:
|
||||
"""The Secret's database-scoped name."""
|
||||
return self._name
|
||||
|
||||
@property
|
||||
def created_at_millis(self) -> int:
|
||||
"""When the Secret was created, in milliseconds since the Unix epoch."""
|
||||
return self._created_at_millis
|
||||
|
||||
@property
|
||||
def updated_at_millis(self) -> int:
|
||||
"""When the Secret's value was last rotated, in epoch milliseconds.
|
||||
|
||||
The only observable that a rotation landed: no API returns a credential,
|
||||
so a caller confirms ``alter_secret`` took effect by watching this move.
|
||||
"""
|
||||
return self._updated_at_millis
|
||||
|
||||
@classmethod
|
||||
def from_json(cls, value: dict) -> "SecretInfo":
|
||||
return cls(
|
||||
name=value["name"],
|
||||
created_at_millis=value["created_at_millis"],
|
||||
updated_at_millis=value["updated_at_millis"],
|
||||
)
|
||||
|
||||
def __repr__(self) -> str:
|
||||
return (
|
||||
f"SecretInfo(name={self._name!r}, "
|
||||
f"created_at_millis={self._created_at_millis!r}, "
|
||||
f"updated_at_millis={self._updated_at_millis!r})"
|
||||
)
|
||||
|
||||
def __eq__(self, other: object) -> bool:
|
||||
return (
|
||||
isinstance(other, SecretInfo)
|
||||
and other._name == self._name
|
||||
and other._created_at_millis == self._created_at_millis
|
||||
and other._updated_at_millis == self._updated_at_millis
|
||||
)
|
||||
|
||||
|
||||
__all__ = [
|
||||
"EnvVarSecret",
|
||||
"SecretInfo",
|
||||
"validate_env_variable",
|
||||
"validate_secret_name",
|
||||
]
|
||||
+165
-35
@@ -186,6 +186,7 @@ if TYPE_CHECKING:
|
||||
CompactionStats,
|
||||
Tag,
|
||||
AddColumnsResult,
|
||||
FunctionErrors,
|
||||
RefreshColumnResult,
|
||||
AddResult,
|
||||
AlterColumnsResult,
|
||||
@@ -634,26 +635,43 @@ def _align_field_types(
|
||||
return new_fields
|
||||
|
||||
|
||||
def _align_list_value_field(
|
||||
value_field: pa.Field, target_value_field: pa.Field
|
||||
) -> pa.Field:
|
||||
# A list has exactly one child, so the inferred child name ("item") aligns
|
||||
# positionally and adopts the table's child name; pa.Table.cast renames it.
|
||||
return _align_field(value_field, target_value_field).with_name(
|
||||
target_value_field.name
|
||||
)
|
||||
def _align_container_child(child: pa.Field, target_child: pa.Field) -> pa.Field:
|
||||
# A list has one child, a map one key and one item, so an inferred child name
|
||||
# ("item") aligns positionally and adopts the table's; pa.Table.cast renames it.
|
||||
return _align_field(child, target_child).with_name(target_child.name)
|
||||
|
||||
|
||||
def _arrow_json_storage_type(input_type: pa.DataType) -> Optional[pa.DataType]:
|
||||
"""The storage type arrow.json would use for ``input_type``.
|
||||
|
||||
Returns None if the type cannot hold JSON text.
|
||||
"""
|
||||
if pa.types.is_string(input_type) or pa.types.is_string_view(input_type):
|
||||
return pa.string()
|
||||
if pa.types.is_large_string(input_type):
|
||||
return pa.large_string()
|
||||
return None
|
||||
|
||||
|
||||
def _align_field(field: pa.Field, target_field: pa.Field) -> pa.Field:
|
||||
# Preserve arrow.json input until it reaches Lance. LanceDB exposes stored
|
||||
# JSON columns as lance.json (JSONB-backed LargeBinary), but casting the
|
||||
# input to that storage type here merely relabels the raw JSON bytes as
|
||||
# LanceDB exposes stored JSON columns as lance.json (JSONB-backed LargeBinary), but
|
||||
# casting the input to that storage type here merely relabels the raw JSON bytes as
|
||||
# JSONB. Lance must see arrow.json so it can perform the JSONB encoding.
|
||||
if (
|
||||
_field_extension_name(field) == "arrow.json"
|
||||
and _field_extension_name(target_field) == "lance.json"
|
||||
):
|
||||
return field
|
||||
if _field_extension_name(target_field) == "lance.json":
|
||||
if _field_extension_name(field) == "arrow.json":
|
||||
return field
|
||||
# Plain JSON text, which is what pyarrow infers for a column of `str`, only
|
||||
# needs the arrow.json label.
|
||||
json_storage = _arrow_json_storage_type(field.type)
|
||||
if json_storage is not None:
|
||||
# Labelled through metadata rather than pa.json_(), which only exists on
|
||||
# newer PyArrow; Lance reads the extension name off the field either way.
|
||||
return pa.field(
|
||||
field.name,
|
||||
json_storage,
|
||||
field.nullable,
|
||||
{"ARROW:extension:name": "arrow.json"},
|
||||
)
|
||||
if pa.types.is_struct(target_field.type):
|
||||
if pa.types.is_struct(field.type):
|
||||
new_type = pa.struct(
|
||||
@@ -667,7 +685,7 @@ def _align_field(field: pa.Field, target_field: pa.Field) -> pa.Field:
|
||||
elif pa.types.is_list(target_field.type):
|
||||
if _is_list_like(field.type):
|
||||
new_type = pa.list_(
|
||||
_align_list_value_field(
|
||||
_align_container_child(
|
||||
field.type.value_field, target_field.type.value_field
|
||||
)
|
||||
)
|
||||
@@ -676,7 +694,7 @@ def _align_field(field: pa.Field, target_field: pa.Field) -> pa.Field:
|
||||
elif pa.types.is_large_list(target_field.type):
|
||||
if _is_list_like(field.type):
|
||||
new_type = pa.large_list(
|
||||
_align_list_value_field(
|
||||
_align_container_child(
|
||||
field.type.value_field, target_field.type.value_field
|
||||
)
|
||||
)
|
||||
@@ -685,13 +703,28 @@ def _align_field(field: pa.Field, target_field: pa.Field) -> pa.Field:
|
||||
elif pa.types.is_fixed_size_list(target_field.type):
|
||||
if _is_list_like(field.type):
|
||||
new_type = pa.list_(
|
||||
_align_list_value_field(
|
||||
_align_container_child(
|
||||
field.type.value_field, target_field.type.value_field
|
||||
),
|
||||
target_field.type.list_size,
|
||||
)
|
||||
else:
|
||||
new_type = target_field.type
|
||||
elif pa.types.is_map(target_field.type):
|
||||
if pa.types.is_map(field.type):
|
||||
# A map has exactly one key and one item field, so like a list's child they
|
||||
# align positionally and adopt the table's names.
|
||||
new_type = pa.map_(
|
||||
_align_container_child(
|
||||
field.type.key_field, target_field.type.key_field
|
||||
),
|
||||
_align_container_child(
|
||||
field.type.item_field, target_field.type.item_field
|
||||
),
|
||||
keys_sorted=target_field.type.keys_sorted,
|
||||
)
|
||||
else:
|
||||
new_type = target_field.type
|
||||
else:
|
||||
new_type = target_field.type
|
||||
return pa.field(field.name, new_type, field.nullable, target_field.metadata)
|
||||
@@ -2188,10 +2221,10 @@ class Table(ABC):
|
||||
Declaring one therefore costs the same on a large table as on an
|
||||
empty one.
|
||||
|
||||
A refresh does not revisit rows it has already filled, so mutating
|
||||
an input leaves the value computed at fill time; recomputing means
|
||||
dropping the column and declaring it again. While a declaration
|
||||
reads a column, that column cannot be renamed, retyped or dropped.
|
||||
A refresh also recomputes the rows whose inputs changed since they
|
||||
were computed, so a mutated input is reflected by the next refresh.
|
||||
While a declaration reads a column, that column cannot be renamed,
|
||||
retyped or dropped.
|
||||
|
||||
On LanceDB Cloud and Enterprise the expression is planned by the
|
||||
server, and the refresh runs as a server job -- see
|
||||
@@ -2211,7 +2244,7 @@ class Table(ABC):
|
||||
>>> table.add_columns(computed={"doubled": "x * 2"})
|
||||
AddColumnsResult(version=2)
|
||||
>>> table.refresh_column("doubled")
|
||||
RefreshColumnResult(rows_filled=2, version=3)
|
||||
RefreshColumnResult(rows_filled=2, version=4)
|
||||
>>> table.to_arrow().sort_by("x").to_pandas()
|
||||
x doubled
|
||||
0 1 2
|
||||
@@ -2225,8 +2258,8 @@ class Table(ABC):
|
||||
|
||||
Declared with ``add_columns(computed=...)``, a column starts empty and
|
||||
gets its values here. Rows appended since the last refresh are filled
|
||||
by the next one; rows already filled are left as they are, so the call
|
||||
is idempotent and does not observe a mutated input.
|
||||
by the next one, and rows whose inputs changed since they were computed
|
||||
are recomputed; everything else is left as it is.
|
||||
|
||||
Local tables only: a remote refresh runs as a server job, through
|
||||
[`refresh_column_async`][lancedb.table.Table.refresh_column_async].
|
||||
@@ -2276,6 +2309,47 @@ class Table(ABC):
|
||||
'finished'
|
||||
"""
|
||||
|
||||
@abstractmethod
|
||||
def function_errors(
|
||||
self,
|
||||
job_id: Optional[str] = None,
|
||||
column: Optional[str] = None,
|
||||
limit: Optional[int] = None,
|
||||
) -> "FunctionErrors":
|
||||
"""
|
||||
The per-row errors Function refreshes recorded on this table.
|
||||
|
||||
A refresh running under a skip policy records each row it skipped
|
||||
with the input that failed and the error. This lists those records,
|
||||
newest job first, plus a summary for any fragment whose per-row
|
||||
detail was capped. LanceDB Cloud and Enterprise only; reading errors
|
||||
needs read access to the table, since a message carries the value
|
||||
that failed.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
job_id: str, optional
|
||||
Only errors recorded by this job.
|
||||
column: str, optional
|
||||
Only errors on this column.
|
||||
limit: int, optional
|
||||
At most this many records (server default 10000, cap 100000).
|
||||
|
||||
Returns
|
||||
-------
|
||||
FunctionErrors
|
||||
``records``, ``fragments`` and ``truncated``, the last saying
|
||||
whether the listing stopped at its limit.
|
||||
|
||||
Examples
|
||||
--------
|
||||
>>> errors = table.function_errors(column="embedding") # doctest: +SKIP
|
||||
>>> for record in errors.records: # doctest: +SKIP
|
||||
... print(record.job_id, record.row_offset, record.error_message)
|
||||
>>> if errors.truncated: # doctest: +SKIP
|
||||
... print("listing stopped at the limit")
|
||||
"""
|
||||
|
||||
@abstractmethod
|
||||
def alter_columns(self, *alterations: Iterable[Dict[str, str]]):
|
||||
"""
|
||||
@@ -4318,17 +4392,30 @@ class LanceTable(Table):
|
||||
return LOOP.run(self._table.add_columns(transforms, computed=computed))
|
||||
|
||||
def refresh_column(self, column: str) -> "RefreshColumnResult":
|
||||
"""Fill a computed column's unfilled rows. See
|
||||
"""Fill a computed column's unfilled rows and recompute those whose
|
||||
inputs changed. See
|
||||
[`AsyncTable.refresh_column`][lancedb.AsyncTable.refresh_column]."""
|
||||
return LOOP.run(self._table.refresh_column(column))
|
||||
|
||||
def refresh_column_async(self, column: str) -> Job[RefreshColumnJobResult]:
|
||||
"""Fill a computed column's unfilled rows, returning a handle to the
|
||||
refresh job. See
|
||||
"""Fill a computed column's unfilled rows and recompute those whose
|
||||
inputs changed, returning a handle to the refresh job. See
|
||||
[`Table.refresh_column_async`][lancedb.table.Table.refresh_column_async].
|
||||
"""
|
||||
return Job(LOOP.run(self._table.refresh_column_async(column)))
|
||||
|
||||
def function_errors(
|
||||
self,
|
||||
job_id: Optional[str] = None,
|
||||
column: Optional[str] = None,
|
||||
limit: Optional[int] = None,
|
||||
) -> "FunctionErrors":
|
||||
"""The per-row errors Function refreshes recorded on this table. See
|
||||
[`Table.function_errors`][lancedb.table.Table.function_errors]."""
|
||||
return LOOP.run(
|
||||
self._table.function_errors(job_id=job_id, column=column, limit=limit)
|
||||
)
|
||||
|
||||
def alter_columns(
|
||||
self, *alterations: Iterable[Dict[str, str]]
|
||||
) -> AlterColumnsResult:
|
||||
@@ -6312,10 +6399,10 @@ class AsyncTable:
|
||||
them from
|
||||
[`refresh_column`][lancedb.table.AsyncTable.refresh_column].
|
||||
|
||||
A refresh does not revisit rows it has already filled, so mutating
|
||||
an input leaves the value computed at fill time. While a
|
||||
declaration reads a column, that column cannot be renamed, retyped
|
||||
or dropped.
|
||||
A refresh also recomputes the rows whose inputs changed since they
|
||||
were computed, so a mutated input is reflected by the next refresh.
|
||||
While a declaration reads a column, that column cannot be renamed,
|
||||
retyped or dropped.
|
||||
|
||||
On LanceDB Cloud and Enterprise the expression is planned by
|
||||
the server. Cannot be combined with ``transforms``.
|
||||
@@ -6377,8 +6464,8 @@ class AsyncTable:
|
||||
|
||||
Declared with ``add_columns(computed=...)``, a column starts empty and
|
||||
gets its values here. Rows appended since the last refresh are filled
|
||||
by the next one; rows already filled are left as they are, so the call
|
||||
is idempotent and does not observe a mutated input.
|
||||
by the next one, and rows whose inputs changed since they were computed
|
||||
are recomputed; everything else is left as it is.
|
||||
|
||||
Local tables only: a remote refresh runs as a server job, through
|
||||
[`refresh_column_async`][lancedb.table.Table.refresh_column_async].
|
||||
@@ -6395,6 +6482,49 @@ class AsyncTable:
|
||||
"""
|
||||
return await self._inner.refresh_column(column)
|
||||
|
||||
async def function_errors(
|
||||
self,
|
||||
job_id: Optional[str] = None,
|
||||
column: Optional[str] = None,
|
||||
limit: Optional[int] = None,
|
||||
) -> "FunctionErrors":
|
||||
"""
|
||||
The per-row errors Function refreshes recorded on this table.
|
||||
|
||||
A refresh running under a skip policy records each row it skipped
|
||||
with the input that failed and the error. This lists those records,
|
||||
newest job first, plus a summary for any fragment whose per-row
|
||||
detail was capped. LanceDB Cloud and Enterprise only; reading errors
|
||||
needs read access to the table, since a message carries the value
|
||||
that failed.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
job_id: str, optional
|
||||
Only errors recorded by this job.
|
||||
column: str, optional
|
||||
Only errors on this column.
|
||||
limit: int, optional
|
||||
At most this many records (server default 10000, cap 100000).
|
||||
|
||||
Returns
|
||||
-------
|
||||
FunctionErrors
|
||||
``records``, ``fragments`` and ``truncated``, the last saying
|
||||
whether the listing stopped at its limit.
|
||||
|
||||
Examples
|
||||
--------
|
||||
>>> errors = await table.function_errors(column="embedding") # doctest: +SKIP
|
||||
>>> for record in errors.records: # doctest: +SKIP
|
||||
... print(record.job_id, record.row_offset, record.error_message)
|
||||
>>> if errors.truncated: # doctest: +SKIP
|
||||
... print("listing stopped at the limit")
|
||||
"""
|
||||
return await self._inner.function_errors(
|
||||
job_id=job_id, column=column, limit=limit
|
||||
)
|
||||
|
||||
async def refresh_column_async(
|
||||
self, column: str
|
||||
) -> AsyncJob[RefreshColumnJobResult]:
|
||||
|
||||
@@ -751,6 +751,80 @@ def test_fetch_blobs_preserves_null_and_empty_values():
|
||||
assert blobs[3].as_py() == b"present"
|
||||
|
||||
|
||||
def test_add_all_null_list_to_blob_column():
|
||||
table = _blob_table("all_null_add", [{"id": 1, "image": None}])
|
||||
|
||||
hits = table.search().to_arrow()
|
||||
blobs = table.fetch_blobs("image", hits)
|
||||
assert len(blobs) == 1
|
||||
assert blobs[0].as_py() is None
|
||||
|
||||
|
||||
def test_add_all_null_list_to_blob_column_with_sanitizer():
|
||||
db = lancedb.connect("memory:///")
|
||||
schema = pa.schema([pa.field("id", pa.int64()), lancedb.blob("image")])
|
||||
table = db.create_table("all_null_sanitized_add", schema=schema)
|
||||
|
||||
table.add([{"id": 1, "image": None}], on_bad_vectors="fill")
|
||||
|
||||
hits = table.search().to_arrow()
|
||||
blobs = table.fetch_blobs("image", hits)
|
||||
assert len(blobs) == 1
|
||||
assert blobs[0].as_py() is None
|
||||
|
||||
|
||||
def test_add_all_null_list_to_nested_blob_column():
|
||||
db = lancedb.connect("memory:///")
|
||||
blob_field = lancedb.blob("image")
|
||||
info_field = pa.field("info", pa.struct([blob_field]))
|
||||
info = pa.StructArray.from_arrays(
|
||||
[_blob_array("image", [b"seed"])], fields=[blob_field]
|
||||
)
|
||||
seed = pa.Table.from_arrays(
|
||||
[pa.array([0], type=pa.int64()), info],
|
||||
schema=pa.schema([pa.field("id", pa.int64()), info_field]),
|
||||
)
|
||||
table = db.create_table("nested_null_add", data=seed)
|
||||
|
||||
table.add([{"id": 1, "info": {"image": None}}])
|
||||
table.add([{"id": 2, "info": {"image": None}}], on_bad_vectors="fill")
|
||||
|
||||
hits = table.search().where("id > 0").to_arrow()
|
||||
blobs = table.fetch_blobs("info.image", hits)
|
||||
assert len(blobs) == 2
|
||||
assert all(blob.as_py() is None for blob in blobs)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("large_list", [False, True], ids=["list", "large_list"])
|
||||
def test_add_list_of_dicts_to_blob_list_column(large_list):
|
||||
db = lancedb.connect("memory:///")
|
||||
blob_field = lancedb.blob("image")
|
||||
blob_values = _blob_array("image", [b"seed"])
|
||||
if large_list:
|
||||
items_field = pa.field("items", pa.large_list(blob_field))
|
||||
items = pa.LargeListArray.from_arrays(
|
||||
pa.array([0, 1], type=pa.int64()), blob_values
|
||||
)
|
||||
else:
|
||||
items_field = pa.field("items", pa.list_(blob_field))
|
||||
items = pa.ListArray.from_arrays(pa.array([0, 1], type=pa.int32()), blob_values)
|
||||
seed = pa.Table.from_arrays(
|
||||
[pa.array([0], type=pa.int64()), items],
|
||||
schema=pa.schema([pa.field("id", pa.int64()), items_field]),
|
||||
)
|
||||
table = db.create_table(f"blob_{large_list}_list_add", data=seed)
|
||||
|
||||
table.add([{"id": 1, "items": [None]}])
|
||||
table.add(
|
||||
[{"id": 2, "items": [b"a", None]}],
|
||||
on_bad_vectors="fill",
|
||||
)
|
||||
|
||||
ids = table.search().select(["id"]).to_arrow()["id"].to_pylist()
|
||||
assert sorted(ids) == [0, 1, 2]
|
||||
assert pa.types.is_large_list(table.schema.field("items").type) is large_list
|
||||
|
||||
|
||||
def test_fetch_blob_ranges_aligns_repeated_ranges_and_nulls():
|
||||
table = _blob_table(
|
||||
"range_alignment",
|
||||
|
||||
@@ -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/$/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)
|
||||
@@ -13,7 +13,9 @@ from lancedb.functions import (
|
||||
FunctionBinding,
|
||||
FunctionVersion,
|
||||
PythonRuntimeSpec,
|
||||
SecretBinding,
|
||||
RefreshColumnResult,
|
||||
SecretReference,
|
||||
)
|
||||
from lancedb.table import AsyncTable
|
||||
|
||||
@@ -37,6 +39,22 @@ def job_result(name: str) -> dict:
|
||||
return json.loads(fixture(name))["result"]
|
||||
|
||||
|
||||
def assert_no_secret_values(value):
|
||||
"""No client value models a resolved credential, at any nesting depth."""
|
||||
if isinstance(value, dict):
|
||||
for key, child in value.items():
|
||||
assert key not in {
|
||||
"secret_value",
|
||||
"secret_values",
|
||||
"resolved_secret",
|
||||
"resolved_secrets",
|
||||
}
|
||||
assert_no_secret_values(child)
|
||||
elif isinstance(value, list):
|
||||
for child in value:
|
||||
assert_no_secret_values(child)
|
||||
|
||||
|
||||
def test_public_function_values_are_in_api_reference():
|
||||
docs = Path(__file__).parents[3] / "docs" / "src" / "python" / "python.md"
|
||||
rendered = docs.read_text()
|
||||
@@ -93,16 +111,27 @@ def test_function_version_identity_is_immutable_and_exact():
|
||||
value = job_result("remote_function_job.json")
|
||||
version = FunctionVersion.from_json(json.dumps(value))
|
||||
assert version.name == "embed"
|
||||
assert version.version == "fv_01K3EXACT"
|
||||
assert version.version == "1"
|
||||
assert version.image.manifest_digest.startswith("sha256:")
|
||||
assert version.version != version.image.manifest_digest
|
||||
assert list(version.secret_bindings) == [
|
||||
SecretBinding(
|
||||
kind="env", variable="HF_TOKEN", secret_ref=SecretReference(name="hf-prod")
|
||||
)
|
||||
]
|
||||
|
||||
with pytest.raises((TypeError, ValueError)):
|
||||
version.version = "fv_changed"
|
||||
version.version = "1"
|
||||
with pytest.raises(TypeError, match="immutable"):
|
||||
version.runtime.env["TOKENIZERS_PARALLELISM"] = "true"
|
||||
version.image.descriptor["format_version"] = "changed"
|
||||
|
||||
changed = dict(value)
|
||||
changed["version"] = "fv_changed"
|
||||
changed["version"] = "2"
|
||||
assert FunctionVersion(**changed) != version
|
||||
assert FunctionVersion(**changed).image == version.image
|
||||
for invalid in [version.image.manifest_digest, "0", "01", "-1", str(2**64)]:
|
||||
with pytest.raises(ValueError):
|
||||
FunctionVersion(**{**value, "version": invalid})
|
||||
|
||||
|
||||
def test_function_version_binds_named_columns_as_one_immutable_application():
|
||||
@@ -137,7 +166,7 @@ def test_function_version_binding_validates_names_and_direct_columns():
|
||||
def test_function_version_keeps_named_struct_outputs_in_one_application():
|
||||
value = job_result("remote_function_job.json")
|
||||
value["name"] = "text_features"
|
||||
value["version"] = "fv_multi_output"
|
||||
value["version"] = "1"
|
||||
value["signature"] = {
|
||||
"inputs": [
|
||||
{"name": "title", "arrow_type": "utf8", "nullable": True},
|
||||
@@ -182,14 +211,14 @@ def test_function_version_keeps_named_struct_outputs_in_one_application():
|
||||
def test_unknown_fields_and_discriminators_are_forward_decodable():
|
||||
value = job_result("remote_function_job.json")
|
||||
value["future_version_metadata"] = {"retention_class": "catalog"}
|
||||
value["runtime"] = {"kind": "wasm", "module_digest": "sha256:wasm"}
|
||||
value["image"]["descriptor"]["future_interface"] = {"kind": "wasm"}
|
||||
value["signature"]["output"]["kind"] = "future_output_shape"
|
||||
|
||||
version = FunctionVersion.from_json(json.dumps(value))
|
||||
assert version.runtime.kind == "wasm"
|
||||
assert version.runtime.python_version is None
|
||||
assert version.runtime.environment is None
|
||||
assert json.loads(version.to_canonical_json())["runtime"] == {"kind": "wasm"}
|
||||
assert version.image.descriptor["future_interface"] == {"kind": "wasm"}
|
||||
assert json.loads(version.to_canonical_json())["image"]["descriptor"][
|
||||
"future_interface"
|
||||
] == {"kind": "wasm"}
|
||||
assert version.signature.output.kind == "future_output_shape"
|
||||
|
||||
|
||||
@@ -222,7 +251,7 @@ def test_function_application_uses_rename_columns_only():
|
||||
|
||||
def test_binding_and_refresh_result_keep_stable_remote_fields():
|
||||
binding = FunctionBinding.from_json(fixture("remote_function_binding.json"))
|
||||
assert binding.function.version == "fv_01K3TEXT"
|
||||
assert binding.function.version == "1"
|
||||
assert [output.output_ordinal for output in binding.outputs] == [0, 1]
|
||||
assert binding.input_schema is not None
|
||||
assert binding.output_schema is not None
|
||||
@@ -276,6 +305,31 @@ def test_refresh_result_rejects_non_u64_values(field):
|
||||
RefreshColumnResult.from_json(json.dumps(value))
|
||||
|
||||
|
||||
def test_canonical_client_values_carry_bindings_and_no_credentials():
|
||||
"""A binding names a Secret; the credential behind it has no client field."""
|
||||
version = FunctionVersion.from_json(
|
||||
json.dumps(job_result("remote_function_job.json"))
|
||||
)
|
||||
canonical = json.loads(version.to_canonical_json())
|
||||
assert canonical["secret_bindings"] == [
|
||||
{"kind": "env", "variable": "HF_TOKEN", "secret_ref": {"name": "hf-prod"}}
|
||||
]
|
||||
assert_no_secret_values(canonical)
|
||||
|
||||
|
||||
def test_a_version_without_bindings_omits_the_field_in_both_directions():
|
||||
"""A Function that binds nothing carries no ``secret_bindings`` key.
|
||||
|
||||
Absent decodes as an empty list, and an empty list serializes back to
|
||||
absent.
|
||||
"""
|
||||
value = job_result("remote_function_job.json")
|
||||
del value["secret_bindings"]
|
||||
version = FunctionVersion.from_json(json.dumps(value))
|
||||
assert list(version.secret_bindings) == []
|
||||
assert "secret_bindings" not in json.loads(version.to_canonical_json())
|
||||
|
||||
|
||||
class _FunctionDeclarationInner:
|
||||
def __init__(self):
|
||||
self.calls = []
|
||||
@@ -339,7 +393,16 @@ def test_rename_requires_named_struct_and_keeps_partial_mapping_immutable():
|
||||
scalar = FunctionApplication.from_json(
|
||||
json.dumps(
|
||||
{
|
||||
"function": {"name": "embed", "version": "fv_exact"},
|
||||
"function": {
|
||||
"name": "embed",
|
||||
"version": "1",
|
||||
"object_id": "fixture",
|
||||
"location": "memory:///fixture",
|
||||
"manifest_digest": (
|
||||
"sha256:"
|
||||
"7e22f815b6648e14f093a3979a8e5a2082fa773ebe1ec84b135cae7e84d6f8e6"
|
||||
),
|
||||
},
|
||||
"inputs": [],
|
||||
"output": {
|
||||
"kind": "scalar",
|
||||
|
||||
@@ -11,6 +11,7 @@ import types
|
||||
from datetime import date
|
||||
import http.server
|
||||
import json
|
||||
import os
|
||||
from pathlib import Path
|
||||
import subprocess
|
||||
import sys
|
||||
@@ -24,11 +25,14 @@ import pytest
|
||||
import lancedb
|
||||
from lancedb.functions import (
|
||||
PythonRuntimeSpec,
|
||||
SecretBinding,
|
||||
SecretReference,
|
||||
UdfDefinition,
|
||||
_canonical_arrow_type,
|
||||
_GRAMMAR_PRIMITIVES,
|
||||
udf,
|
||||
)
|
||||
from lancedb.secrets import EnvVarSecret
|
||||
|
||||
THRESHOLD = 20
|
||||
_CACHE = None
|
||||
@@ -45,6 +49,11 @@ FIXTURES = (
|
||||
)
|
||||
|
||||
|
||||
FUNCTION_VERSION = json.loads(
|
||||
(FIXTURES / "remote_function_version.canonical.json").read_text()
|
||||
)["version"]
|
||||
|
||||
|
||||
@udf(
|
||||
pip=["numpy>=2"],
|
||||
env={"MODE": "test"},
|
||||
@@ -54,6 +63,15 @@ def normalize_score(value: float) -> float:
|
||||
return value / 100.0
|
||||
|
||||
|
||||
@udf(
|
||||
pip=["openai==3.7.0"],
|
||||
env={"MODE": "test"},
|
||||
python_version="3.12",
|
||||
)
|
||||
def analyze_caption(caption: str) -> str:
|
||||
return caption.strip()
|
||||
|
||||
|
||||
def test_scalar_udf_matches_shared_registration_golden_and_remains_callable():
|
||||
assert isinstance(normalize_score, UdfDefinition)
|
||||
assert normalize_score(25.0) == 0.25
|
||||
@@ -70,6 +88,286 @@ def test_scalar_udf_matches_shared_registration_golden_and_remains_callable():
|
||||
}
|
||||
|
||||
|
||||
def test_secret_bound_udf_matches_its_shared_registration_golden():
|
||||
assert analyze_caption(" hello ") == "hello"
|
||||
bound = analyze_caption.bind_secrets(
|
||||
[EnvVarSecret(secret_name="openai-prod", env_variable="OPENAI_API_KEY")]
|
||||
)
|
||||
assert (
|
||||
bound.to_canonical_json()
|
||||
== (FIXTURES / "remote_function_secret_registration_request.canonical.json")
|
||||
.read_text()
|
||||
.strip()
|
||||
)
|
||||
|
||||
|
||||
def test_a_namespaced_binding_records_the_path_and_the_name():
|
||||
"""A binding names the parts, so nothing has to be parsed back out.
|
||||
|
||||
A root binding carries no path at all: the field is absent rather than an
|
||||
empty list, so a binding states a namespace only when it has one.
|
||||
"""
|
||||
root = EnvVarSecret(secret_name="openai-prod", env_variable="OPENAI_API_KEY")
|
||||
assert root.secret_namespace_path == []
|
||||
|
||||
nested = EnvVarSecret(
|
||||
secret_name="openai-prod",
|
||||
env_variable="OPENAI_API_KEY",
|
||||
secret_namespace_path=["prod", "vision"],
|
||||
)
|
||||
assert nested.secret_namespace_path == ["prod", "vision"]
|
||||
assert nested != root
|
||||
|
||||
bound = analyze_caption.bind_secrets([nested])
|
||||
assert list(bound.secret_bindings) == [
|
||||
SecretBinding(
|
||||
kind="env",
|
||||
variable="OPENAI_API_KEY",
|
||||
secret_ref=SecretReference(
|
||||
name="openai-prod", namespace_path=("prod", "vision")
|
||||
),
|
||||
)
|
||||
]
|
||||
|
||||
at_root = analyze_caption.bind_secrets([root])
|
||||
assert list(at_root.secret_bindings) == [
|
||||
SecretBinding(
|
||||
kind="env",
|
||||
variable="OPENAI_API_KEY",
|
||||
secret_ref=SecretReference(name="openai-prod"),
|
||||
)
|
||||
]
|
||||
# A root binding carries no path at all on the wire.
|
||||
canonical = json.loads(at_root.to_canonical_json())
|
||||
assert canonical["secret_bindings"] == [
|
||||
{
|
||||
"kind": "env",
|
||||
"variable": "OPENAI_API_KEY",
|
||||
"secret_ref": {"name": "openai-prod"},
|
||||
}
|
||||
]
|
||||
|
||||
|
||||
def test_a_namespace_path_is_validated_locally():
|
||||
# The charset is the service's, not a delimiter's: a reference is never
|
||||
# joined, so a segment cannot make anything parse two ways.
|
||||
with pytest.raises(ValueError):
|
||||
EnvVarSecret(
|
||||
secret_name="openai-prod",
|
||||
env_variable="K",
|
||||
secret_namespace_path=["with$delim"],
|
||||
)
|
||||
with pytest.raises(ValueError):
|
||||
EnvVarSecret(
|
||||
secret_name="openai-prod", env_variable="K", secret_namespace_path=["a/b"]
|
||||
)
|
||||
# A bare string is a plausible mistake with the wrong meaning.
|
||||
with pytest.raises(TypeError):
|
||||
EnvVarSecret(
|
||||
secret_name="openai-prod", env_variable="K", secret_namespace_path="prod"
|
||||
)
|
||||
|
||||
|
||||
def test_an_unbound_request_carries_no_binding_at_all():
|
||||
"""Binding is a registration-time decision, so the definition holds none.
|
||||
|
||||
The decorator declares nothing about secrets, which is what makes the PRD's
|
||||
claim true: a Function's source and its registration request are identical
|
||||
whether or not a credential is later bound to it.
|
||||
"""
|
||||
unbound = json.loads(analyze_caption.registration_request.to_canonical_json())
|
||||
assert "secret_bindings" not in unbound
|
||||
assert "OPENAI_API_KEY" not in json.dumps(unbound)
|
||||
|
||||
|
||||
def test_binding_a_secret_leaves_the_packaged_artifact_untouched():
|
||||
"""The artifact is source bytes and nothing else, with or without secrets."""
|
||||
bound = analyze_caption.bind_secrets(
|
||||
[EnvVarSecret(secret_name="openai-prod", env_variable="OPENAI_API_KEY")]
|
||||
)
|
||||
assert bound.artifact == analyze_caption.registration_request.artifact
|
||||
assert bound.artifact.digest == analyze_caption.registration_request.artifact.digest
|
||||
|
||||
|
||||
def test_a_function_declaring_no_secret_is_registered_exactly_as_before():
|
||||
"""The compatibility claim: nothing about the no-secret path moves."""
|
||||
assert (
|
||||
normalize_score.bind_secrets(None).to_canonical_json()
|
||||
== normalize_score.registration_request.to_canonical_json()
|
||||
)
|
||||
assert (
|
||||
"secret_bindings"
|
||||
not in normalize_score.registration_request.to_canonical_json()
|
||||
)
|
||||
|
||||
|
||||
def test_a_function_binds_each_variable_once():
|
||||
with pytest.raises(ValueError, match="binds each environment variable once"):
|
||||
analyze_caption.bind_secrets(
|
||||
[
|
||||
EnvVarSecret(secret_name="openai-prod", env_variable="OPENAI_API_KEY"),
|
||||
EnvVarSecret(
|
||||
secret_name="openai-staging", env_variable="OPENAI_API_KEY"
|
||||
),
|
||||
]
|
||||
)
|
||||
|
||||
|
||||
def test_bindings_may_not_collide_with_plain_configuration():
|
||||
"""`env` is stored with the Function; a Secret is not. Refuse, do not pick."""
|
||||
with pytest.raises(ValueError, match="must be disjoint"):
|
||||
analyze_caption.bind_secrets(
|
||||
[EnvVarSecret(secret_name="mode-prod", env_variable="MODE")]
|
||||
)
|
||||
|
||||
|
||||
def test_a_binding_envelope_reaches_the_service_for_it_to_judge():
|
||||
"""Binding rules are the service's: it owns the runtime the names land in.
|
||||
|
||||
The client sends what it was given, so a rule it duplicated could disagree
|
||||
with the service's without either side noticing. What is checked here is
|
||||
that the envelope arrives intact -- the shape the service judges is the
|
||||
shape the caller wrote.
|
||||
"""
|
||||
with _mock_remote_function_catalog() as (host, state):
|
||||
db = lancedb.connect(
|
||||
"db://dev",
|
||||
api_key="fake",
|
||||
host_override=host,
|
||||
client_config={"retry_config": {"retries": 0}},
|
||||
)
|
||||
bindings = [
|
||||
EnvVarSecret(secret_name=f"secret-{index}", env_variable=f"TOKEN_{index}")
|
||||
for index in range(17)
|
||||
]
|
||||
db.create_function(normalize_score, secrets=bindings)
|
||||
|
||||
sent = state["requests"][0][1]
|
||||
assert len(sent["secret_bindings"]) == 17
|
||||
assert {
|
||||
"kind": "env",
|
||||
"variable": "TOKEN_0",
|
||||
"secret_ref": {"name": "secret-0"},
|
||||
} in sent["secret_bindings"]
|
||||
|
||||
|
||||
_SECRET_DEBUG_LOG_SOURCE = """
|
||||
import http.server
|
||||
import json
|
||||
import threading
|
||||
|
||||
import lancedb
|
||||
|
||||
|
||||
class Handler(http.server.BaseHTTPRequestHandler):
|
||||
def log_message(self, *args):
|
||||
pass
|
||||
|
||||
def do_POST(self):
|
||||
self.rfile.read(int(self.headers.get("Content-Length", "0")))
|
||||
payload = json.dumps({}).encode()
|
||||
self.send_response(200)
|
||||
self.send_header("Content-Type", "application/json")
|
||||
self.send_header("Content-Length", str(len(payload)))
|
||||
self.end_headers()
|
||||
self.wfile.write(payload)
|
||||
|
||||
|
||||
server = http.server.ThreadingHTTPServer(("localhost", 0), Handler)
|
||||
threading.Thread(target=server.serve_forever, daemon=True).start()
|
||||
try:
|
||||
db = lancedb.connect(
|
||||
"db://dev",
|
||||
api_key="API_KEY_SENTINEL",
|
||||
host_override="http://localhost:%d" % server.server_address[1],
|
||||
client_config={"retry_config": {"retries": 0}},
|
||||
)
|
||||
db.create_secret("openai-prod", "SECRET_VALUE_SENTINEL")
|
||||
finally:
|
||||
server.shutdown()
|
||||
"""
|
||||
|
||||
|
||||
def test_a_credential_never_reaches_a_debug_log(tmp_path):
|
||||
"""The logger sees the serialized body, so no value-side redaction reaches it.
|
||||
|
||||
Runs in a subprocess because the Rust logger reads ``LANCEDB_LOG`` once, at
|
||||
import.
|
||||
"""
|
||||
script = tmp_path / "write_secret.py"
|
||||
script.write_text(_SECRET_DEBUG_LOG_SOURCE)
|
||||
|
||||
result = subprocess.run(
|
||||
[sys.executable, str(script)],
|
||||
check=True,
|
||||
capture_output=True,
|
||||
text=True,
|
||||
env={**os.environ, "LANCEDB_LOG": "debug"},
|
||||
)
|
||||
output = result.stdout + result.stderr
|
||||
|
||||
# Without this the test passes when debug logging is simply off.
|
||||
assert "Sending request_id=" in output, output
|
||||
assert "SECRET_VALUE_SENTINEL" not in output
|
||||
assert "API_KEY_SENTINEL" not in output
|
||||
|
||||
|
||||
def test_a_credential_value_is_rejected_in_the_binding_position():
|
||||
"""The one mistake the typed binding exists to stop."""
|
||||
with pytest.raises(TypeError, match="EnvVarSecret"):
|
||||
analyze_caption.bind_secrets(["sk-live-0001"])
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("secret", "variable", "message"),
|
||||
[
|
||||
("openai-prod", "not-a-var", "invalid environment variable name"),
|
||||
("openai-prod", "API-TOKEN", "invalid environment variable name"),
|
||||
("not a name", "API_TOKEN", "invalid Secret name"),
|
||||
("openai$prod", "API_TOKEN", "invalid Secret name"),
|
||||
],
|
||||
)
|
||||
def test_a_binding_validates_both_names_locally(secret, variable, message):
|
||||
with pytest.raises(ValueError, match=message):
|
||||
EnvVarSecret(secret_name=secret, env_variable=variable)
|
||||
|
||||
|
||||
def test_a_secret_name_admits_what_a_namespace_name_does():
|
||||
"""A Secret has to be nameable wherever a namespace already is.
|
||||
|
||||
LanceDB namespace and table names are `[A-Za-z0-9_.-]` with no rule about
|
||||
which character comes first, so a name may lead with `_`, `-` or `.`.
|
||||
Anything narrower here would leave Secrets unaddressable inside namespaces
|
||||
that already exist -- the reason periods are admitted is the reason the
|
||||
edges are too.
|
||||
"""
|
||||
for name in ["openai.prod.v1", ".hidden", "_internal", "-lead", "trailing."]:
|
||||
binding = EnvVarSecret(secret_name=name, env_variable="OPENAI_API_KEY")
|
||||
assert binding.secret_name == name
|
||||
|
||||
for name in ["", "with/slash", "with$delimiter", "a" * 256]:
|
||||
with pytest.raises(ValueError, match="invalid Secret name"):
|
||||
EnvVarSecret(secret_name=name, env_variable="OPENAI_API_KEY")
|
||||
|
||||
# A namespace segment follows the same rule, and LanceDB already admits
|
||||
# these shapes as namespace names -- so a Secret is addressable inside one.
|
||||
for segment in [".hidden", "_internal", "-lead", "trailing."]:
|
||||
binding = EnvVarSecret(
|
||||
secret_name="openai-prod",
|
||||
env_variable="OPENAI_API_KEY",
|
||||
secret_namespace_path=[segment],
|
||||
)
|
||||
assert binding.secret_namespace_path == [segment]
|
||||
|
||||
for segment in ["", "with/slash", "with$delimiter"]:
|
||||
with pytest.raises(ValueError, match="invalid namespace path segment"):
|
||||
EnvVarSecret(
|
||||
secret_name="openai-prod",
|
||||
env_variable="OPENAI_API_KEY",
|
||||
secret_namespace_path=[segment],
|
||||
)
|
||||
|
||||
|
||||
def _main_udf_source(
|
||||
*, threshold: int = 20, input_annotation: str = "int", comparison: str = ">="
|
||||
) -> str:
|
||||
@@ -1199,11 +1497,11 @@ def test_local_function_catalog_operations_are_not_supported(tmp_path):
|
||||
with pytest.raises(NotImplementedError, match=message):
|
||||
db.create_function_async(normalize_score)
|
||||
with pytest.raises(NotImplementedError, match=message):
|
||||
db.get_function("normalize_score", version="fv_exact")
|
||||
db.get_function("normalize_score", version=FUNCTION_VERSION)
|
||||
with pytest.raises(NotImplementedError, match=message):
|
||||
db.list_functions()
|
||||
with pytest.raises(NotImplementedError, match=message):
|
||||
db.drop_function("normalize_score", version="fv_exact")
|
||||
db.drop_function("normalize_score", version=FUNCTION_VERSION)
|
||||
|
||||
|
||||
@contextlib.contextmanager
|
||||
@@ -1227,18 +1525,34 @@ def _mock_remote_function_catalog():
|
||||
body = json.loads(self.rfile.read(length) or b"{}")
|
||||
state["requests"].append((self.path, body))
|
||||
status = 200
|
||||
if self.path == "/v1/function/normalize_score/create":
|
||||
# `{id}` is the Function name, so match on the shape rather than on
|
||||
# one name: these tests register more than one Function.
|
||||
parts = self.path.strip("/").split("/")
|
||||
function_action = (
|
||||
(urllib.parse.unquote(parts[2]), parts[3])
|
||||
if len(parts) == 4 and parts[0] == "v1" and parts[1] == "function"
|
||||
else (None, None)
|
||||
)
|
||||
secret_action = (
|
||||
(urllib.parse.unquote(parts[2]), parts[3])
|
||||
if len(parts) == 4 and parts[0] == "v1" and parts[1] == "secret"
|
||||
else (None, None)
|
||||
)
|
||||
if function_action[1] == "create":
|
||||
state["version"] = {
|
||||
"name": "normalize_score",
|
||||
"version": "fv_exact",
|
||||
"artifact": {
|
||||
key: body["artifact"][key]
|
||||
for key in ("kind", "digest", "entrypoint")
|
||||
},
|
||||
"version": FUNCTION_VERSION,
|
||||
"object_id": "fixture",
|
||||
"location": "memory:///fixture",
|
||||
"metadata": {},
|
||||
"disabled": False,
|
||||
"image": json.loads(
|
||||
(
|
||||
FIXTURES / "remote_function_version.canonical.json"
|
||||
).read_text()
|
||||
)["image"],
|
||||
"signature": body["signature"],
|
||||
"runtime": body["runtime"],
|
||||
"runtime_digest": "sha256:runtime",
|
||||
"environment_digest": "sha256:environment",
|
||||
"secret_bindings": body.get("secret_bindings", []),
|
||||
"created_at": "2026-08-21T00:00:00Z",
|
||||
}
|
||||
response = {"job_id": "job-register"}
|
||||
@@ -1252,11 +1566,20 @@ def _mock_remote_function_catalog():
|
||||
"result": state["version"],
|
||||
}
|
||||
elif self.path == "/v1/function/normalize_score/describe":
|
||||
assert body == {"version": "fv_exact"}
|
||||
assert body == {"version": FUNCTION_VERSION}
|
||||
response = state["version"]
|
||||
elif self.path == "/v1/function/normalize_score/drop":
|
||||
assert body == {"version": "fv_exact"}
|
||||
assert body == {"version": FUNCTION_VERSION}
|
||||
response = {"dropped": True}
|
||||
elif secret_action[1] in ("create", "alter"):
|
||||
# The Secret is the path identifier, so the body is the value.
|
||||
assert set(body) == {"value"}
|
||||
assert secret_action[0] == "openai-prod"
|
||||
response = {}
|
||||
elif secret_action[1] == "drop":
|
||||
assert secret_action[0] == "openai-prod"
|
||||
assert body == {}
|
||||
response = {}
|
||||
else:
|
||||
status = 404
|
||||
response = {"error": "not found"}
|
||||
@@ -1269,6 +1592,16 @@ def _mock_remote_function_catalog():
|
||||
for key, values in urllib.parse.parse_qs(url.query).items()
|
||||
}
|
||||
state["requests"].append((url.path, query))
|
||||
if url.path == "/v1/namespace/$/secret/list":
|
||||
if "page_token" not in query:
|
||||
self._write_response(
|
||||
200,
|
||||
{"secrets": [{"name": "openai-prod"}], "page_token": "next"},
|
||||
)
|
||||
else:
|
||||
assert query["page_token"] == "next"
|
||||
self._write_response(200, {"secrets": [{"name": "hf-prod"}]})
|
||||
return
|
||||
if url.path != "/v1/namespace/$/function/list":
|
||||
self._write_response(404, {"error": "not found"})
|
||||
return
|
||||
@@ -1278,7 +1611,7 @@ def _mock_remote_function_catalog():
|
||||
"functions": [
|
||||
{
|
||||
"name": "normalize_score",
|
||||
"version": "fv_exact",
|
||||
"version": FUNCTION_VERSION,
|
||||
"definition": state["version"],
|
||||
}
|
||||
],
|
||||
@@ -1314,7 +1647,7 @@ def test_remote_registration_job_and_exact_version_reopen_round_trip():
|
||||
|
||||
assert created == reopened
|
||||
assert reopened.name == "normalize_score"
|
||||
assert reopened.version == "fv_exact"
|
||||
assert reopened.version == FUNCTION_VERSION
|
||||
create_request = state["requests"][0][1]
|
||||
expected_request = json.loads(
|
||||
normalize_score.registration_request.to_canonical_json()
|
||||
@@ -1323,6 +1656,92 @@ def test_remote_registration_job_and_exact_version_reopen_round_trip():
|
||||
assert create_request == expected_request
|
||||
|
||||
|
||||
def test_remote_registration_sends_bindings_and_never_a_credential():
|
||||
with _mock_remote_function_catalog() as (host, state):
|
||||
db = lancedb.connect(
|
||||
"db://dev",
|
||||
api_key="fake",
|
||||
host_override=host,
|
||||
client_config={"retry_config": {"retries": 0}},
|
||||
)
|
||||
created = db.create_function(
|
||||
analyze_caption,
|
||||
secrets=[
|
||||
EnvVarSecret(secret_name="openai-prod", env_variable="OPENAI_API_KEY")
|
||||
],
|
||||
)
|
||||
|
||||
assert list(created.secret_bindings) == [
|
||||
SecretBinding(
|
||||
kind="env",
|
||||
variable="OPENAI_API_KEY",
|
||||
secret_ref=SecretReference(name="openai-prod"),
|
||||
)
|
||||
]
|
||||
path, create_request = state["requests"][0]
|
||||
assert path == "/v1/function/analyze_caption/create"
|
||||
assert create_request["secret_bindings"] == [
|
||||
{
|
||||
"kind": "env",
|
||||
"variable": "OPENAI_API_KEY",
|
||||
"secret_ref": {"name": "openai-prod"},
|
||||
}
|
||||
]
|
||||
# The request names a Secret and carries nothing that could be one. The
|
||||
# Function's own name is the path identifier rather than a body field, so
|
||||
# it is the one key the body does not repeat.
|
||||
expected = json.loads(
|
||||
analyze_caption.bind_secrets(
|
||||
[EnvVarSecret(secret_name="openai-prod", env_variable="OPENAI_API_KEY")]
|
||||
).to_canonical_json()
|
||||
)
|
||||
assert expected.pop("name") == "analyze_caption"
|
||||
assert create_request == expected
|
||||
|
||||
|
||||
def test_remote_secret_verbs_round_trip():
|
||||
with _mock_remote_function_catalog() as (host, state):
|
||||
db = lancedb.connect(
|
||||
"db://dev",
|
||||
api_key="fake",
|
||||
host_override=host,
|
||||
client_config={"retry_config": {"retries": 0}},
|
||||
)
|
||||
assert db.create_secret("openai-prod", "sk-live-0001") is None
|
||||
assert db.alter_secret("openai-prod", "sk-live-0002") is None
|
||||
assert db.list_secrets() == ["openai-prod", "hf-prod"]
|
||||
assert db.drop_secret("openai-prod") is None
|
||||
|
||||
routes = [path for path, _ in state["requests"]]
|
||||
assert routes == [
|
||||
"/v1/secret/openai-prod/create",
|
||||
"/v1/secret/openai-prod/alter",
|
||||
"/v1/namespace/$/secret/list",
|
||||
"/v1/namespace/$/secret/list",
|
||||
"/v1/secret/openai-prod/drop",
|
||||
]
|
||||
# The Secret is the path identifier, so the body is the value alone.
|
||||
assert state["requests"][0][1] == {"value": "sk-live-0001"}
|
||||
# Listing is a GET: the first page asks for nothing, the second resumes on
|
||||
# the token the server handed back, and neither carries a body.
|
||||
assert state["requests"][2][1] == {}
|
||||
assert state["requests"][3][1] == {"page_token": "next"}
|
||||
|
||||
|
||||
def test_building_a_binding_contacts_no_server():
|
||||
"""A binding is a local value: it says nothing about whether the Secret exists.
|
||||
|
||||
Existence is the server's answer at registration, where a mistyped name is a
|
||||
clear error rather than a client-side check that was already stale.
|
||||
"""
|
||||
with _mock_remote_function_catalog() as (_host, state):
|
||||
binding = EnvVarSecret(secret_name="openai-prod", env_variable="OPENAI_API_KEY")
|
||||
assert binding.secret_name == "openai-prod"
|
||||
assert binding.env_variable == "OPENAI_API_KEY"
|
||||
|
||||
assert state["requests"] == []
|
||||
|
||||
|
||||
def test_blocking_remote_registration_returns_function_version():
|
||||
with _mock_remote_function_catalog() as (host, state):
|
||||
db = lancedb.connect(
|
||||
@@ -1334,7 +1753,7 @@ def test_blocking_remote_registration_returns_function_version():
|
||||
created = db.create_function(normalize_score)
|
||||
|
||||
assert created.name == "normalize_score"
|
||||
assert created.version == "fv_exact"
|
||||
assert created.version == FUNCTION_VERSION
|
||||
assert [path for path, _ in state["requests"]] == [
|
||||
"/v1/function/normalize_score/create",
|
||||
"/v1/jobs/describe",
|
||||
@@ -1392,12 +1811,12 @@ def test_remote_drop_function_sends_exact_version():
|
||||
host_override=host,
|
||||
client_config={"retry_config": {"retries": 0}},
|
||||
)
|
||||
assert db.drop_function("normalize_score", version="fv_exact") is True
|
||||
assert db.drop_function("normalize_score", version=FUNCTION_VERSION) is True
|
||||
|
||||
assert state["requests"] == [
|
||||
(
|
||||
"/v1/function/normalize_score/drop",
|
||||
{"version": "fv_exact"},
|
||||
{"version": FUNCTION_VERSION},
|
||||
)
|
||||
]
|
||||
|
||||
@@ -1411,11 +1830,13 @@ async def test_async_remote_drop_function_sends_exact_version():
|
||||
host_override=host,
|
||||
client_config={"retry_config": {"retries": 0}},
|
||||
)
|
||||
assert await db.drop_function("normalize_score", version="fv_exact") is True
|
||||
assert (
|
||||
await db.drop_function("normalize_score", version=FUNCTION_VERSION) is True
|
||||
)
|
||||
|
||||
assert state["requests"] == [
|
||||
(
|
||||
"/v1/function/normalize_score/drop",
|
||||
{"version": "fv_exact"},
|
||||
{"version": FUNCTION_VERSION},
|
||||
)
|
||||
]
|
||||
|
||||
@@ -1,9 +1,15 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright The LanceDB Authors
|
||||
|
||||
import contextlib
|
||||
import http.server
|
||||
import json
|
||||
import threading
|
||||
|
||||
import lancedb
|
||||
import pytest
|
||||
from lancedb.materialized_view import MaterializedViewDefinition
|
||||
from lancedb.remote.db import RemoteDBConnection
|
||||
|
||||
|
||||
STABLE_ROW_IDS = {"new_table_enable_stable_row_ids": "true"}
|
||||
@@ -22,6 +28,152 @@ def make_db(tmp_path):
|
||||
return db
|
||||
|
||||
|
||||
@contextlib.contextmanager
|
||||
def mock_remote_materialized_views():
|
||||
requests = []
|
||||
|
||||
class Handler(http.server.BaseHTTPRequestHandler):
|
||||
def log_message(self, *args):
|
||||
pass
|
||||
|
||||
def do_GET(self):
|
||||
requests.append(self.path)
|
||||
encoded = json.dumps({"views": ["daily_sales"]}).encode()
|
||||
self.send_response(200)
|
||||
self.send_header("Content-Type", "application/json")
|
||||
self.send_header("Content-Length", str(len(encoded)))
|
||||
self.end_headers()
|
||||
self.wfile.write(encoded)
|
||||
|
||||
with http.server.HTTPServer(("localhost", 0), Handler) as server:
|
||||
thread = threading.Thread(target=server.serve_forever)
|
||||
thread.start()
|
||||
try:
|
||||
yield f"http://localhost:{server.server_address[1]}", requests
|
||||
finally:
|
||||
server.shutdown()
|
||||
thread.join()
|
||||
|
||||
|
||||
@contextlib.contextmanager
|
||||
def mock_remote_materialized_view_create():
|
||||
requests = []
|
||||
|
||||
class Handler(http.server.BaseHTTPRequestHandler):
|
||||
def log_message(self, *args):
|
||||
pass
|
||||
|
||||
def do_POST(self):
|
||||
length = int(self.headers.get("Content-Length", 0))
|
||||
body = json.loads(self.rfile.read(length) or b"{}")
|
||||
requests.append((self.path, body))
|
||||
job_id = "mv-drop-123" if self.path.endswith("/drop") else "mv-create-123"
|
||||
encoded = json.dumps({"job_id": job_id}).encode()
|
||||
self.send_response(202)
|
||||
self.send_header("Content-Type", "application/json")
|
||||
self.send_header("Content-Length", str(len(encoded)))
|
||||
self.end_headers()
|
||||
self.wfile.write(encoded)
|
||||
|
||||
with http.server.HTTPServer(("localhost", 0), Handler) as server:
|
||||
thread = threading.Thread(target=server.serve_forever)
|
||||
thread.start()
|
||||
try:
|
||||
yield f"http://localhost:{server.server_address[1]}", requests
|
||||
finally:
|
||||
server.shutdown()
|
||||
thread.join()
|
||||
|
||||
|
||||
def test_remote_list_uses_namespace_route():
|
||||
with mock_remote_materialized_views() as (host, requests):
|
||||
db = lancedb.connect(
|
||||
"db://dev",
|
||||
api_key="fake",
|
||||
host_override=host,
|
||||
client_config={"retry_config": {"retries": 0}},
|
||||
)
|
||||
assert db.list_materialized_views() == ["daily_sales"]
|
||||
assert requests == ["/v1/namespace/$/materialized_view/list"]
|
||||
|
||||
|
||||
def test_remote_create_async_returns_server_job():
|
||||
with mock_remote_materialized_view_create() as (host, requests):
|
||||
db = lancedb.connect(
|
||||
"db://dev",
|
||||
api_key="fake",
|
||||
host_override=host,
|
||||
client_config={"retry_config": {"retries": 0}},
|
||||
)
|
||||
job = db.create_materialized_view_async("adults", "people", where="age >= 18")
|
||||
assert job.id == "mv-create-123"
|
||||
assert requests == [
|
||||
(
|
||||
"/v1/materialized_view/adults/create",
|
||||
{"query": 'SELECT * FROM "people" WHERE age >= 18', "with_no_data": False},
|
||||
)
|
||||
]
|
||||
|
||||
|
||||
def test_remote_drop_async_returns_server_job():
|
||||
with mock_remote_materialized_view_create() as (host, requests):
|
||||
db = lancedb.connect(
|
||||
"db://dev",
|
||||
api_key="fake",
|
||||
host_override=host,
|
||||
client_config={"retry_config": {"retries": 0}},
|
||||
)
|
||||
job = db.drop_materialized_view_async("adults")
|
||||
assert job.id == "mv-drop-123"
|
||||
assert requests == [("/v1/materialized_view/adults/drop", {})]
|
||||
|
||||
|
||||
def test_sync_remote_create_uses_public_async_connection():
|
||||
calls = []
|
||||
|
||||
class StubAsyncTable:
|
||||
name = "adults"
|
||||
|
||||
class StubAsyncMaterializedView:
|
||||
table = StubAsyncTable()
|
||||
|
||||
class StrictAsyncConnection:
|
||||
async def create_materialized_view(
|
||||
self,
|
||||
name,
|
||||
source,
|
||||
*,
|
||||
select=None,
|
||||
where=None,
|
||||
limit=None,
|
||||
with_no_data=False,
|
||||
):
|
||||
calls.append((name, source, select, where, limit, with_no_data))
|
||||
return StubAsyncMaterializedView()
|
||||
|
||||
async def drop_materialized_view(self, name, *, namespace_path=None):
|
||||
calls.append(("drop", name, namespace_path))
|
||||
|
||||
db = RemoteDBConnection.__new__(RemoteDBConnection)
|
||||
db._conn = StrictAsyncConnection()
|
||||
db.db_name = "example"
|
||||
db.serialize = lambda: "{}"
|
||||
|
||||
view = db.create_materialized_view(
|
||||
"adults",
|
||||
"people",
|
||||
select=["name"],
|
||||
where="age >= 18",
|
||||
limit=10,
|
||||
with_no_data=True,
|
||||
)
|
||||
assert view.name == "adults"
|
||||
assert calls == [("adults", "people", ["name"], "age >= 18", 10, True)]
|
||||
|
||||
db.drop_materialized_view("adults", namespace_path=["analytics"])
|
||||
assert calls[-1] == ("drop", "adults", ["analytics"])
|
||||
|
||||
|
||||
def test_create_refresh_and_query(tmp_path):
|
||||
db = make_db(tmp_path)
|
||||
view = db.create_materialized_view(
|
||||
@@ -31,16 +183,34 @@ def test_create_refresh_and_query(tmp_path):
|
||||
where="age >= 18",
|
||||
)
|
||||
assert view.name == "adults"
|
||||
assert view.table.count_rows() == 0
|
||||
|
||||
result = view.refresh()
|
||||
assert result.mode == "rebuild"
|
||||
assert result.rows_written == 2
|
||||
assert view.table.count_rows() == 2
|
||||
|
||||
rows = view.table.search().to_list()
|
||||
assert sorted(row["shout"] for row in rows) == ["ADA", "GRACE"]
|
||||
|
||||
|
||||
def test_create_and_refresh_jobs(tmp_path):
|
||||
db = make_db(tmp_path)
|
||||
create_job = db.create_materialized_view_async(
|
||||
"adults", "people", where="age >= 18", with_no_data=True
|
||||
)
|
||||
assert create_job.id is None
|
||||
assert create_job.wait() is None
|
||||
|
||||
view = db.open_materialized_view("adults")
|
||||
refresh_job = view.refresh_async()
|
||||
assert refresh_job.id is None
|
||||
result = refresh_job.wait()
|
||||
assert result.mode == "rebuild"
|
||||
assert result.rows_written == 2
|
||||
assert view.table.count_rows() == 2
|
||||
|
||||
drop_job = db.drop_materialized_view_async("adults")
|
||||
assert drop_job.id is None
|
||||
assert drop_job.wait() is None
|
||||
assert "adults" not in db.list_materialized_views()
|
||||
|
||||
|
||||
def test_definition_round_trips(tmp_path):
|
||||
db = make_db(tmp_path)
|
||||
db.create_materialized_view("adults", "people", where="age >= 18")
|
||||
@@ -56,7 +226,7 @@ def test_definition_round_trips(tmp_path):
|
||||
|
||||
def test_incremental_refresh_after_append(tmp_path):
|
||||
db = make_db(tmp_path)
|
||||
view = db.create_materialized_view("copy", "people")
|
||||
view = db.create_materialized_view("copy", "people", with_no_data=True)
|
||||
view.refresh()
|
||||
|
||||
db.open_table("people").add([{"name": "alan", "age": 41}])
|
||||
@@ -70,7 +240,7 @@ def test_incremental_refresh_after_append(tmp_path):
|
||||
|
||||
def test_incremental_refresh_after_update(tmp_path):
|
||||
db = make_db(tmp_path)
|
||||
view = db.create_materialized_view("copy", "people")
|
||||
view = db.create_materialized_view("copy", "people", with_no_data=True)
|
||||
view.refresh()
|
||||
|
||||
db.open_table("people").update(where="name = 'kid'", values={"age": 8})
|
||||
@@ -87,7 +257,7 @@ def test_legacy_storage_source_update_rebuilds(tmp_path):
|
||||
storage_options={**STABLE_ROW_IDS, "new_table_data_storage_version": "legacy"},
|
||||
)
|
||||
db.create_table("people", [{"name": "ada", "age": 36}, {"name": "kid", "age": 7}])
|
||||
view = db.create_materialized_view("copy", "people")
|
||||
view = db.create_materialized_view("copy", "people", with_no_data=True)
|
||||
view.refresh()
|
||||
|
||||
db.open_table("people").update(where="name = 'kid'", values={"age": 8})
|
||||
@@ -104,6 +274,11 @@ def test_list_and_not_a_view(tmp_path):
|
||||
assert db.list_materialized_views() == ["adults"]
|
||||
with pytest.raises(ValueError, match="not a materialized view"):
|
||||
db.open_materialized_view("people")
|
||||
with pytest.raises(ValueError, match="not a materialized view"):
|
||||
db.drop_materialized_view("people")
|
||||
|
||||
db.drop_materialized_view("adults")
|
||||
assert db.list_materialized_views() == []
|
||||
|
||||
|
||||
def test_invalid_expression_fails_at_create(tmp_path):
|
||||
@@ -119,7 +294,10 @@ async def test_async_create_refresh_and_open(tmp_path):
|
||||
await db.create_table("people", [{"name": "ada", "age": 36}])
|
||||
|
||||
view = await db.create_materialized_view(
|
||||
"shouts", "people", select=[("shout", "upper(name)")]
|
||||
"shouts",
|
||||
"people",
|
||||
select=[("shout", "upper(name)")],
|
||||
with_no_data=True,
|
||||
)
|
||||
result = await view.refresh()
|
||||
assert result.mode == "rebuild"
|
||||
@@ -131,11 +309,35 @@ async def test_async_create_refresh_and_open(tmp_path):
|
||||
assert await db.list_materialized_views() == ["shouts"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_create_and_refresh_jobs(tmp_path):
|
||||
db = await lancedb.connect_async(tmp_path, storage_options=STABLE_ROW_IDS)
|
||||
await db.create_table("people", [{"name": "ada", "age": 36}])
|
||||
|
||||
create_job = await db.create_materialized_view_async(
|
||||
"adults", "people", with_no_data=True
|
||||
)
|
||||
assert create_job.id is None
|
||||
assert await create_job.wait() is None
|
||||
|
||||
view = await db.open_materialized_view("adults")
|
||||
refresh_job = await view.refresh_async()
|
||||
assert refresh_job.id is None
|
||||
result = await refresh_job.wait()
|
||||
assert result.mode == "rebuild"
|
||||
assert result.rows_written == 1
|
||||
|
||||
drop_job = await db.drop_materialized_view_async("adults")
|
||||
assert drop_job.id is None
|
||||
assert await drop_job.wait() is None
|
||||
assert "adults" not in await db.list_materialized_views()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_incremental(tmp_path):
|
||||
db = await lancedb.connect_async(tmp_path, storage_options=STABLE_ROW_IDS)
|
||||
await db.create_table("people", [{"name": "ada", "age": 36}])
|
||||
view = await db.create_materialized_view("copy", "people")
|
||||
view = await db.create_materialized_view("copy", "people", with_no_data=True)
|
||||
await view.refresh()
|
||||
|
||||
table = await db.open_table("people")
|
||||
@@ -157,7 +359,10 @@ def test_bare_select_names_are_quoted(tmp_path):
|
||||
db.create_table("odd_names", [{"order item": "widget", "select": 2}])
|
||||
|
||||
view = db.create_materialized_view(
|
||||
"quoted", "odd_names", select=["order item", "select"]
|
||||
"quoted",
|
||||
"odd_names",
|
||||
select=["order item", "select"],
|
||||
with_no_data=True,
|
||||
)
|
||||
result = view.refresh()
|
||||
assert result.rows_written == 1
|
||||
@@ -166,19 +371,6 @@ def test_bare_select_names_are_quoted(tmp_path):
|
||||
assert rows[0]["select"] == 2
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_remote_is_refused_without_network():
|
||||
db = await lancedb.connect_async(
|
||||
"db://nowhere", api_key="sk_test", region="us-east-1"
|
||||
)
|
||||
with pytest.raises(NotImplementedError, match="local"):
|
||||
await db.create_materialized_view("v", "src")
|
||||
with pytest.raises(NotImplementedError, match="local"):
|
||||
await db.open_materialized_view("v")
|
||||
with pytest.raises(NotImplementedError, match="local"):
|
||||
await db.list_materialized_views()
|
||||
|
||||
|
||||
def test_scalar_select_is_one_column(tmp_path):
|
||||
db = make_db(tmp_path)
|
||||
view = db.create_materialized_view("just_name", "people", select="name")
|
||||
@@ -235,6 +427,17 @@ def test_namespace_connection_materialized_views(tmp_path):
|
||||
with pytest.raises(ValueError, match="not a materialized view"):
|
||||
db.open_materialized_view("people")
|
||||
|
||||
create_job = db.create_materialized_view_async(
|
||||
"job_view", "people", with_no_data=True
|
||||
)
|
||||
assert create_job.wait() is None
|
||||
refresh_job = db.open_materialized_view("job_view").refresh_async()
|
||||
assert refresh_job.wait().rows_written == 2
|
||||
|
||||
assert db.drop_materialized_view_async("job_view").wait() is None
|
||||
db.drop_materialized_view("adults")
|
||||
assert db.list_materialized_views() == []
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_namespace_connection_materialized_views(tmp_path):
|
||||
@@ -267,6 +470,19 @@ async def test_async_namespace_connection_materialized_views(tmp_path):
|
||||
)
|
||||
assert handle._namespace_path == through_namespace._namespace_path
|
||||
|
||||
create_job = await db.create_materialized_view_async(
|
||||
"job_view", "people", with_no_data=True
|
||||
)
|
||||
assert await create_job.wait() is None
|
||||
job_view = await db.open_materialized_view("job_view")
|
||||
refresh_job = await job_view.refresh_async()
|
||||
assert (await refresh_job.wait()).rows_written == 2
|
||||
|
||||
drop_job = await db.drop_materialized_view_async("job_view")
|
||||
assert await drop_job.wait() is None
|
||||
await db.drop_materialized_view("adults")
|
||||
assert await db.list_materialized_views() == []
|
||||
|
||||
|
||||
def test_namespaced_select_kind_is_read_and_unknown_kinds_are_refused():
|
||||
import json
|
||||
|
||||
@@ -0,0 +1,149 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright The LanceDB Authors
|
||||
|
||||
"""Namespace operations must not require the optional ``pylance`` dependency.
|
||||
|
||||
The sync ``lancedb.connect()`` connection used to route namespace operations
|
||||
through the Python ``lance_namespace`` client, whose ``dir`` implementation
|
||||
lives in ``lance.namespace`` (shipped by the optional ``pylance`` extra). On an
|
||||
install without that extra, even ``db.list_namespaces()`` failed, while the
|
||||
async API worked because it goes straight to the native Rust connection.
|
||||
|
||||
The ``without_pylance`` fixture below simulates a missing ``pylance`` so these
|
||||
tests fail on any environment when the native routing regresses. The ground
|
||||
truth remains the "Test without pylance or pandas" CI job, which runs this file
|
||||
with ``pylance`` and ``pandas`` actually uninstalled -- so nothing here may
|
||||
import either at module scope.
|
||||
"""
|
||||
|
||||
import sys
|
||||
from importlib import import_module
|
||||
from importlib.abc import MetaPathFinder
|
||||
|
||||
import lancedb
|
||||
import pyarrow as pa
|
||||
import pytest
|
||||
|
||||
|
||||
def _is_lance(module_name: str) -> bool:
|
||||
# "lance_namespace" is a separate, non-optional package -- leave it alone.
|
||||
return module_name == "lance" or module_name.startswith("lance.")
|
||||
|
||||
|
||||
class _BlockLanceImports(MetaPathFinder):
|
||||
def find_spec(self, fullname, path=None, target=None):
|
||||
if _is_lance(fullname):
|
||||
raise ModuleNotFoundError(f"No module named {fullname!r}", name=fullname)
|
||||
return None
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def without_pylance(monkeypatch):
|
||||
"""Make ``lance`` unimportable, as on an install without the pylance extra.
|
||||
|
||||
Uninstalled means absent from ``sys.modules`` too, not merely unimportable:
|
||||
LanceDB has code that branches on ``"lance" in sys.modules``, so a fixture
|
||||
that poisons the entry instead of removing it would fail tests that a real
|
||||
install passes.
|
||||
"""
|
||||
for name in list(sys.modules):
|
||||
if _is_lance(name):
|
||||
monkeypatch.delitem(sys.modules, name)
|
||||
monkeypatch.setattr(sys, "meta_path", [_BlockLanceImports(), *sys.meta_path])
|
||||
|
||||
|
||||
def _schema() -> pa.Schema:
|
||||
return pa.schema([pa.field("id", pa.int64())])
|
||||
|
||||
|
||||
def test_fixture_matches_an_uninstalled_pylance(without_pylance):
|
||||
"""Guard the guard: the other tests are meaningless if lance stays importable."""
|
||||
assert "lance" not in sys.modules
|
||||
with pytest.raises(ModuleNotFoundError):
|
||||
import_module("lance.namespace")
|
||||
|
||||
|
||||
def test_list_namespaces_on_sync_connection(tmp_path, without_pylance):
|
||||
"""The original reproducer: this alone used to raise on a sync connection."""
|
||||
db = lancedb.connect(tmp_path)
|
||||
assert db.list_namespaces().namespaces == []
|
||||
|
||||
|
||||
def test_sync_namespace_lifecycle(tmp_path, without_pylance):
|
||||
db = lancedb.connect(tmp_path)
|
||||
|
||||
db.create_namespace(["child"])
|
||||
assert db.list_namespaces().namespaces == ["child"]
|
||||
assert db.list_namespaces(namespace_path=["child"]).namespaces == []
|
||||
db.describe_namespace(["child"])
|
||||
|
||||
db.create_namespace(["child", "grandchild"])
|
||||
assert db.list_namespaces(namespace_path=["child"]).namespaces == ["grandchild"]
|
||||
|
||||
db.drop_namespace(["child", "grandchild"])
|
||||
db.drop_namespace(["child"])
|
||||
assert db.list_namespaces().namespaces == []
|
||||
|
||||
|
||||
def test_sync_namespaced_table_lifecycle(tmp_path, without_pylance):
|
||||
db = lancedb.connect(tmp_path)
|
||||
db.create_namespace(["child"])
|
||||
|
||||
table = db.create_table("tbl", schema=_schema(), namespace_path=["child"])
|
||||
assert table.namespace == ["child"]
|
||||
table.add([{"id": 1}])
|
||||
|
||||
assert db.list_tables(namespace_path=["child"]).tables == ["tbl"]
|
||||
assert db.list_tables().tables == []
|
||||
|
||||
opened = db.open_table("tbl", namespace_path=["child"])
|
||||
assert opened.namespace == ["child"]
|
||||
assert opened.count_rows() == 1
|
||||
assert opened.search().limit(5).to_arrow().num_rows == 1
|
||||
|
||||
db.drop_table("tbl", namespace_path=["child"])
|
||||
assert db.list_tables(namespace_path=["child"]).tables == []
|
||||
db.drop_namespace(["child"])
|
||||
|
||||
|
||||
def test_sync_root_table_lifecycle(tmp_path, without_pylance):
|
||||
"""Root-namespace tables share the namespace plumbing, so cover them too."""
|
||||
db = lancedb.connect(tmp_path)
|
||||
|
||||
table = db.create_table("tbl", schema=_schema())
|
||||
table.add([{"id": 1}])
|
||||
|
||||
assert db.table_names() == ["tbl"]
|
||||
assert "tbl" in db
|
||||
assert db["tbl"].count_rows() == 1
|
||||
|
||||
db.drop_table("tbl")
|
||||
assert db.table_names() == []
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_namespace_lifecycle(tmp_path, without_pylance):
|
||||
db = await lancedb.connect_async(tmp_path)
|
||||
|
||||
await db.create_namespace(["child"])
|
||||
assert (await db.list_namespaces()).namespaces == ["child"]
|
||||
|
||||
table = await db.create_table("tbl", schema=_schema(), namespace_path=["child"])
|
||||
await table.add([{"id": 1}])
|
||||
assert (await db.list_tables(namespace_path=["child"])).tables == ["tbl"]
|
||||
assert await table.count_rows() == 1
|
||||
|
||||
await db.drop_table("tbl", namespace_path=["child"])
|
||||
await db.drop_namespace(["child"])
|
||||
assert (await db.list_namespaces()).namespaces == []
|
||||
|
||||
|
||||
def test_namespace_client_still_requires_pylance(tmp_path, without_pylance):
|
||||
"""``namespace_client()`` is the one namespace API that opts into pylance.
|
||||
|
||||
It hands out a Python ``LanceNamespace``, so it cannot be served natively.
|
||||
Pinned here so the boundary stays explicit and the error stays actionable.
|
||||
"""
|
||||
db = lancedb.connect(tmp_path)
|
||||
with pytest.raises(ValueError, match="lance.namespace.DirectoryNamespace"):
|
||||
db.namespace_client()
|
||||
@@ -982,6 +982,71 @@ def test_remote_refresh_async_returns_typed_terminal_result():
|
||||
assert result.version == 8
|
||||
|
||||
|
||||
def test_remote_function_errors_lists_the_rows_a_refresh_skipped():
|
||||
listing = {
|
||||
"records": [
|
||||
{
|
||||
"job_id": "j-7",
|
||||
"fragment_id": 3,
|
||||
"row_offset": 9,
|
||||
"column": "embedding",
|
||||
"function": "embed",
|
||||
"function_version": "2",
|
||||
"table_version": 11,
|
||||
"error_type": "ValueError",
|
||||
"error_message": "bad input 'x'",
|
||||
"created_at_millis": 1700000000000,
|
||||
}
|
||||
],
|
||||
"fragments": [
|
||||
{
|
||||
"job_id": "j-7",
|
||||
"fragment_id": 4,
|
||||
"rows_skipped": 500,
|
||||
"rows_recorded": 100,
|
||||
}
|
||||
],
|
||||
"truncated": True,
|
||||
}
|
||||
bodies = []
|
||||
|
||||
def handler(request):
|
||||
content_len = int(request.headers.get("Content-Length", 0))
|
||||
body = request.rfile.read(content_len) if content_len > 0 else b""
|
||||
if request.path == "/v1/table/test/errors":
|
||||
bodies.append(json.loads(body))
|
||||
request.send_response(200)
|
||||
request.send_header("Content-Type", "application/json")
|
||||
request.end_headers()
|
||||
request.wfile.write(json.dumps(listing).encode())
|
||||
elif request.path == "/v1/table/test/describe/":
|
||||
request.send_response(200)
|
||||
request.send_header("Content-Type", "application/json")
|
||||
request.end_headers()
|
||||
request.wfile.write(
|
||||
json.dumps({"version": 1, "schema": {"fields": []}}).encode()
|
||||
)
|
||||
else:
|
||||
request.send_response(404)
|
||||
request.end_headers()
|
||||
|
||||
with mock_lancedb_connection(handler) as db:
|
||||
table = db.open_table("test")
|
||||
errors = table.function_errors(job_id="j-7", column="embedding", limit=2)
|
||||
everything = table.function_errors()
|
||||
|
||||
assert bodies == [{"job_id": "j-7", "column": "embedding", "limit": 2}, {}]
|
||||
assert errors.truncated is True
|
||||
assert [r.error_message for r in errors.records] == ["bad input 'x'"]
|
||||
assert errors.records[0].row_offset == 9
|
||||
assert errors.records[0].function_version == "2"
|
||||
assert (errors.fragments[0].rows_skipped, errors.fragments[0].rows_recorded) == (
|
||||
500,
|
||||
100,
|
||||
)
|
||||
assert everything.truncated is True
|
||||
|
||||
|
||||
def test_remote_job_wait_raises_on_failure():
|
||||
from lancedb.exceptions import JobFailedError
|
||||
from lancedb.index import BTree
|
||||
|
||||
@@ -21,6 +21,7 @@ from lancedb.rerankers import (
|
||||
OpenaiReranker,
|
||||
JinaReranker,
|
||||
AnswerdotaiRerankers,
|
||||
TypeSafeReranker,
|
||||
VoyageAIReranker,
|
||||
MRRReranker,
|
||||
WatsonxReranker,
|
||||
@@ -101,7 +102,15 @@ def get_test_table(tmp_path):
|
||||
return table, MyTable
|
||||
|
||||
|
||||
def _run_test_reranker(reranker, table, query, query_vector, schema):
|
||||
def _run_test_reranker(
|
||||
reranker, table, query, query_vector, schema, deterministic=True
|
||||
):
|
||||
"""Exercise a reranker across search types.
|
||||
|
||||
Set ``deterministic=False`` for rerankers whose scores can vary between
|
||||
identical calls, such as remote model APIs; repeated searches are then not
|
||||
expected to return the same order.
|
||||
"""
|
||||
# Hybrid search setting
|
||||
result1 = (
|
||||
table.search(query, query_type="hybrid", vector_column_name="vector")
|
||||
@@ -113,7 +122,8 @@ def _run_test_reranker(reranker, table, query, query_vector, schema):
|
||||
.rerank(reranker=reranker)
|
||||
.to_pydantic(schema)
|
||||
)
|
||||
assert result1 == result2
|
||||
if deterministic:
|
||||
assert result1 == result2
|
||||
|
||||
query_vector = table.to_pandas()["vector"][0]
|
||||
result = (
|
||||
@@ -203,7 +213,9 @@ def _run_test_reranker(reranker, table, query, query_vector, schema):
|
||||
)
|
||||
assert len(result_deduped) <= 20
|
||||
result_arrow = reranker.rerank_multivector([rs1.to_arrow(), rs2.to_arrow()], query)
|
||||
assert len(result) == 20 and result == result_arrow
|
||||
assert len(result) == 20 and len(result_arrow) == 20
|
||||
if deterministic:
|
||||
assert result == result_arrow
|
||||
|
||||
|
||||
def _run_test_hybrid_reranker(reranker, tmp_path):
|
||||
@@ -515,6 +527,95 @@ def test_voyageai_reranker(tmp_path):
|
||||
_run_test_reranker(reranker, table, "single player experience", None, schema)
|
||||
|
||||
|
||||
class _FakeTypeSafeClient:
|
||||
"""Stands in for ``typesafe_sdk.TypeSafeClient``; scores by word overlap."""
|
||||
|
||||
def __init__(self):
|
||||
self.requests = []
|
||||
|
||||
def system_one(self, state, questions, model):
|
||||
self.requests.append((state, questions, model))
|
||||
query_words = set(state["query"].lower().split())
|
||||
doc_words = set(state["document"].lower().split())
|
||||
noul = len(query_words & doc_words) / len(query_words)
|
||||
answers = {key: type("NoulAnswer", (), {"noul": noul})() for key in questions}
|
||||
return type("SystemOneResponse", (), {"answers": answers})()
|
||||
|
||||
|
||||
def test_typesafe_reranker_with_fake_client(tmp_path):
|
||||
reranker = TypeSafeReranker(max_concurrency=4)
|
||||
reranker._client = _FakeTypeSafeClient()
|
||||
table, schema = get_test_table(tmp_path)
|
||||
_run_test_reranker(reranker, table, "single player experience", None, schema)
|
||||
|
||||
state, questions, model = reranker._client.requests[0]
|
||||
assert model == "jev-latest"
|
||||
assert set(state) == {"query", "document"}
|
||||
assert questions == {
|
||||
"relevance": {
|
||||
"type": "noul",
|
||||
"instructions": reranker.instructions,
|
||||
"criteria": reranker.criteria,
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
def test_typesafe_reranker_scores_each_row():
|
||||
reranker = TypeSafeReranker(
|
||||
column="body",
|
||||
instructions="Is this about cats?",
|
||||
criteria={},
|
||||
return_score="all",
|
||||
)
|
||||
reranker._client = _FakeTypeSafeClient()
|
||||
results = pa.table(
|
||||
{
|
||||
"body": ["dogs bark", "cats purr and cats nap", None, "cats"],
|
||||
"_distance": [0.1, 0.2, 0.3, 0.4],
|
||||
}
|
||||
)
|
||||
|
||||
reranked = reranker.rerank_vector("cats nap", results)
|
||||
|
||||
assert reranked["body"].to_pylist() == [
|
||||
"cats purr and cats nap",
|
||||
"cats",
|
||||
"dogs bark",
|
||||
None,
|
||||
]
|
||||
assert reranked["_relevance_score"].to_pylist() == [1.0, 0.5, 0.0, 0.0]
|
||||
assert reranked["_distance"].to_pylist() == [0.2, 0.4, 0.1, 0.3]
|
||||
# Null documents are scored 0 without a request.
|
||||
assert len(reranker._client.requests) == 3
|
||||
assert reranker._client.requests[0][1] == {
|
||||
"relevance": {"type": "noul", "instructions": "Is this about cats?"}
|
||||
}
|
||||
|
||||
|
||||
def test_typesafe_reranker_rejects_invalid_arguments():
|
||||
with pytest.raises(ValueError, match="criteria keys"):
|
||||
TypeSafeReranker(criteria={"yes": "relevant"})
|
||||
with pytest.raises(ValueError, match="max_concurrency"):
|
||||
TypeSafeReranker(max_concurrency=0)
|
||||
|
||||
|
||||
@pytest.mark.skipif(
|
||||
os.environ.get("TYPESAFE_API_KEY") is None, reason="TYPESAFE_API_KEY not set"
|
||||
)
|
||||
def test_typesafe_reranker(tmp_path):
|
||||
pytest.importorskip("typesafe_sdk")
|
||||
reranker = TypeSafeReranker()
|
||||
table, schema = get_test_table(tmp_path)
|
||||
_run_test_reranker(
|
||||
reranker,
|
||||
table,
|
||||
"single player experience",
|
||||
None,
|
||||
schema,
|
||||
deterministic=False,
|
||||
)
|
||||
|
||||
|
||||
def test_empty_result_reranker():
|
||||
pytest.importorskip("sentence_transformers")
|
||||
db = lancedb.connect("memory://")
|
||||
|
||||
@@ -2979,28 +2979,29 @@ async def test_merge_insert_async(mem_db_async: AsyncConnection):
|
||||
assert (await table.to_arrow()).sort_by("a") == expected
|
||||
|
||||
|
||||
def _json_arrow_table(schema, rows):
|
||||
json_type = schema.field("j").type
|
||||
json_values = pa.ExtensionArray.from_storage(
|
||||
json_type,
|
||||
pa.array([value for _, value in rows], type=json_type.storage_type),
|
||||
)
|
||||
return pa.Table.from_arrays(
|
||||
[pa.array([row_id for row_id, _ in rows]), json_values], schema=schema
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.skipif(not hasattr(pa, "json_"), reason="requires PyArrow JSON type")
|
||||
@pytest.mark.asyncio
|
||||
async def test_merge_insert_encodes_json(mem_db_async: AsyncConnection):
|
||||
json_type = pa.json_()
|
||||
schema = pa.schema([pa.field("id", pa.string()), pa.field("j", json_type)])
|
||||
|
||||
def json_table(rows):
|
||||
json_values = pa.ExtensionArray.from_storage(
|
||||
json_type,
|
||||
pa.array([value for _, value in rows], type=json_type.storage_type),
|
||||
)
|
||||
return pa.Table.from_arrays(
|
||||
[pa.array([row_id for row_id, _ in rows]), json_values], schema=schema
|
||||
)
|
||||
schema = pa.schema([pa.field("id", pa.string()), pa.field("j", pa.json_())])
|
||||
|
||||
table = await mem_db_async.create_table("json_merge", schema=schema)
|
||||
await table.add(json_table([("a", '{"k": 1}'), ("b", '{"k": 9}')]))
|
||||
await table.add(_json_arrow_table(schema, [("a", '{"k": 1}'), ("b", '{"k": 9}')]))
|
||||
|
||||
await (
|
||||
table.merge_insert("id")
|
||||
.when_matched_update_all()
|
||||
.execute(json_table([("a", '{"k": 2}')]))
|
||||
.execute(_json_arrow_table(schema, [("a", '{"k": 2}')]))
|
||||
)
|
||||
|
||||
rows = sorted(await table.query().to_list(), key=lambda row: row["id"])
|
||||
@@ -3015,20 +3016,176 @@ async def test_merge_insert_encodes_json(mem_db_async: AsyncConnection):
|
||||
@pytest.mark.skipif(not hasattr(pa, "json_"), reason="requires PyArrow JSON type")
|
||||
@pytest.mark.asyncio
|
||||
async def test_add_sanitization_encodes_json(mem_db_async: AsyncConnection):
|
||||
json_type = pa.json_()
|
||||
schema = pa.schema([pa.field("id", pa.string()), pa.field("j", json_type)])
|
||||
json_values = pa.ExtensionArray.from_storage(
|
||||
json_type, pa.array(['{"k": 3}'], type=json_type.storage_type)
|
||||
)
|
||||
data = pa.Table.from_arrays([pa.array(["c"]), json_values], schema=schema)
|
||||
schema = pa.schema([pa.field("id", pa.string()), pa.field("j", pa.json_())])
|
||||
|
||||
table = await mem_db_async.create_table("json_add", schema=schema)
|
||||
await table.add(data, on_bad_vectors="fill")
|
||||
await table.add(
|
||||
_json_arrow_table(schema, [("c", '{"k": 3}')]), on_bad_vectors="fill"
|
||||
)
|
||||
|
||||
rows = await table.query().where("json_extract(j, '$.k') = '3'").to_list()
|
||||
assert rows == [{"id": "c", "j": '{"k":3}'}]
|
||||
|
||||
|
||||
@pytest.mark.skipif(not hasattr(pa, "json_"), reason="requires PyArrow JSON type")
|
||||
@pytest.mark.asyncio
|
||||
async def test_add_all_null_json_batch(mem_db_async: AsyncConnection):
|
||||
"""A batch of dicts whose json values are all None infers as pa.null(), which used
|
||||
to fail with a `json` vs `large_binary` schema mismatch. A row-at-a-time insert of
|
||||
an optional json column always looks like this."""
|
||||
schema = pa.schema([pa.field("id", pa.string()), pa.field("j", pa.json_())])
|
||||
table = await mem_db_async.create_table("json_nulls", schema=schema)
|
||||
|
||||
await table.add([{"id": "a", "j": None}])
|
||||
assert await table.count_rows() == 1
|
||||
|
||||
# ... and again once real JSON has been written.
|
||||
await table.add(_json_arrow_table(schema, [("b", '{"k": 9}')]))
|
||||
await table.add([{"id": "c", "j": None}])
|
||||
|
||||
rows = sorted(await table.query().to_list(), key=lambda row: row["id"])
|
||||
assert rows == [
|
||||
{"id": "a", "j": None},
|
||||
{"id": "b", "j": '{"k":9}'},
|
||||
{"id": "c", "j": None},
|
||||
]
|
||||
|
||||
# The nulls must not disturb reads of the column.
|
||||
filtered = await table.query().where("json_extract(j, '$.k') = '9'").to_list()
|
||||
assert filtered == [{"id": "b", "j": '{"k":9}'}]
|
||||
assert len(await table.query().where("j IS NULL").to_list()) == 2
|
||||
|
||||
|
||||
@pytest.mark.skipif(not hasattr(pa, "json_"), reason="requires PyArrow JSON type")
|
||||
def test_add_all_null_json_batch_sync(mem_db: DBConnection):
|
||||
schema = pa.schema([pa.field("id", pa.string()), pa.field("j", pa.json_())])
|
||||
table = mem_db.create_table("json_nulls_sync", schema=schema)
|
||||
|
||||
table.add([{"id": "a", "j": None}])
|
||||
|
||||
assert table.count_rows() == 1
|
||||
assert table.to_arrow()["j"].to_pylist() == [None]
|
||||
|
||||
|
||||
@pytest.mark.skipif(not hasattr(pa, "json_"), reason="requires PyArrow JSON type")
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
("values", "expected"),
|
||||
[
|
||||
([None], [None]),
|
||||
([None, '{"k": 1}'], [None, '{"k":1}']),
|
||||
(['{"k": 2}'], ['{"k":2}']),
|
||||
],
|
||||
)
|
||||
async def test_add_list_of_dicts_to_json_column(
|
||||
mem_db_async: AsyncConnection, values, expected
|
||||
):
|
||||
schema = pa.schema([pa.field("id", pa.int64()), pa.field("value", pa.json_())])
|
||||
table = await mem_db_async.create_table("json_list_add", schema=schema)
|
||||
|
||||
await table.add([{"id": idx, "value": value} for idx, value in enumerate(values)])
|
||||
|
||||
rows = (await table.to_arrow()).sort_by("id").to_pylist()
|
||||
assert [row["value"] for row in rows] == expected
|
||||
|
||||
|
||||
@pytest.mark.skipif(not hasattr(pa, "json_"), reason="requires PyArrow JSON type")
|
||||
@pytest.mark.asyncio
|
||||
async def test_add_list_of_dicts_to_nested_json_column(
|
||||
mem_db_async: AsyncConnection,
|
||||
):
|
||||
json_field = pa.field("value", pa.json_())
|
||||
info_field = pa.field("info", pa.struct([json_field]))
|
||||
info = pa.StructArray.from_arrays(
|
||||
[pa.array(['{"seed": 0}'], type=pa.json_())], fields=[json_field]
|
||||
)
|
||||
seed = pa.Table.from_arrays(
|
||||
[pa.array([0], type=pa.int64()), info],
|
||||
schema=pa.schema([pa.field("id", pa.int64()), info_field]),
|
||||
)
|
||||
table = await mem_db_async.create_table("nested_json_list_add", data=seed)
|
||||
|
||||
await table.add([{"id": 1, "info": {"value": '{"k": 1}'}}])
|
||||
await table.add([{"id": 2, "info": {"value": '{"k": 2}'}}], on_bad_vectors="fill")
|
||||
|
||||
rows = (await table.to_arrow()).sort_by("id").to_pylist()
|
||||
assert rows == [
|
||||
{"id": 0, "info": {"value": '{"seed":0}'}},
|
||||
{"id": 1, "info": {"value": '{"k":1}'}},
|
||||
{"id": 2, "info": {"value": '{"k":2}'}},
|
||||
]
|
||||
|
||||
|
||||
@pytest.mark.skipif(not hasattr(pa, "json_"), reason="requires PyArrow JSON type")
|
||||
@pytest.mark.asyncio
|
||||
async def test_add_list_of_dicts_to_json_list_column(mem_db_async: AsyncConnection):
|
||||
"""JSON inside a list must be JSONB-encoded, not stored as the raw text.
|
||||
|
||||
Storing raw text appends without error but leaves the column unreadable, so the
|
||||
round trip is checked with a filter as well as by value.
|
||||
"""
|
||||
schema = pa.schema(
|
||||
[
|
||||
pa.field("id", pa.int64()),
|
||||
pa.field("docs", pa.list_(pa.field("item", pa.json_()))),
|
||||
]
|
||||
)
|
||||
table = await mem_db_async.create_table("json_list_add", schema=schema)
|
||||
|
||||
await table.add([{"id": 1, "docs": ['{"k": 1}', '{"k": 2}']}])
|
||||
await table.add([{"id": 2, "docs": ['{"k": 3}']}], on_bad_vectors="fill")
|
||||
|
||||
rows = (await table.to_arrow()).sort_by("id").to_pylist()
|
||||
assert rows == [
|
||||
{"id": 1, "docs": ['{"k":1}', '{"k":2}']},
|
||||
{"id": 2, "docs": ['{"k":3}']},
|
||||
]
|
||||
|
||||
matched = await table.query().where("json_extract(docs[1], '$.k') = 3").to_arrow()
|
||||
assert matched.column("id").to_pylist() == [2]
|
||||
|
||||
|
||||
@pytest.mark.skipif(not hasattr(pa, "json_"), reason="requires PyArrow JSON type")
|
||||
@pytest.mark.asyncio
|
||||
async def test_add_map_of_json_values(mem_db_async: AsyncConnection):
|
||||
"""JSON in a map's values needs the same encoding a list's items do."""
|
||||
schema = pa.schema(
|
||||
[
|
||||
pa.field("id", pa.int64()),
|
||||
pa.field("m", pa.map_(pa.string(), pa.json_())),
|
||||
]
|
||||
)
|
||||
table = await mem_db_async.create_table(
|
||||
"json_map_add",
|
||||
schema=schema,
|
||||
storage_options={"new_table_data_storage_version": "2.2"},
|
||||
)
|
||||
|
||||
def batch(row_id: int, text: str) -> pa.Table:
|
||||
return pa.table(
|
||||
{
|
||||
"id": pa.array([row_id], type=pa.int64()),
|
||||
"m": pa.array([[("k", text)]], type=pa.map_(pa.string(), pa.string())),
|
||||
}
|
||||
)
|
||||
|
||||
await table.add(batch(1, '{"x": 1}'))
|
||||
await table.add(batch(2, '{"x": 2}'), on_bad_vectors="fill")
|
||||
|
||||
rows = (await table.to_arrow()).sort_by("id").to_pylist()
|
||||
assert rows == [
|
||||
{"id": 1, "m": [("k", '{"x":1}')]},
|
||||
{"id": 2, "m": [("k", '{"x":2}')]},
|
||||
]
|
||||
|
||||
matched = (
|
||||
await table.query()
|
||||
.where("json_extract(element_at(m, 'k')[1], '$.x') = '2'")
|
||||
.to_arrow()
|
||||
)
|
||||
assert matched.column("id").to_pylist() == [2]
|
||||
|
||||
|
||||
def test_create_with_embedding_function(mem_db: DBConnection):
|
||||
class MyTable(LanceModel):
|
||||
text: str
|
||||
@@ -4170,6 +4327,13 @@ async def test_computed_column_async(tmp_path):
|
||||
assert (await table.to_arrow())["tripled"].to_pylist() == [9]
|
||||
|
||||
|
||||
def test_function_errors_are_remote_only(tmp_path):
|
||||
db = lancedb.connect(tmp_path)
|
||||
table = db.create_table("t", [{"x": 1}])
|
||||
with pytest.raises(NotImplementedError, match="LanceDB Cloud and Enterprise"):
|
||||
table.function_errors()
|
||||
|
||||
|
||||
def test_refresh_column_async_returns_job(tmp_path):
|
||||
db = lancedb.connect(tmp_path)
|
||||
table = db.create_table("computed_job", [{"x": 1}, {"x": 2}])
|
||||
@@ -4183,13 +4347,14 @@ def test_refresh_column_async_returns_job(tmp_path):
|
||||
assert result.rows_failed == 0
|
||||
assert result.rows_remaining == 0
|
||||
assert result.source_version == 2
|
||||
assert result.published_version == 3
|
||||
# The fill lands at 3; the stamp recording its inputs is published at 4.
|
||||
assert result.published_version == 4
|
||||
assert job.status() == "finished"
|
||||
assert sorted(table.to_arrow()["doubled"].to_pylist()) == [2, 4]
|
||||
|
||||
no_op = table.refresh_column_async("doubled").wait()
|
||||
assert no_op.rows_assigned == 0
|
||||
assert no_op.source_version == 3
|
||||
assert no_op.source_version == 4
|
||||
assert no_op.published_version is None
|
||||
|
||||
# Bad input raises at the call, not through the job.
|
||||
@@ -4208,6 +4373,6 @@ async def test_refresh_column_async_job_async_table(tmp_path):
|
||||
assert isinstance(result, lancedb.RefreshColumnResult)
|
||||
assert result.rows_assigned == 1
|
||||
assert result.source_version == 2
|
||||
assert result.published_version == 3
|
||||
assert result.published_version == 4
|
||||
assert await job.status() == "finished"
|
||||
assert (await table.to_arrow())["tripled"].to_pylist() == [9]
|
||||
|
||||
@@ -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()?,
|
||||
})
|
||||
})
|
||||
}
|
||||
+146
-3
@@ -381,7 +381,7 @@ impl Connection {
|
||||
})
|
||||
}
|
||||
|
||||
#[pyo3(signature = (name, source, projections=None, filter=None, limit=None))]
|
||||
#[pyo3(signature = (name, source, projections=None, filter=None, limit=None, with_no_data=false))]
|
||||
pub fn create_materialized_view(
|
||||
self_: PyRef<'_, Self>,
|
||||
name: String,
|
||||
@@ -389,6 +389,7 @@ impl Connection {
|
||||
projections: Option<Vec<(String, String)>>,
|
||||
filter: Option<String>,
|
||||
limit: Option<u64>,
|
||||
with_no_data: bool,
|
||||
) -> PyResult<Bound<'_, PyAny>> {
|
||||
let inner = self_.get_inner()?.clone();
|
||||
future_into_py(self_.py(), async move {
|
||||
@@ -402,16 +403,79 @@ impl Connection {
|
||||
if let Some(limit) = limit {
|
||||
builder = builder.limit(limit);
|
||||
}
|
||||
let view = builder.execute().await.infer_error()?;
|
||||
builder = builder.with_no_data(with_no_data);
|
||||
let view = Box::pin(builder.execute()).await.infer_error()?;
|
||||
Ok(Table::new(view.table().clone()))
|
||||
})
|
||||
}
|
||||
|
||||
#[pyo3(signature = (name, source, projections=None, filter=None, limit=None, with_no_data=false))]
|
||||
pub fn create_materialized_view_async(
|
||||
self_: PyRef<'_, Self>,
|
||||
name: String,
|
||||
source: String,
|
||||
projections: Option<Vec<(String, String)>>,
|
||||
filter: Option<String>,
|
||||
limit: Option<u64>,
|
||||
with_no_data: bool,
|
||||
) -> PyResult<Bound<'_, PyAny>> {
|
||||
let inner = self_.get_inner()?.clone();
|
||||
future_into_py(self_.py(), async move {
|
||||
let mut builder = inner.create_materialized_view(name, source);
|
||||
if let Some(projections) = projections {
|
||||
builder = builder.select(projections);
|
||||
}
|
||||
if let Some(filter) = filter {
|
||||
builder = builder.only_if(filter);
|
||||
}
|
||||
if let Some(limit) = limit {
|
||||
builder = builder.limit(limit);
|
||||
}
|
||||
let job = Box::pin(builder.with_no_data(with_no_data).execute_async())
|
||||
.await
|
||||
.infer_error()?;
|
||||
Ok(crate::job::Job::new(job))
|
||||
})
|
||||
}
|
||||
|
||||
pub fn list_materialized_views(self_: PyRef<'_, Self>) -> PyResult<Bound<'_, PyAny>> {
|
||||
let inner = self_.get_inner()?.clone();
|
||||
future_into_py(self_.py(), async move {
|
||||
let views = inner.list_materialized_views().await.infer_error()?;
|
||||
Ok(views.into_iter().map(|view| view.name).collect::<Vec<_>>())
|
||||
Ok(views)
|
||||
})
|
||||
}
|
||||
|
||||
#[pyo3(signature = (name, namespace_path=None))]
|
||||
pub fn drop_materialized_view(
|
||||
self_: PyRef<'_, Self>,
|
||||
name: String,
|
||||
namespace_path: Option<Vec<String>>,
|
||||
) -> PyResult<Bound<'_, PyAny>> {
|
||||
let inner = self_.get_inner()?.clone();
|
||||
let namespace_path = namespace_path.unwrap_or_default();
|
||||
future_into_py(self_.py(), async move {
|
||||
inner
|
||||
.drop_materialized_view(name, &namespace_path)
|
||||
.await
|
||||
.infer_error()
|
||||
})
|
||||
}
|
||||
|
||||
#[pyo3(signature = (name, namespace_path=None))]
|
||||
pub fn drop_materialized_view_async(
|
||||
self_: PyRef<'_, Self>,
|
||||
name: String,
|
||||
namespace_path: Option<Vec<String>>,
|
||||
) -> PyResult<Bound<'_, PyAny>> {
|
||||
let inner = self_.get_inner()?.clone();
|
||||
let namespace_path = namespace_path.unwrap_or_default();
|
||||
future_into_py(self_.py(), async move {
|
||||
inner
|
||||
.drop_materialized_view_async(name, &namespace_path)
|
||||
.await
|
||||
.infer_error()
|
||||
.map(crate::job::Job::new)
|
||||
})
|
||||
}
|
||||
|
||||
@@ -704,6 +768,85 @@ impl Connection {
|
||||
})
|
||||
}
|
||||
|
||||
#[pyo3(signature = (name, value, namespace_path=None))]
|
||||
pub fn create_secret(
|
||||
self_: PyRef<'_, Self>,
|
||||
name: String,
|
||||
value: String,
|
||||
namespace_path: Option<Vec<String>>,
|
||||
) -> PyResult<Bound<'_, PyAny>> {
|
||||
let inner = self_.get_inner()?.clone();
|
||||
let namespace_path = namespace_path.unwrap_or_default();
|
||||
future_into_py(self_.py(), async move {
|
||||
inner
|
||||
.create_secret(name, value, &namespace_path)
|
||||
.await
|
||||
.infer_error()
|
||||
})
|
||||
}
|
||||
|
||||
#[pyo3(signature = (name, value, namespace_path=None))]
|
||||
pub fn alter_secret(
|
||||
self_: PyRef<'_, Self>,
|
||||
name: String,
|
||||
value: String,
|
||||
namespace_path: Option<Vec<String>>,
|
||||
) -> PyResult<Bound<'_, PyAny>> {
|
||||
let inner = self_.get_inner()?.clone();
|
||||
let namespace_path = namespace_path.unwrap_or_default();
|
||||
future_into_py(self_.py(), async move {
|
||||
inner
|
||||
.alter_secret(name, value, &namespace_path)
|
||||
.await
|
||||
.infer_error()
|
||||
})
|
||||
}
|
||||
|
||||
#[pyo3(signature = (namespace_path=None))]
|
||||
pub fn list_secrets(
|
||||
self_: PyRef<'_, Self>,
|
||||
namespace_path: Option<Vec<String>>,
|
||||
) -> PyResult<Bound<'_, PyAny>> {
|
||||
let inner = self_.get_inner()?.clone();
|
||||
let namespace_path = namespace_path.unwrap_or_default();
|
||||
future_into_py(self_.py(), async move {
|
||||
inner.list_secrets(&namespace_path).await.infer_error()
|
||||
})
|
||||
}
|
||||
|
||||
#[pyo3(signature = (name, namespace_path=None))]
|
||||
pub fn drop_secret(
|
||||
self_: PyRef<'_, Self>,
|
||||
name: String,
|
||||
namespace_path: Option<Vec<String>>,
|
||||
) -> PyResult<Bound<'_, PyAny>> {
|
||||
let inner = self_.get_inner()?.clone();
|
||||
let namespace_path = namespace_path.unwrap_or_default();
|
||||
future_into_py(self_.py(), async move {
|
||||
inner.drop_secret(name, &namespace_path).await.infer_error()
|
||||
})
|
||||
}
|
||||
|
||||
/// Name and timestamps as a plain tuple. `SecretInfo` carries no value, so
|
||||
/// there is none to filter out here. Timestamps stay integers rather than
|
||||
/// going through a string, so the caller can compare two without parsing.
|
||||
#[pyo3(signature = (name, namespace_path=None))]
|
||||
pub fn describe_secret(
|
||||
self_: PyRef<'_, Self>,
|
||||
name: String,
|
||||
namespace_path: Option<Vec<String>>,
|
||||
) -> PyResult<Bound<'_, PyAny>> {
|
||||
let inner = self_.get_inner()?.clone();
|
||||
let namespace_path = namespace_path.unwrap_or_default();
|
||||
future_into_py(self_.py(), async move {
|
||||
let info = inner
|
||||
.describe_secret(name, &namespace_path)
|
||||
.await
|
||||
.infer_error()?;
|
||||
Ok((info.name, info.created_at_millis, info.updated_at_millis))
|
||||
})
|
||||
}
|
||||
|
||||
pub fn list_jobs(self_: PyRef<'_, Self>) -> PyResult<Bound<'_, PyAny>> {
|
||||
let inner = self_.get_inner()?.clone();
|
||||
future_into_py(self_.py(), async move {
|
||||
|
||||
@@ -29,7 +29,10 @@ impl<T> PythonErrorExt<T> for std::result::Result<T, LanceError> {
|
||||
LanceError::InvalidInput { .. }
|
||||
| LanceError::InvalidTableName { .. }
|
||||
| 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())),
|
||||
|
||||
+35
-5
@@ -8,19 +8,21 @@ use expr::{PyExpr, expr_col, expr_func, expr_lit};
|
||||
use index::IndexConfig;
|
||||
use permutation::{PyAsyncPermutationBuilder, PyPermutationReader};
|
||||
use pyo3::{
|
||||
Bound, PyResult, Python, pymodule,
|
||||
types::{PyModule, PyModuleMethods},
|
||||
Bound, PyResult, Python, pyfunction, pymodule,
|
||||
types::{PyAnyMethods, PyModule, PyModuleMethods},
|
||||
wrap_pyfunction,
|
||||
};
|
||||
use query::{FTSQuery, HybridQuery, Query, VectorQuery};
|
||||
use session::Session;
|
||||
use table::{
|
||||
AddColumnsResult, AddResult, AlterColumnsResult, DeleteResult, DropColumnsResult, FtsToken,
|
||||
LsmWriteSpec, MergeResult, PyBlobFile, RefreshColumnResult, RefreshMaterializedViewResult,
|
||||
Table, UpdateFieldMetadataResult, UpdateResult,
|
||||
FunctionErrorFragment, FunctionErrorRecord, FunctionErrors, LsmWriteSpec, MergeResult,
|
||||
PyBlobFile, RefreshColumnResult, RefreshMaterializedViewResult, Table,
|
||||
UpdateFieldMetadataResult, UpdateResult,
|
||||
};
|
||||
|
||||
pub mod arrow;
|
||||
pub mod catalog;
|
||||
pub mod connection;
|
||||
pub mod error;
|
||||
pub mod expr;
|
||||
@@ -38,15 +40,36 @@ pub mod sql;
|
||||
pub mod table;
|
||||
pub mod util;
|
||||
|
||||
/// Shut down the shared Tokio runtime (see `runtime::shutdown`).
|
||||
///
|
||||
/// Registered below as a Python `atexit` callback rather than called
|
||||
/// directly: `atexit` runs while the interpreter is still fully valid,
|
||||
/// which is the coordinated, bounded exit the runtime otherwise never gets.
|
||||
///
|
||||
/// Runs the actual wait with the GIL released (`Python::detach`): shutdown
|
||||
/// blocks the calling thread waiting on the runtime's own worker threads,
|
||||
/// and if any in-flight task needs the GIL to finish (e.g. one that calls
|
||||
/// back into Python), holding it here while waiting on that same task would
|
||||
/// deadlock rather than time out.
|
||||
#[pyfunction]
|
||||
fn shutdown_runtime(py: Python<'_>) {
|
||||
py.detach(|| runtime::shutdown(std::time::Duration::from_secs(5)));
|
||||
}
|
||||
|
||||
#[pymodule]
|
||||
pub fn _lancedb(_py: Python, m: &Bound<'_, PyModule>) -> PyResult<()> {
|
||||
pub fn _lancedb(py: Python, m: &Bound<'_, PyModule>) -> PyResult<()> {
|
||||
let env = Env::new()
|
||||
.filter_or("LANCEDB_LOG", "warn")
|
||||
.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>()?;
|
||||
m.add_class::<crate::oauth::PySessionStatus>()?;
|
||||
m.add_class::<crate::oauth::PySessionLogout>()?;
|
||||
m.add_class::<crate::job::Job>()?;
|
||||
m.add_class::<crate::job::JobInfo>()?;
|
||||
m.add_class::<crate::job::JobDescription>()?;
|
||||
@@ -62,6 +85,9 @@ pub fn _lancedb(_py: Python, m: &Bound<'_, PyModule>) -> PyResult<()> {
|
||||
m.add_class::<RecordBatchStream>()?;
|
||||
m.add_class::<AddColumnsResult>()?;
|
||||
m.add_class::<RefreshColumnResult>()?;
|
||||
m.add_class::<FunctionErrors>()?;
|
||||
m.add_class::<FunctionErrorRecord>()?;
|
||||
m.add_class::<FunctionErrorFragment>()?;
|
||||
m.add_class::<RefreshMaterializedViewResult>()?;
|
||||
m.add_class::<AlterColumnsResult>()?;
|
||||
m.add_class::<UpdateFieldMetadataResult>()?;
|
||||
@@ -95,5 +121,9 @@ pub fn _lancedb(_py: Python, m: &Bound<'_, PyModule>) -> PyResult<()> {
|
||||
m.add_function(wrap_pyfunction!(expr_lit, m)?)?;
|
||||
m.add_function(wrap_pyfunction!(expr_func, m)?)?;
|
||||
m.add("__version__", env!("CARGO_PKG_VERSION"))?;
|
||||
// Give the shared runtime a coordinated, bounded shutdown at normal
|
||||
// process exit -- see `shutdown_runtime` and `runtime::shutdown` for why.
|
||||
py.import("atexit")?
|
||||
.call_method1("register", (wrap_pyfunction!(shutdown_runtime, m)?,))?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
+319
-6
@@ -1,10 +1,33 @@
|
||||
// SPDX-License-Identifier: Apache-2.0
|
||||
// SPDX-FileCopyrightText: Copyright The LanceDB Authors
|
||||
|
||||
use pyo3::FromPyObject;
|
||||
use std::path::PathBuf;
|
||||
use std::sync::Arc;
|
||||
|
||||
use pyo3::{FromPyObject, PyResult, Python, pyclass, pymethods};
|
||||
|
||||
use crate::error::PythonErrorExt;
|
||||
use crate::runtime::future_into_py;
|
||||
use lancedb::error::Error;
|
||||
use lancedb::remote::oauth::{OAuthConfig, OAuthFlow};
|
||||
use lancedb::remote::oauth::{AuthorizationCodeOptions, ClientAuthMethod, OAuthConfig, OAuthFlow};
|
||||
use lancedb::remote::{OAuthSession, SessionLogout, SessionStatus, TokenCacheOptions};
|
||||
|
||||
/// Python-side persistent token cache options, extracted via FromPyObject.
|
||||
/// Maps to `lancedb.remote.oauth.TokenCacheOptions` Python dataclass.
|
||||
#[derive(FromPyObject, Default)]
|
||||
pub struct PyTokenCacheOptions {
|
||||
pub cache_dir: Option<String>,
|
||||
pub lock_timeout_secs: Option<u64>,
|
||||
}
|
||||
|
||||
impl From<PyTokenCacheOptions> for TokenCacheOptions {
|
||||
fn from(py: PyTokenCacheOptions) -> Self {
|
||||
Self {
|
||||
cache_dir: py.cache_dir.map(PathBuf::from),
|
||||
lock_timeout_secs: py.lock_timeout_secs,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Python-side OAuth configuration, extracted via FromPyObject.
|
||||
/// Maps to `lancedb.remote.oauth.OAuthConfig` Python dataclass.
|
||||
@@ -13,10 +36,19 @@ pub struct PyOAuthConfig {
|
||||
pub issuer_url: String,
|
||||
pub client_id: String,
|
||||
pub scopes: Vec<String>,
|
||||
/// Optional resource indicator for authorization and token requests.
|
||||
pub resource: Option<String>,
|
||||
/// Optional provider-specific audience for authorization and token requests.
|
||||
pub audience: Option<String>,
|
||||
pub flow: String,
|
||||
pub client_secret: Option<String>,
|
||||
pub client_auth_method: Option<String>,
|
||||
pub redirect_uri: Option<String>,
|
||||
pub callback_port: Option<u16>,
|
||||
pub use_pkce: bool,
|
||||
pub managed_identity_client_id: Option<String>,
|
||||
pub refresh_buffer_secs: Option<u64>,
|
||||
pub token_cache: Option<PyTokenCacheOptions>,
|
||||
}
|
||||
|
||||
impl TryFrom<PyOAuthConfig> for OAuthConfig {
|
||||
@@ -25,6 +57,17 @@ impl TryFrom<PyOAuthConfig> for OAuthConfig {
|
||||
fn try_from(py: PyOAuthConfig) -> Result<Self, Self::Error> {
|
||||
let flow = match py.flow.as_str() {
|
||||
"client_credentials" => OAuthFlow::ClientCredentials,
|
||||
"authorization_code" => {
|
||||
let mut options = AuthorizationCodeOptions::new().use_pkce(py.use_pkce);
|
||||
if let Some(redirect_uri) = py.redirect_uri {
|
||||
options = options.redirect_uri(redirect_uri);
|
||||
}
|
||||
if let Some(callback_port) = py.callback_port {
|
||||
options = options.callback_port(callback_port);
|
||||
}
|
||||
OAuthFlow::AuthorizationCode(options)
|
||||
}
|
||||
"device_code" => OAuthFlow::DeviceCode,
|
||||
"azure_managed_identity" => OAuthFlow::AzureManagedIdentity {
|
||||
client_id: py.managed_identity_client_id,
|
||||
},
|
||||
@@ -35,13 +78,182 @@ impl TryFrom<PyOAuthConfig> for OAuthConfig {
|
||||
}
|
||||
};
|
||||
|
||||
let client_auth_method = match py.client_auth_method.as_deref() {
|
||||
Some("none") => Some(ClientAuthMethod::None),
|
||||
Some("client_secret_basic") => Some(ClientAuthMethod::ClientSecretBasic),
|
||||
Some("client_secret_post") => Some(ClientAuthMethod::ClientSecretPost),
|
||||
None => None,
|
||||
Some(other) => {
|
||||
return Err(Error::InvalidInput {
|
||||
message: format!("Unknown OAuth client auth method: {other}"),
|
||||
});
|
||||
}
|
||||
};
|
||||
|
||||
Ok(Self {
|
||||
issuer_url: py.issuer_url,
|
||||
client_id: py.client_id,
|
||||
client_secret: py.client_secret,
|
||||
client_auth_method,
|
||||
scopes: py.scopes,
|
||||
resource: py.resource,
|
||||
audience: py.audience,
|
||||
flow,
|
||||
refresh_buffer_secs: py.refresh_buffer_secs,
|
||||
token_cache: py.token_cache.map(TokenCacheOptions::from),
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
/// Wrapper around [`lancedb::remote::SessionStatus`] exposing safe metadata.
|
||||
#[pyclass(name = "SessionStatus", skip_from_py_object)]
|
||||
#[derive(Clone)]
|
||||
pub struct PySessionStatus {
|
||||
inner: SessionStatus,
|
||||
}
|
||||
|
||||
#[pymethods]
|
||||
impl PySessionStatus {
|
||||
/// Whether a cached session exists that can obtain tokens without
|
||||
/// interactive authentication.
|
||||
#[getter]
|
||||
pub fn refreshable(&self) -> bool {
|
||||
self.inner.refreshable
|
||||
}
|
||||
|
||||
/// Canonical issuer URL of the cached session.
|
||||
#[getter]
|
||||
pub fn issuer_url(&self) -> String {
|
||||
self.inner.issuer_url.clone()
|
||||
}
|
||||
|
||||
/// Client ID of the cached session.
|
||||
#[getter]
|
||||
pub fn client_id(&self) -> String {
|
||||
self.inner.client_id.clone()
|
||||
}
|
||||
|
||||
/// Canonical (sorted, de-duplicated) scopes of the cached session.
|
||||
#[getter]
|
||||
pub fn scopes(&self) -> Vec<String> {
|
||||
self.inner.scopes.clone()
|
||||
}
|
||||
|
||||
/// Resource indicator used to obtain the cached session.
|
||||
#[getter]
|
||||
pub fn resource(&self) -> Option<String> {
|
||||
self.inner.resource.clone()
|
||||
}
|
||||
|
||||
/// Provider-specific audience used to obtain the cached session.
|
||||
#[getter]
|
||||
pub fn audience(&self) -> Option<String> {
|
||||
self.inner.audience.clone()
|
||||
}
|
||||
|
||||
/// Flow that produced the cached session.
|
||||
#[getter]
|
||||
pub fn flow(&self) -> String {
|
||||
self.inner.flow.clone()
|
||||
}
|
||||
|
||||
/// When the cached session was obtained, as Unix seconds.
|
||||
#[getter]
|
||||
pub fn obtained_at(&self) -> Option<u64> {
|
||||
self.inner.obtained_at
|
||||
}
|
||||
|
||||
pub fn __repr__(&self) -> String {
|
||||
format!(
|
||||
"SessionStatus(refreshable={}, issuer_url='{}', client_id='{}', flow='{}')",
|
||||
self.inner.refreshable, self.inner.issuer_url, self.inner.client_id, self.inner.flow
|
||||
)
|
||||
}
|
||||
}
|
||||
|
||||
impl From<SessionStatus> for PySessionStatus {
|
||||
fn from(inner: SessionStatus) -> Self {
|
||||
Self { inner }
|
||||
}
|
||||
}
|
||||
|
||||
/// Wrapper around [`lancedb::remote::SessionLogout`].
|
||||
#[pyclass(name = "SessionLogout", skip_from_py_object)]
|
||||
#[derive(Clone)]
|
||||
pub struct PySessionLogout {
|
||||
inner: SessionLogout,
|
||||
}
|
||||
|
||||
#[pymethods]
|
||||
impl PySessionLogout {
|
||||
/// Whether a cached credential was removed.
|
||||
#[getter]
|
||||
pub fn removed(&self) -> bool {
|
||||
self.inner.removed
|
||||
}
|
||||
|
||||
pub fn __repr__(&self) -> String {
|
||||
format!("SessionLogout(removed={})", self.inner.removed)
|
||||
}
|
||||
}
|
||||
|
||||
impl From<SessionLogout> for PySessionLogout {
|
||||
fn from(inner: SessionLogout) -> Self {
|
||||
Self { inner }
|
||||
}
|
||||
}
|
||||
|
||||
/// Wrapper around [`lancedb::remote::OAuthSession`].
|
||||
#[pyclass(name = "OAuthSession", skip_from_py_object)]
|
||||
#[derive(Clone)]
|
||||
pub struct PyOAuthSession {
|
||||
inner: Arc<OAuthSession>,
|
||||
}
|
||||
|
||||
#[pymethods]
|
||||
impl PyOAuthSession {
|
||||
/// Create a session manager for the given OAuth configuration.
|
||||
///
|
||||
/// The configuration must set ``token_cache`` options and use a flow that
|
||||
/// supports persistent sessions (authorization code or device code).
|
||||
#[new]
|
||||
pub fn new(config: PyOAuthConfig) -> PyResult<Self> {
|
||||
let config: OAuthConfig = config.try_into().infer_error()?;
|
||||
let inner = OAuthSession::new(config).infer_error()?;
|
||||
Ok(Self {
|
||||
inner: Arc::new(inner),
|
||||
})
|
||||
}
|
||||
|
||||
/// Eagerly run the configured authentication flow and store the session.
|
||||
pub fn login<'py>(&self, py: Python<'py>) -> PyResult<pyo3::Bound<'py, pyo3::PyAny>> {
|
||||
let inner = Arc::clone(&self.inner);
|
||||
future_into_py(py, async move {
|
||||
inner.login().await.map(PySessionStatus::from).infer_error()
|
||||
})
|
||||
}
|
||||
|
||||
/// Report whether a matching cached session exists, with safe metadata.
|
||||
pub fn status<'py>(&self, py: Python<'py>) -> PyResult<pyo3::Bound<'py, pyo3::PyAny>> {
|
||||
let inner = Arc::clone(&self.inner);
|
||||
future_into_py(py, async move {
|
||||
inner
|
||||
.status()
|
||||
.await
|
||||
.map(PySessionStatus::from)
|
||||
.infer_error()
|
||||
})
|
||||
}
|
||||
|
||||
/// Remove the matching local cached credential.
|
||||
pub fn logout<'py>(&self, py: Python<'py>) -> PyResult<pyo3::Bound<'py, pyo3::PyAny>> {
|
||||
let inner = Arc::clone(&self.inner);
|
||||
future_into_py(py, async move {
|
||||
inner
|
||||
.logout()
|
||||
.await
|
||||
.map(PySessionLogout::from)
|
||||
.infer_error()
|
||||
})
|
||||
}
|
||||
}
|
||||
@@ -50,16 +262,30 @@ impl TryFrom<PyOAuthConfig> for OAuthConfig {
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn test_unknown_oauth_flow_returns_invalid_input() {
|
||||
let config = PyOAuthConfig {
|
||||
fn base_config() -> PyOAuthConfig {
|
||||
PyOAuthConfig {
|
||||
issuer_url: "https://issuer.example.com".to_string(),
|
||||
client_id: "client-id".to_string(),
|
||||
scopes: vec!["scope".to_string()],
|
||||
flow: "typo".to_string(),
|
||||
flow: "device_code".to_string(),
|
||||
client_secret: None,
|
||||
client_auth_method: None,
|
||||
redirect_uri: None,
|
||||
callback_port: None,
|
||||
use_pkce: true,
|
||||
managed_identity_client_id: None,
|
||||
refresh_buffer_secs: None,
|
||||
resource: None,
|
||||
audience: None,
|
||||
token_cache: None,
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_unknown_oauth_flow_returns_invalid_input() {
|
||||
let config = PyOAuthConfig {
|
||||
flow: "typo".to_string(),
|
||||
..base_config()
|
||||
};
|
||||
|
||||
let err = OAuthConfig::try_from(config).unwrap_err();
|
||||
@@ -69,4 +295,91 @@ mod tests {
|
||||
if message == "Unknown OAuth flow type: typo"
|
||||
));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_authorization_code_conversion_preserves_options() {
|
||||
let config = PyOAuthConfig {
|
||||
flow: "authorization_code".to_string(),
|
||||
client_secret: Some("secret".to_string()),
|
||||
redirect_uri: Some("http://127.0.0.1:9000/callback".to_string()),
|
||||
callback_port: Some(9000),
|
||||
use_pkce: false,
|
||||
resource: Some("urn:resource".into()),
|
||||
audience: Some("audience".into()),
|
||||
..base_config()
|
||||
};
|
||||
|
||||
let converted = OAuthConfig::try_from(config).unwrap();
|
||||
let OAuthFlow::AuthorizationCode(options) = converted.flow else {
|
||||
panic!("expected authorization code flow");
|
||||
};
|
||||
assert_eq!(
|
||||
options.redirect_uri.as_deref(),
|
||||
Some("http://127.0.0.1:9000/callback")
|
||||
);
|
||||
assert_eq!(options.callback_port, Some(9000));
|
||||
assert!(!options.use_pkce);
|
||||
assert_eq!(converted.resource.as_deref(), Some("urn:resource"));
|
||||
assert_eq!(converted.audience.as_deref(), Some("audience"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_device_code_conversion() {
|
||||
let config = base_config();
|
||||
let converted = OAuthConfig::try_from(config).unwrap();
|
||||
assert!(matches!(converted.flow, OAuthFlow::DeviceCode));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_client_auth_method_conversion() {
|
||||
for (value, expected) in [
|
||||
("none", ClientAuthMethod::None),
|
||||
("client_secret_basic", ClientAuthMethod::ClientSecretBasic),
|
||||
("client_secret_post", ClientAuthMethod::ClientSecretPost),
|
||||
] {
|
||||
let config = PyOAuthConfig {
|
||||
client_auth_method: Some(value.to_string()),
|
||||
..base_config()
|
||||
};
|
||||
|
||||
let converted = OAuthConfig::try_from(config).unwrap();
|
||||
assert_eq!(converted.client_auth_method, Some(expected));
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_unknown_client_auth_method_returns_invalid_input() {
|
||||
let config = PyOAuthConfig {
|
||||
client_auth_method: Some("typo".to_string()),
|
||||
..base_config()
|
||||
};
|
||||
|
||||
let err = OAuthConfig::try_from(config).unwrap_err();
|
||||
assert!(matches!(
|
||||
err,
|
||||
Error::InvalidInput { message }
|
||||
if message == "Unknown OAuth client auth method: typo"
|
||||
));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_token_cache_conversion() {
|
||||
let config = PyOAuthConfig {
|
||||
resource: None,
|
||||
audience: None,
|
||||
token_cache: Some(PyTokenCacheOptions {
|
||||
cache_dir: Some("/tmp/oauth-cache".to_string()),
|
||||
lock_timeout_secs: Some(5),
|
||||
}),
|
||||
..base_config()
|
||||
};
|
||||
|
||||
let converted = OAuthConfig::try_from(config).unwrap();
|
||||
let cache = converted.token_cache.expect("token cache options");
|
||||
assert_eq!(
|
||||
cache.cache_dir.as_deref(),
|
||||
Some(std::path::Path::new("/tmp/oauth-cache"))
|
||||
);
|
||||
assert_eq!(cache.lock_timeout_secs, Some(5));
|
||||
}
|
||||
}
|
||||
|
||||
+460
-30
@@ -4,20 +4,84 @@
|
||||
//! Fork-safe wrapper around tokio + pyo3-async-runtimes.
|
||||
//!
|
||||
//! `pyo3_async_runtimes::tokio` keeps its multi-threaded runtime in a
|
||||
//! `OnceLock` that can never be replaced. Tokio's worker threads do not
|
||||
//! `OnceLock` that can never be replaced. Tokio's worker threads do not
|
||||
//! survive `fork()`, so once a child inherits a "frozen" runtime, every
|
||||
//! `future_into_py` call hangs forever.
|
||||
//! `future_into_py` call hangs forever. Normal (non-fork) process exit has
|
||||
//! its own gap: nothing tells the runtime to shut down, so its worker
|
||||
//! threads keep running, uncoordinated with the interpreter, right up until
|
||||
//! the process ends. If one of them is mid-task exactly as `Py_Finalize`
|
||||
//! starts tearing down interpreter state, it can panic on state that's
|
||||
//! already gone -- and since that happens on a background thread with no
|
||||
//! PyO3-wrapped call frame to catch it, Rust aborts the whole process
|
||||
//! rather than failing that one call. [`shutdown`], registered as a Python
|
||||
//! `atexit` callback, closes that gap by giving the runtime a coordinated,
|
||||
//! bounded exit while the interpreter is still fully valid.
|
||||
//!
|
||||
//! We sidestep the global by routing every future through our own
|
||||
//! [`LanceRuntime`] (a [`pyo3_async_runtimes::generic::Runtime`] impl) backed
|
||||
//! by an [`AtomicPtr`] to a tokio runtime that we own. A `pthread_atfork`
|
||||
//! child handler nulls the pointer; the next `spawn` rebuilds the runtime in
|
||||
//! the child. This mirrors the pattern used in the Lance Python bindings.
|
||||
//! Getting both of these right at once took a few tries; the design here
|
||||
//! rests on three separate mechanisms, each solving one problem the others
|
||||
//! cannot:
|
||||
//!
|
||||
//! **`OUTSTANDING`, not `Arc::strong_count`, decides when the runtime is
|
||||
//! idle.** Early versions tried to infer "is anything still using this
|
||||
//! runtime" from how many `Arc<Runtime>` clones existed. That signal is
|
||||
//! wrong in both directions: a clone taken only for the instant a task is
|
||||
//! *submitted* says nothing about whether that task has actually finished
|
||||
//! running (`Runtime::shutdown_timeout` gives spawned, non-blocking tasks
|
||||
//! no grace period at all -- a task "keeps running until it yields, then is
|
||||
//! dropped" -- so a reclaim landing right after submission would silently
|
||||
//! abandon it before it ever got to run); and a clone held for a task's
|
||||
//! *whole* lifetime can end up making that task the final owner of the
|
||||
//! `Runtime`, so completing it drops the `Runtime` from inside one of its
|
||||
//! own worker threads, which tokio itself forbids ("cannot drop a runtime
|
||||
//! in a context where blocking is not allowed") and panics. `OUTSTANDING`
|
||||
//! is an explicit courtesy counter instead: every top-level `spawn`,
|
||||
//! `spawn_blocking`, or `block_on` call increments it before it starts and
|
||||
//! decrements it (via [`OutstandingGuard`]) when it is truly done, entirely
|
||||
//! decoupled from how many `Arc` clones exist at any instant. `shutdown`
|
||||
//! waits for it to reach zero before ever touching the runtime, which also
|
||||
//! closes a narrower race: because the counter is incremented *before*
|
||||
//! `get_runtime()` is even called, an install already in progress when
|
||||
//! `shutdown` runs is never invisible to it the way an empty slot would be.
|
||||
//! If the counter never reaches zero within the bound, `shutdown` stops
|
||||
//! waiting and forces the retirement attempt anyway -- silently returning
|
||||
//! with the runtime and its workers still fully alive would just recreate
|
||||
//! the exact race this function exists to close, for any call slower than
|
||||
//! the grace period.
|
||||
//!
|
||||
//! **Tasks never hold an `Arc<Runtime>`.** Because `OUTSTANDING` (not
|
||||
//! reference counting) is what `shutdown` waits on, a top-level task only
|
||||
//! needs to carry an `OutstandingGuard` -- a token whose `Drop` is a plain
|
||||
//! atomic decrement -- not a clone of the runtime itself. That is what
|
||||
//! makes it impossible for a task's completion to become the final,
|
||||
//! worker-thread-side drop of the `Runtime`: nothing a task holds ever
|
||||
//! *is* the `Runtime`.
|
||||
//!
|
||||
//! **`atfork_child` touches nothing but a plain counter.** `future_into_py`
|
||||
//! spawns a task that, once running, spawns a second one to do the real
|
||||
//! work and awaits its `JoinHandle`; if a nested call re-resolved "the
|
||||
//! current runtime" independently, a reclaim landing between the two calls
|
||||
//! could bind them to different instances. `spawn`/`spawn_blocking` close
|
||||
//! that with `Handle::try_current`: a call already running on one of our
|
||||
//! worker threads stays pinned to that instance, so only the first,
|
||||
//! outermost call of a chain ever consults [`get_runtime`]. That leaves
|
||||
//! fork as the other place identity can change, and it has to be handled
|
||||
//! without ever calling into `ArcSwapOption` from the child handler itself
|
||||
//! -- `swap`/`compare_and_swap` reconcile reader "debts" internally (via
|
||||
//! thread-local state, and potentially an allocation), none of which is
|
||||
//! safe to run in a forked child that may have inherited another thread's
|
||||
//! lock mid-acquisition. `atfork_child` therefore does nothing but bump a
|
||||
//! bare `GENERATION` counter; [`get_runtime`] compares the generation its
|
||||
//! installed runtime was built in against the live counter on every call,
|
||||
//! from ordinary (non-signal) context, and treats a mismatch as "stale,
|
||||
//! rebuild" -- exactly the check `atfork_child` used to perform directly.
|
||||
|
||||
use std::future::Future;
|
||||
use std::pin::Pin;
|
||||
use std::sync::atomic::{AtomicBool, AtomicPtr, Ordering};
|
||||
use std::sync::Arc;
|
||||
use std::sync::atomic::{AtomicBool, AtomicU64, Ordering};
|
||||
use std::time::{Duration, Instant};
|
||||
|
||||
use arc_swap::ArcSwapOption;
|
||||
use pyo3::{Bound, PyAny, PyResult, Python, conversion::IntoPyObject};
|
||||
use pyo3_async_runtimes::{
|
||||
TaskLocals,
|
||||
@@ -25,8 +89,31 @@ use pyo3_async_runtimes::{
|
||||
};
|
||||
use tokio::{runtime, task};
|
||||
|
||||
static RUNTIME: AtomicPtr<runtime::Runtime> = AtomicPtr::new(std::ptr::null_mut());
|
||||
static RUNTIME_INSTALLING: AtomicBool = AtomicBool::new(false);
|
||||
/// A runtime tagged with the fork generation it was built in, so a stale
|
||||
/// (post-fork, dead-worker-threads) instance can be told apart from a live
|
||||
/// one without `atfork_child` ever having to touch it directly.
|
||||
struct Tagged {
|
||||
runtime: runtime::Runtime,
|
||||
generation: u64,
|
||||
}
|
||||
|
||||
impl std::ops::Deref for Tagged {
|
||||
type Target = runtime::Runtime;
|
||||
fn deref(&self) -> &runtime::Runtime {
|
||||
&self.runtime
|
||||
}
|
||||
}
|
||||
|
||||
static RUNTIME: ArcSwapOption<Tagged> = ArcSwapOption::const_empty();
|
||||
/// Bumped only by `atfork_child`, and only ever read elsewhere. This is the
|
||||
/// entire fork-safety mechanism: no lock, no allocation, no thread-local
|
||||
/// access -- just one atomic add, which is all a `pthread_atfork` child
|
||||
/// handler is ever safe to do.
|
||||
static GENERATION: AtomicU64 = AtomicU64::new(0);
|
||||
/// Count of top-level `spawn`/`spawn_blocking`/`block_on` calls that have
|
||||
/// started but not yet finished. See the module docs for why this, and not
|
||||
/// `Arc::strong_count`, is what `shutdown` waits on.
|
||||
static OUTSTANDING: AtomicU64 = AtomicU64::new(0);
|
||||
static ATFORK_INSTALLED: AtomicBool = AtomicBool::new(false);
|
||||
|
||||
fn create_runtime() -> runtime::Runtime {
|
||||
@@ -37,23 +124,67 @@ fn create_runtime() -> runtime::Runtime {
|
||||
.expect("Failed to build tokio runtime")
|
||||
}
|
||||
|
||||
fn get_runtime() -> &'static runtime::Runtime {
|
||||
/// Get a live, owned handle to the shared runtime, rebuilding it if the
|
||||
/// installed one predates the most recent `fork()`.
|
||||
fn get_runtime() -> Arc<Tagged> {
|
||||
let current_gen = GENERATION.load(Ordering::SeqCst);
|
||||
loop {
|
||||
let ptr = RUNTIME.load(Ordering::SeqCst);
|
||||
if !ptr.is_null() {
|
||||
return unsafe { &*ptr };
|
||||
let existing = RUNTIME.load_full();
|
||||
if let Some(existing) = &existing
|
||||
&& existing.generation == current_gen
|
||||
{
|
||||
return Arc::clone(existing);
|
||||
}
|
||||
if !RUNTIME_INSTALLING.fetch_or(true, Ordering::SeqCst) {
|
||||
break;
|
||||
if !ATFORK_INSTALLED.fetch_or(true, Ordering::SeqCst) {
|
||||
install_atfork();
|
||||
}
|
||||
std::thread::yield_now();
|
||||
// Built optimistically, outside any lock: on the rare race where two
|
||||
// threads both find the slot empty (or stale), one candidate wins
|
||||
// the compare-and-swap below and the other is simply dropped here,
|
||||
// tearing down its own (never shared, never used) worker pool the
|
||||
// ordinary way.
|
||||
let candidate = Arc::new(Tagged {
|
||||
runtime: create_runtime(),
|
||||
generation: current_gen,
|
||||
});
|
||||
let previous = RUNTIME.compare_and_swap(&existing, Some(Arc::clone(&candidate)));
|
||||
let won = match (&*previous, &existing) {
|
||||
(None, None) => true,
|
||||
(Some(prev), Some(exist)) => Arc::ptr_eq(prev, exist),
|
||||
_ => false,
|
||||
};
|
||||
if won {
|
||||
if let Some(stale) = existing {
|
||||
// A prior generation's runtime: its worker threads are dead
|
||||
// in this process (they do not survive fork), so dropping it
|
||||
// normally would try to join them and hang. Leak it instead.
|
||||
std::mem::forget(stale);
|
||||
}
|
||||
return candidate;
|
||||
}
|
||||
// Someone else's candidate (or a concurrent shutdown) won; go around
|
||||
// and reload.
|
||||
}
|
||||
if !ATFORK_INSTALLED.fetch_or(true, Ordering::SeqCst) {
|
||||
install_atfork();
|
||||
}
|
||||
|
||||
/// RAII token tracked by [`OUTSTANDING`]. Held for the duration of a
|
||||
/// top-level `block_on` call, or moved into a top-level `spawn`/
|
||||
/// `spawn_blocking` task so it decrements only once that task's *entire*
|
||||
/// body -- including anything it nested-spawns and awaits -- has run to
|
||||
/// completion or been dropped without completing.
|
||||
struct OutstandingGuard;
|
||||
|
||||
impl OutstandingGuard {
|
||||
fn new() -> Self {
|
||||
OUTSTANDING.fetch_add(1, Ordering::SeqCst);
|
||||
Self
|
||||
}
|
||||
}
|
||||
|
||||
impl Drop for OutstandingGuard {
|
||||
fn drop(&mut self) {
|
||||
OUTSTANDING.fetch_sub(1, Ordering::SeqCst);
|
||||
}
|
||||
let new_ptr = Box::into_raw(Box::new(create_runtime()));
|
||||
RUNTIME.store(new_ptr, Ordering::SeqCst);
|
||||
unsafe { &*new_ptr }
|
||||
}
|
||||
|
||||
/// Block the current thread on a future using the shared runtime.
|
||||
@@ -62,16 +193,78 @@ fn get_runtime() -> &'static runtime::Runtime {
|
||||
/// building a namespace client). Must not be called from within the runtime's
|
||||
/// own worker threads.
|
||||
pub fn block_on<F: std::future::Future>(fut: F) -> F::Output {
|
||||
let _guard = OutstandingGuard::new();
|
||||
get_runtime().block_on(fut)
|
||||
}
|
||||
|
||||
/// Runs in async-signal context after `fork()` in the child. We can only
|
||||
/// touch atomics here; we deliberately leak the previous runtime because
|
||||
/// dropping a tokio `Runtime` would try to join its (now-dead) worker
|
||||
/// threads and hang.
|
||||
/// Gracefully quiesce the shared runtime, meant to run at normal process exit.
|
||||
///
|
||||
/// Waits (bounded by `timeout`) for [`OUTSTANDING`] to reach zero -- i.e.
|
||||
/// for every top-level call already under way to actually finish, not just
|
||||
/// for `Arc::strong_count` to look low -- before ever touching the runtime.
|
||||
/// If the bound elapses first, it stops waiting and attempts retirement
|
||||
/// anyway: leaving the runtime and its worker threads untouched would just
|
||||
/// recreate the exact race this function exists to close, for any call
|
||||
/// slower than the grace period.
|
||||
///
|
||||
/// Retirement itself removes the runtime from the slot and calls
|
||||
/// `shutdown_timeout` rather than a bare `drop`: dropping a tokio `Runtime`
|
||||
/// waits (in the worst case indefinitely) for its worker threads to join,
|
||||
/// whereas `shutdown_timeout` gives real in-flight work -- a connection
|
||||
/// pool's keep-alive, a graceful close -- a bounded chance to finish first,
|
||||
/// then forcibly ends whatever has not. Reclaiming can only proceed once
|
||||
/// `Arc::try_unwrap` proves no other reference remains; if some transient
|
||||
/// `get_runtime()` caller is, at that exact instant, still between loading
|
||||
/// the slot and finishing its own call, this abandons the runtime instead
|
||||
/// of forcing the issue -- the same trade `atfork_child` already makes.
|
||||
///
|
||||
/// Neither of the two ways this can fail to cleanly retire the runtime --
|
||||
/// the wait timing out, or `try_unwrap` losing that race -- has any other
|
||||
/// signal to report through (`shutdown_timeout` itself returns nothing),
|
||||
/// so both log a warning instead of failing silently.
|
||||
pub fn shutdown(timeout: Duration) {
|
||||
let deadline = Instant::now() + timeout;
|
||||
loop {
|
||||
let outstanding = OUTSTANDING.load(Ordering::SeqCst);
|
||||
if outstanding == 0 {
|
||||
break;
|
||||
}
|
||||
if Instant::now() >= deadline {
|
||||
log::warn!(
|
||||
"lancedb: runtime shutdown timed out with {outstanding} call(s) still in flight; forcing shutdown anyway, some in-flight work may be abandoned"
|
||||
);
|
||||
break;
|
||||
}
|
||||
std::thread::sleep(Duration::from_millis(1));
|
||||
}
|
||||
let Some(current) = RUNTIME.load_full() else {
|
||||
return;
|
||||
};
|
||||
RUNTIME.compare_and_swap(&Some(Arc::clone(¤t)), None);
|
||||
match Arc::try_unwrap(current) {
|
||||
Ok(tagged) => {
|
||||
tagged
|
||||
.runtime
|
||||
.shutdown_timeout(deadline.saturating_duration_since(Instant::now()));
|
||||
}
|
||||
Err(_) => {
|
||||
// Some transient `get_runtime()` caller is, at this exact
|
||||
// instant, still between loading the slot and finishing its own
|
||||
// call: we have no owned handle to call `shutdown_timeout` on,
|
||||
// and no way to force one (tokio has no shutdown API over a
|
||||
// shared reference). Nothing more to do but say so.
|
||||
log::warn!(
|
||||
"lancedb: runtime shutdown could not obtain exclusive ownership; the shared runtime was left running"
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Runs in async-signal context after `fork()` in the child. Touches
|
||||
/// nothing but a plain atomic add -- see the module docs for why even
|
||||
/// `ArcSwapOption::swap` is not safe to call here.
|
||||
extern "C" fn atfork_child() {
|
||||
RUNTIME.store(std::ptr::null_mut(), Ordering::SeqCst);
|
||||
RUNTIME_INSTALLING.store(false, Ordering::SeqCst);
|
||||
GENERATION.fetch_add(1, Ordering::SeqCst);
|
||||
}
|
||||
|
||||
#[cfg(not(windows))]
|
||||
@@ -102,11 +295,26 @@ impl Runtime for LanceRuntime {
|
||||
type JoinError = LanceJoinError;
|
||||
type JoinHandle = Pin<Box<dyn Future<Output = Result<(), Self::JoinError>> + Send>>;
|
||||
|
||||
/// `pyo3_async_runtimes::generic::future_into_py` spawns a task that,
|
||||
/// once it starts running, spawns a second one to do the real work and
|
||||
/// awaits its `JoinHandle`. `Handle::try_current` pins that nested call
|
||||
/// to whatever runtime is already executing it, so only the first,
|
||||
/// outermost call of a chain -- one running on a thread outside any
|
||||
/// runtime -- ever consults [`get_runtime`] or [`OUTSTANDING`].
|
||||
fn spawn<F>(fut: F) -> Self::JoinHandle
|
||||
where
|
||||
F: Future<Output = ()> + Send + 'static,
|
||||
{
|
||||
let handle = get_runtime().spawn(fut);
|
||||
let handle = match tokio::runtime::Handle::try_current() {
|
||||
Ok(handle) => handle.spawn(fut),
|
||||
Err(_) => {
|
||||
let guard = OutstandingGuard::new();
|
||||
get_runtime().spawn(async move {
|
||||
let _guard = guard;
|
||||
fut.await;
|
||||
})
|
||||
}
|
||||
};
|
||||
Box::pin(async move { handle.await.map_err(LanceJoinError) })
|
||||
}
|
||||
|
||||
@@ -114,7 +322,16 @@ impl Runtime for LanceRuntime {
|
||||
where
|
||||
F: FnOnce() + Send + 'static,
|
||||
{
|
||||
let handle = get_runtime().spawn_blocking(f);
|
||||
let handle = match tokio::runtime::Handle::try_current() {
|
||||
Ok(handle) => handle.spawn_blocking(f),
|
||||
Err(_) => {
|
||||
let guard = OutstandingGuard::new();
|
||||
get_runtime().spawn_blocking(move || {
|
||||
let _guard = guard;
|
||||
f();
|
||||
})
|
||||
}
|
||||
};
|
||||
Box::pin(async move { handle.await.map_err(LanceJoinError) })
|
||||
}
|
||||
}
|
||||
@@ -149,3 +366,216 @@ where
|
||||
{
|
||||
pyo3_async_runtimes::generic::future_into_py::<LanceRuntime, _, T>(py, fut)
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
// RUNTIME, GENERATION, OUTSTANDING, and ATFORK_INSTALLED are process-wide
|
||||
// statics, and Rust's test harness runs tests in parallel by default,
|
||||
// so separate test functions below would otherwise race each other
|
||||
// through this shared state (this reproduced in CI: one test's
|
||||
// in-flight task got reclaimed by a *different* test's concurrent
|
||||
// `shutdown()` call, and another observed `OUTSTANDING` left non-zero
|
||||
// by a still-running sibling). Every test takes this lock first so
|
||||
// only one of them touches the shared runtime state at a time; a
|
||||
// poisoned lock (a previous test's genuine failure) is still honored
|
||||
// rather than cascading into every later test as an unrelated panic.
|
||||
static TEST_LOCK: std::sync::Mutex<()> = std::sync::Mutex::new(());
|
||||
|
||||
fn lock_runtime_state_for_test() -> std::sync::MutexGuard<'static, ()> {
|
||||
TEST_LOCK
|
||||
.lock()
|
||||
.unwrap_or_else(|poisoned| poisoned.into_inner())
|
||||
}
|
||||
|
||||
// A task's own completion must never be the final drop of the shared
|
||||
// `Runtime`: tasks carry only an `OutstandingGuard` (a plain counter
|
||||
// token), never an `Arc<Runtime>`, specifically so this can't happen.
|
||||
// Getting this wrong panics ("cannot drop a runtime in a context where
|
||||
// blocking is not allowed") -- this reproduced unprompted, twice, in a
|
||||
// single run of `test_nested_spawn_survives_concurrent_shutdown` under
|
||||
// an earlier design that held the `Arc` for a task's whole lifetime.
|
||||
#[test]
|
||||
#[allow(unused_must_use)] // fire-and-forget spawn, same as future_into_py itself
|
||||
fn test_top_level_task_survives_concurrent_shutdown_reclaim() {
|
||||
use std::sync::mpsc;
|
||||
let _lock = lock_runtime_state_for_test();
|
||||
|
||||
for _ in 0..50 {
|
||||
let (tx, rx) = mpsc::channel::<()>();
|
||||
|
||||
LanceRuntime::spawn(async move {
|
||||
for _ in 0..5 {
|
||||
task::yield_now().await;
|
||||
}
|
||||
let _ = tx.send(());
|
||||
});
|
||||
|
||||
shutdown(Duration::from_secs(2));
|
||||
|
||||
rx.recv_timeout(Duration::from_secs(5))
|
||||
.expect("top-level task was abandoned by a concurrent shutdown reclaim");
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_shutdown_stops_and_the_runtime_rebuilds_lazily_after() {
|
||||
let _lock = lock_runtime_state_for_test();
|
||||
|
||||
// No runtime created yet in this process: shutdown must be a no-op,
|
||||
// not a null-pointer dereference.
|
||||
shutdown(Duration::from_secs(1));
|
||||
|
||||
// Force the runtime into existence, then shut it down. Bounded by
|
||||
// the timeout, so a hang here means shutdown itself is broken, not
|
||||
// that the test is slow.
|
||||
assert_eq!(block_on(async { 1 + 1 }), 2);
|
||||
shutdown(Duration::from_secs(5));
|
||||
|
||||
// A caller after shutdown -- e.g. a stray call racing with the
|
||||
// atexit callback -- must get a fresh, working runtime rather than
|
||||
// a dangling reference into the one just torn down.
|
||||
assert_eq!(block_on(async { 2 + 2 }), 4);
|
||||
|
||||
// Shutting down twice in a row (e.g. atexit firing more than once)
|
||||
// must not panic or double-free.
|
||||
shutdown(Duration::from_secs(5));
|
||||
shutdown(Duration::from_secs(5));
|
||||
}
|
||||
|
||||
// Adapted from the concurrent reproducer that found the original bug:
|
||||
// many threads hammering the runtime while shutdown races them, not
|
||||
// just the sequential rebuild path above. Before every reader held its
|
||||
// own `Arc`, this could dereference a runtime `shutdown` had already
|
||||
// freed, or hang forever. Repeated, since a race is not guaranteed to
|
||||
// show up on any single attempt.
|
||||
#[test]
|
||||
fn test_shutdown_is_safe_concurrently_with_live_callers() {
|
||||
use std::sync::Barrier;
|
||||
use std::sync::atomic::AtomicBool as StopFlag;
|
||||
let _lock = lock_runtime_state_for_test();
|
||||
|
||||
for _ in 0..50 {
|
||||
let barrier = Arc::new(Barrier::new(9));
|
||||
let stop = Arc::new(StopFlag::new(false));
|
||||
let workers: Vec<_> = (0..8)
|
||||
.map(|_| {
|
||||
let barrier = Arc::clone(&barrier);
|
||||
let stop = Arc::clone(&stop);
|
||||
std::thread::spawn(move || {
|
||||
barrier.wait();
|
||||
while !stop.load(Ordering::Relaxed) {
|
||||
assert_eq!(block_on(async { 1 + 1 }), 2);
|
||||
}
|
||||
})
|
||||
})
|
||||
.collect();
|
||||
|
||||
barrier.wait();
|
||||
shutdown(Duration::from_millis(50));
|
||||
stop.store(true, Ordering::Relaxed);
|
||||
for worker in workers {
|
||||
worker.join().unwrap();
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Reproduces the actual bug mechanism, not just concurrent `block_on`
|
||||
// traffic: `future_into_py` spawns an outer task that, once running,
|
||||
// spawns a second (inner) one for the real work and awaits its
|
||||
// `JoinHandle` (see `pyo3_async_runtimes::generic::future_into_py_with_locals`).
|
||||
// Before `spawn`/`spawn_blocking` pinned nested calls to whatever
|
||||
// runtime is already executing them, a reclaim landing between the
|
||||
// outer and inner spawn could bind them to two different runtime
|
||||
// instances -- and if the outer task's own runtime was the one torn
|
||||
// down while it awaited the inner task, it never resumed. Every wait
|
||||
// here is bounded so a reintroduced bug fails this test instead of
|
||||
// hanging the suite.
|
||||
#[test]
|
||||
#[allow(unused_must_use)] // fire-and-forget outer spawn, same as future_into_py itself
|
||||
fn test_nested_spawn_survives_concurrent_shutdown() {
|
||||
use std::sync::mpsc;
|
||||
let _lock = lock_runtime_state_for_test();
|
||||
|
||||
for _ in 0..50 {
|
||||
let stop = Arc::new(AtomicBool::new(false));
|
||||
let (done_tx, done_rx) = mpsc::channel::<()>();
|
||||
|
||||
let workers: Vec<_> = (0..8)
|
||||
.map(|_| {
|
||||
let stop = Arc::clone(&stop);
|
||||
let done_tx = done_tx.clone();
|
||||
std::thread::spawn(move || {
|
||||
while !stop.load(Ordering::Relaxed) {
|
||||
let (tx, rx) = mpsc::sync_channel::<()>(1);
|
||||
LanceRuntime::spawn(async move {
|
||||
let inner = LanceRuntime::spawn(async {
|
||||
let _ = 1 + 1;
|
||||
});
|
||||
let _ = inner.await;
|
||||
let _ = tx.send(());
|
||||
});
|
||||
// A hang here across a concurrent shutdown is
|
||||
// exactly the bug this test exists to catch.
|
||||
let _ = rx.recv_timeout(Duration::from_secs(2));
|
||||
}
|
||||
let _ = done_tx.send(());
|
||||
})
|
||||
})
|
||||
.collect();
|
||||
|
||||
std::thread::sleep(Duration::from_millis(10));
|
||||
shutdown(Duration::from_millis(50));
|
||||
stop.store(true, Ordering::Relaxed);
|
||||
|
||||
for _ in 0..8 {
|
||||
done_rx
|
||||
.recv_timeout(Duration::from_secs(10))
|
||||
.expect("worker hung after shutdown raced a nested spawn");
|
||||
}
|
||||
for worker in workers {
|
||||
worker.join().unwrap();
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Forces the exact interleaving the install-race finding described: an
|
||||
// installer registers as outstanding (as `spawn`'s outermost branch
|
||||
// does, before it ever calls `get_runtime()`) and then pauses, while a
|
||||
// concurrent `shutdown()` must not decide "nothing here" and return
|
||||
// before that install actually completes and is retired in turn.
|
||||
#[test]
|
||||
fn test_shutdown_waits_for_a_racing_install() {
|
||||
use std::sync::mpsc;
|
||||
let _lock = lock_runtime_state_for_test();
|
||||
|
||||
// Clean slate: no runtime installed, OUTSTANDING at zero.
|
||||
shutdown(Duration::from_secs(5));
|
||||
|
||||
let (installer_ready_tx, installer_ready_rx) = mpsc::channel::<()>();
|
||||
let (proceed_tx, proceed_rx) = mpsc::channel::<()>();
|
||||
|
||||
let installer = std::thread::spawn(move || {
|
||||
let _guard = OutstandingGuard::new();
|
||||
installer_ready_tx.send(()).unwrap();
|
||||
proceed_rx.recv().unwrap();
|
||||
assert_eq!(get_runtime().block_on(async { 1 + 1 }), 2);
|
||||
});
|
||||
|
||||
installer_ready_rx.recv().unwrap();
|
||||
let shutdown_thread = std::thread::spawn(|| shutdown(Duration::from_secs(5)));
|
||||
// Give shutdown's polling loop several chances to (wrongly) observe
|
||||
// an idle runtime before the installer is allowed to proceed.
|
||||
std::thread::sleep(Duration::from_millis(50));
|
||||
proceed_tx.send(()).unwrap();
|
||||
|
||||
installer.join().unwrap();
|
||||
shutdown_thread.join().unwrap();
|
||||
|
||||
// shutdown must have waited for the install to finish and then
|
||||
// retired it, not returned early and left it stranded.
|
||||
assert!(RUNTIME.load_full().is_none());
|
||||
assert_eq!(OUTSTANDING.load(Ordering::SeqCst), 0);
|
||||
}
|
||||
}
|
||||
|
||||
@@ -441,6 +441,117 @@ impl From<lancedb::table::RefreshColumnResult> for RefreshColumnResult {
|
||||
}
|
||||
}
|
||||
|
||||
/// One row a Function refresh skipped, as the server recorded it.
|
||||
#[pyclass(module = "lancedb._lancedb", get_all, from_py_object)]
|
||||
#[derive(Clone, Debug)]
|
||||
pub struct FunctionErrorRecord {
|
||||
pub job_id: String,
|
||||
pub fragment_id: u64,
|
||||
pub row_offset: Option<u32>,
|
||||
pub column: String,
|
||||
pub function: String,
|
||||
pub function_version: String,
|
||||
pub table_version: u64,
|
||||
pub error_type: String,
|
||||
pub error_message: String,
|
||||
pub created_at_millis: i64,
|
||||
}
|
||||
|
||||
#[pymethods]
|
||||
impl FunctionErrorRecord {
|
||||
pub fn __repr__(&self) -> String {
|
||||
format!(
|
||||
"FunctionErrorRecord(job_id={:?}, fragment_id={}, row_offset={:?}, column={:?}, \
|
||||
error_type={:?}, error_message={:?})",
|
||||
self.job_id,
|
||||
self.fragment_id,
|
||||
self.row_offset,
|
||||
self.column,
|
||||
self.error_type,
|
||||
self.error_message
|
||||
)
|
||||
}
|
||||
}
|
||||
|
||||
impl From<lancedb::function::FunctionErrorRecord> for FunctionErrorRecord {
|
||||
fn from(record: lancedb::function::FunctionErrorRecord) -> Self {
|
||||
Self {
|
||||
job_id: record.job_id,
|
||||
fragment_id: record.fragment_id,
|
||||
row_offset: record.row_offset,
|
||||
column: record.column,
|
||||
function: record.function,
|
||||
function_version: record.function_version,
|
||||
table_version: record.table_version,
|
||||
error_type: record.error_type,
|
||||
error_message: record.error_message,
|
||||
created_at_millis: record.created_at_millis,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// A fragment whose per-row error detail was capped.
|
||||
#[pyclass(module = "lancedb._lancedb", get_all, from_py_object)]
|
||||
#[derive(Clone, Debug)]
|
||||
pub struct FunctionErrorFragment {
|
||||
pub job_id: String,
|
||||
pub fragment_id: u64,
|
||||
pub rows_skipped: u64,
|
||||
pub rows_recorded: u64,
|
||||
}
|
||||
|
||||
#[pymethods]
|
||||
impl FunctionErrorFragment {
|
||||
pub fn __repr__(&self) -> String {
|
||||
format!(
|
||||
"FunctionErrorFragment(job_id={:?}, fragment_id={}, rows_skipped={}, rows_recorded={})",
|
||||
self.job_id, self.fragment_id, self.rows_skipped, self.rows_recorded
|
||||
)
|
||||
}
|
||||
}
|
||||
|
||||
impl From<lancedb::function::FunctionErrorFragment> for FunctionErrorFragment {
|
||||
fn from(fragment: lancedb::function::FunctionErrorFragment) -> Self {
|
||||
Self {
|
||||
job_id: fragment.job_id,
|
||||
fragment_id: fragment.fragment_id,
|
||||
rows_skipped: fragment.rows_skipped,
|
||||
rows_recorded: fragment.rows_recorded,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// A table's per-row Function errors.
|
||||
#[pyclass(module = "lancedb._lancedb", get_all, from_py_object)]
|
||||
#[derive(Clone, Debug)]
|
||||
pub struct FunctionErrors {
|
||||
pub records: Vec<FunctionErrorRecord>,
|
||||
pub fragments: Vec<FunctionErrorFragment>,
|
||||
pub truncated: bool,
|
||||
}
|
||||
|
||||
#[pymethods]
|
||||
impl FunctionErrors {
|
||||
pub fn __repr__(&self) -> String {
|
||||
format!(
|
||||
"FunctionErrors(records={}, fragments={}, truncated={})",
|
||||
self.records.len(),
|
||||
self.fragments.len(),
|
||||
self.truncated
|
||||
)
|
||||
}
|
||||
}
|
||||
|
||||
impl From<lancedb::function::FunctionErrors> for FunctionErrors {
|
||||
fn from(errors: lancedb::function::FunctionErrors) -> Self {
|
||||
Self {
|
||||
records: errors.records.into_iter().map(Into::into).collect(),
|
||||
fragments: errors.fragments.into_iter().map(Into::into).collect(),
|
||||
truncated: errors.truncated,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[pyclass(get_all, from_py_object)]
|
||||
#[derive(Clone, Debug)]
|
||||
pub struct RefreshMaterializedViewResult {
|
||||
@@ -452,6 +563,17 @@ pub struct RefreshMaterializedViewResult {
|
||||
|
||||
#[pymethods]
|
||||
impl RefreshMaterializedViewResult {
|
||||
#[staticmethod]
|
||||
pub fn from_json(value: &str) -> PyResult<Self> {
|
||||
let result: lancedb::RefreshMaterializedViewResult =
|
||||
serde_json::from_str(value).map_err(|err| {
|
||||
PyValueError::new_err(format!(
|
||||
"failed to decode materialized-view refresh result: {err}"
|
||||
))
|
||||
})?;
|
||||
Ok(Self::from(result))
|
||||
}
|
||||
|
||||
pub fn __repr__(&self) -> String {
|
||||
format!(
|
||||
"RefreshMaterializedViewResult(mode={}, rows_written={}, source_version={}, version={})",
|
||||
@@ -1627,6 +1749,25 @@ impl Table {
|
||||
})
|
||||
}
|
||||
|
||||
#[pyo3(signature = (job_id=None, column=None, limit=None))]
|
||||
pub fn function_errors(
|
||||
self_: PyRef<'_, Self>,
|
||||
job_id: Option<String>,
|
||||
column: Option<String>,
|
||||
limit: Option<usize>,
|
||||
) -> PyResult<Bound<'_, PyAny>> {
|
||||
let inner = self_.inner_ref()?.clone();
|
||||
let request = lancedb::function::FunctionErrorsRequest {
|
||||
job_id,
|
||||
column,
|
||||
limit,
|
||||
};
|
||||
future_into_py(self_.py(), async move {
|
||||
let errors = inner.function_errors(request).await.infer_error()?;
|
||||
Ok(FunctionErrors::from(errors))
|
||||
})
|
||||
}
|
||||
|
||||
#[pyo3(signature = (full=false, source_version=None))]
|
||||
pub fn refresh_materialized_view(
|
||||
self_: PyRef<'_, Self>,
|
||||
@@ -1647,6 +1788,40 @@ impl Table {
|
||||
})
|
||||
}
|
||||
|
||||
#[pyo3(signature = (full=false, source_version=None))]
|
||||
pub fn refresh_materialized_view_async(
|
||||
self_: PyRef<'_, Self>,
|
||||
full: bool,
|
||||
source_version: Option<u64>,
|
||||
) -> PyResult<Bound<'_, PyAny>> {
|
||||
let inner = self_.inner_ref()?.clone();
|
||||
future_into_py(self_.py(), async move {
|
||||
let view = lancedb::MaterializedView::from_table(inner)
|
||||
.await
|
||||
.infer_error()?;
|
||||
let mut builder = view.refresh().full(full);
|
||||
if let Some(version) = source_version {
|
||||
builder = builder.source_version(version);
|
||||
}
|
||||
let job = builder.execute_async().await.infer_error()?;
|
||||
Ok(crate::job::Job::new_typed(job))
|
||||
})
|
||||
}
|
||||
|
||||
pub fn materialized_view_definition(self_: PyRef<'_, Self>) -> PyResult<Bound<'_, PyAny>> {
|
||||
let inner = self_.inner_ref()?.clone();
|
||||
future_into_py(self_.py(), async move {
|
||||
let view = lancedb::MaterializedView::from_table(inner)
|
||||
.await
|
||||
.infer_error()?;
|
||||
serde_json::to_string(view.definition()).map_err(|err| {
|
||||
PyRuntimeError::new_err(format!(
|
||||
"failed to serialize materialized-view definition: {err}"
|
||||
))
|
||||
})
|
||||
})
|
||||
}
|
||||
|
||||
pub fn add_columns_with_schema(
|
||||
self_: PyRef<'_, Self>,
|
||||
schema: PyArrowType<Schema>,
|
||||
|
||||
@@ -1,10 +1,19 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright The LanceDB Authors
|
||||
|
||||
import asyncio
|
||||
import importlib.util
|
||||
import json
|
||||
import os
|
||||
import subprocess
|
||||
import sys
|
||||
import threading
|
||||
import urllib.parse
|
||||
from http.server import BaseHTTPRequestHandler, HTTPServer
|
||||
from pathlib import Path
|
||||
|
||||
import pytest
|
||||
|
||||
|
||||
def _load_oauth_module():
|
||||
oauth_path = (
|
||||
@@ -31,3 +40,341 @@ def test_oauth_config_repr_redacts_client_secret():
|
||||
rendered = repr(config)
|
||||
assert "super-secret" not in rendered
|
||||
assert "client_secret" not in rendered
|
||||
|
||||
|
||||
def test_authorization_code_uses_pkce_by_default():
|
||||
oauth = _load_oauth_module()
|
||||
|
||||
config = oauth.OAuthConfig(
|
||||
issuer_url="https://issuer.example.com",
|
||||
client_id="client-id",
|
||||
scopes=["openid"],
|
||||
flow=oauth.OAuthFlowType.AUTHORIZATION_CODE,
|
||||
)
|
||||
|
||||
assert config.use_pkce is True
|
||||
assert config.redirect_uri is None
|
||||
assert config.callback_port is None
|
||||
|
||||
|
||||
def test_device_code_flow_value():
|
||||
oauth = _load_oauth_module()
|
||||
|
||||
assert oauth.OAuthFlowType.DEVICE_CODE.value == "device_code"
|
||||
|
||||
|
||||
def test_client_auth_method_values():
|
||||
oauth = _load_oauth_module()
|
||||
|
||||
assert oauth.ClientAuthMethod.NONE.value == "none"
|
||||
assert oauth.ClientAuthMethod.CLIENT_SECRET_BASIC.value == "client_secret_basic"
|
||||
assert oauth.ClientAuthMethod.CLIENT_SECRET_POST.value == "client_secret_post"
|
||||
|
||||
|
||||
def test_client_auth_method_defaults_to_none():
|
||||
oauth = _load_oauth_module()
|
||||
|
||||
config = oauth.OAuthConfig(
|
||||
issuer_url="https://issuer.example.com",
|
||||
client_id="client-id",
|
||||
scopes=["openid"],
|
||||
client_auth_method=oauth.ClientAuthMethod.CLIENT_SECRET_POST,
|
||||
)
|
||||
|
||||
assert config.client_auth_method is oauth.ClientAuthMethod.CLIENT_SECRET_POST
|
||||
|
||||
default_config = oauth.OAuthConfig(
|
||||
issuer_url="https://issuer.example.com",
|
||||
client_id="client-id",
|
||||
scopes=["openid"],
|
||||
)
|
||||
assert default_config.client_auth_method is None
|
||||
|
||||
|
||||
def test_token_cache_options_default_to_memory_only():
|
||||
oauth = _load_oauth_module()
|
||||
|
||||
config = oauth.OAuthConfig(
|
||||
issuer_url="https://issuer.example.com",
|
||||
client_id="client-id",
|
||||
scopes=["openid"],
|
||||
)
|
||||
assert config.token_cache is None
|
||||
assert config.resource is None
|
||||
assert config.audience is None
|
||||
|
||||
options = oauth.TokenCacheOptions()
|
||||
assert options.cache_dir is None
|
||||
assert options.lock_timeout_secs is None
|
||||
|
||||
|
||||
def _remote_oauth():
|
||||
pytest.importorskip("lancedb")
|
||||
from lancedb.remote import oauth as remote_oauth
|
||||
|
||||
return remote_oauth
|
||||
|
||||
|
||||
def _device_config(remote_oauth, issuer_url, cache_dir):
|
||||
return remote_oauth.OAuthConfig(
|
||||
issuer_url=issuer_url,
|
||||
client_id="client-id",
|
||||
scopes=["openid"],
|
||||
flow=remote_oauth.OAuthFlowType.DEVICE_CODE,
|
||||
token_cache=remote_oauth.TokenCacheOptions(cache_dir=str(cache_dir)),
|
||||
)
|
||||
|
||||
|
||||
def test_oauth_session_status_and_logout_without_cache_entry(tmp_path):
|
||||
remote_oauth = _remote_oauth()
|
||||
config = _device_config(remote_oauth, "https://issuer.example.com", tmp_path)
|
||||
|
||||
session = remote_oauth.OAuthSession(config)
|
||||
status = asyncio.run(session.status())
|
||||
assert status.refreshable is False
|
||||
assert status.issuer_url == "https://issuer.example.com"
|
||||
assert status.client_id == "client-id"
|
||||
assert status.scopes == ["openid"]
|
||||
assert status.flow == "device_code"
|
||||
assert status.obtained_at is None
|
||||
|
||||
logout = asyncio.run(session.logout())
|
||||
assert logout.removed is False
|
||||
|
||||
|
||||
class _MockIdpState:
|
||||
def __init__(self, port):
|
||||
self.port = port
|
||||
self.lock = threading.Lock()
|
||||
self.device_authorizations = 0
|
||||
self.refresh_grants = 0
|
||||
self.invalid_grant_rejections = 0
|
||||
self.access_tokens_issued = 0
|
||||
self.current_refresh = None
|
||||
self.requests = []
|
||||
|
||||
|
||||
class _MockIdpHandler(BaseHTTPRequestHandler):
|
||||
@property
|
||||
def state(self) -> _MockIdpState:
|
||||
return self.server.state
|
||||
|
||||
def log_message(self, fmt, *args):
|
||||
pass
|
||||
|
||||
def _respond(self, status, payload):
|
||||
body = json.dumps(payload).encode()
|
||||
self.send_response(status)
|
||||
self.send_header("Content-Type", "application/json")
|
||||
self.send_header("Content-Length", str(len(body)))
|
||||
self.end_headers()
|
||||
self.wfile.write(body)
|
||||
|
||||
def do_GET(self):
|
||||
if self.path.startswith("/.well-known/openid-configuration"):
|
||||
base = f"http://127.0.0.1:{self.state.port}"
|
||||
self._respond(
|
||||
200,
|
||||
{
|
||||
"token_endpoint": f"{base}/token",
|
||||
"device_authorization_endpoint": f"{base}/device",
|
||||
},
|
||||
)
|
||||
else:
|
||||
self._respond(404, {})
|
||||
|
||||
def do_POST(self):
|
||||
length = int(self.headers.get("Content-Length", 0))
|
||||
body = self.rfile.read(length).decode()
|
||||
params = urllib.parse.parse_qs(body)
|
||||
with self.state.lock:
|
||||
self.state.requests.append(params)
|
||||
|
||||
if self.path == "/device":
|
||||
with self.state.lock:
|
||||
self.state.device_authorizations += 1
|
||||
base = f"http://127.0.0.1:{self.state.port}"
|
||||
self._respond(
|
||||
200,
|
||||
{
|
||||
"device_code": "device-code",
|
||||
"user_code": "ABCD-EFGH",
|
||||
"verification_uri": f"{base}/verify",
|
||||
"expires_in": 60,
|
||||
"interval": 1,
|
||||
},
|
||||
)
|
||||
return
|
||||
|
||||
if self.path == "/token":
|
||||
grant_type = params.get("grant_type", [""])[0]
|
||||
with self.state.lock:
|
||||
if grant_type == "refresh_token":
|
||||
self.state.refresh_grants += 1
|
||||
offered = params.get("refresh_token", [""])[0]
|
||||
if offered != self.state.current_refresh:
|
||||
self.state.invalid_grant_rejections += 1
|
||||
self._respond(400, {"error": "invalid_grant"})
|
||||
return
|
||||
elif "device_code" not in grant_type:
|
||||
self._respond(400, {"error": "unsupported_grant_type"})
|
||||
return
|
||||
self.state.access_tokens_issued += 1
|
||||
number = self.state.access_tokens_issued
|
||||
refresh = f"refresh-{number}"
|
||||
self.state.current_refresh = refresh
|
||||
self._respond(
|
||||
200,
|
||||
{
|
||||
"access_token": f"access-{number}",
|
||||
"refresh_token": refresh,
|
||||
"expires_in": 3600,
|
||||
},
|
||||
)
|
||||
return
|
||||
|
||||
self._respond(404, {})
|
||||
|
||||
|
||||
def _start_mock_idp() -> tuple[_MockIdpState, HTTPServer]:
|
||||
server = HTTPServer(("127.0.0.1", 0), _MockIdpHandler)
|
||||
state = _MockIdpState(server.server_address[1])
|
||||
server.state = state
|
||||
thread = threading.Thread(target=server.serve_forever, daemon=True)
|
||||
thread.start()
|
||||
return state, server
|
||||
|
||||
|
||||
def _run_subprocess(script: Path, issuer_url: str, cache_dir: Path, target: dict):
|
||||
env = dict(os.environ)
|
||||
env["LANCEDB_OAUTH_BROWSER"] = "/usr/bin/true"
|
||||
result = subprocess.run(
|
||||
[sys.executable, str(script), issuer_url, str(cache_dir), json.dumps(target)],
|
||||
capture_output=True,
|
||||
text=True,
|
||||
timeout=120,
|
||||
env=env,
|
||||
)
|
||||
assert result.returncode == 0, (
|
||||
f"subprocess failed:\nstdout: {result.stdout}\nstderr: {result.stderr}"
|
||||
)
|
||||
return result
|
||||
|
||||
|
||||
LOGIN_SCRIPT = """
|
||||
import asyncio
|
||||
import json
|
||||
import sys
|
||||
|
||||
from lancedb.remote import OAuthConfig, OAuthFlowType, OAuthSession, TokenCacheOptions
|
||||
|
||||
issuer_url, cache_dir = sys.argv[1], sys.argv[2]
|
||||
config = OAuthConfig(
|
||||
issuer_url=issuer_url,
|
||||
client_id="client-id",
|
||||
scopes=["openid"],
|
||||
flow=OAuthFlowType.DEVICE_CODE,
|
||||
token_cache=TokenCacheOptions(cache_dir=cache_dir),
|
||||
**json.loads(sys.argv[3]),
|
||||
)
|
||||
session = OAuthSession(config)
|
||||
status = asyncio.run(session.login())
|
||||
assert status.refreshable, "login must cache a refresh token"
|
||||
assert status.resource == config.resource
|
||||
assert status.audience == config.audience
|
||||
print("LOGIN-OK")
|
||||
"""
|
||||
|
||||
REUSE_SCRIPT = """
|
||||
import asyncio
|
||||
import json
|
||||
import sys
|
||||
|
||||
import lancedb
|
||||
from lancedb.remote import OAuthConfig, OAuthFlowType, OAuthSession, TokenCacheOptions
|
||||
|
||||
issuer_url, cache_dir = sys.argv[1], sys.argv[2]
|
||||
config = OAuthConfig(
|
||||
issuer_url=issuer_url,
|
||||
client_id="client-id",
|
||||
scopes=["openid"],
|
||||
flow=OAuthFlowType.DEVICE_CODE,
|
||||
token_cache=TokenCacheOptions(cache_dir=cache_dir),
|
||||
**json.loads(sys.argv[3]),
|
||||
)
|
||||
|
||||
session = OAuthSession(config)
|
||||
status = asyncio.run(session.status())
|
||||
assert status.refreshable, "second process must see the cached session"
|
||||
|
||||
|
||||
async def main():
|
||||
# Point the database endpoint at a dead port. OAuth headers are fetched
|
||||
# before the request is sent, so a successful refresh proves the second
|
||||
# process reused the cached session; only the database call fails.
|
||||
db = await lancedb.connect_async(
|
||||
"db://e2e",
|
||||
host_override="http://127.0.0.1:1",
|
||||
client_config={"retry_config": {"retries": 0}},
|
||||
oauth_config=config,
|
||||
)
|
||||
try:
|
||||
await db.table_names()
|
||||
except Exception:
|
||||
print("DATABASE-UNREACHABLE-AS-EXPECTED")
|
||||
else:
|
||||
raise AssertionError("expected the database request to fail")
|
||||
|
||||
|
||||
asyncio.run(main())
|
||||
print("REUSE-OK")
|
||||
"""
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"target",
|
||||
[
|
||||
{},
|
||||
{
|
||||
"resource": "https://api.example.com/a?x=1&y=two",
|
||||
"audience": "audience + & / ü",
|
||||
},
|
||||
],
|
||||
)
|
||||
def test_cross_process_session_reuse_without_new_prompt(tmp_path, target):
|
||||
pytest.importorskip("lancedb")
|
||||
state, server = _start_mock_idp()
|
||||
try:
|
||||
issuer_url = f"http://127.0.0.1:{state.port}"
|
||||
login_script = tmp_path / "login.py"
|
||||
login_script.write_text(LOGIN_SCRIPT)
|
||||
reuse_script = tmp_path / "reuse.py"
|
||||
reuse_script.write_text(REUSE_SCRIPT)
|
||||
cache_dir = tmp_path / "oauth-cache"
|
||||
|
||||
result = _run_subprocess(login_script, issuer_url, cache_dir, target)
|
||||
assert "LOGIN-OK" in result.stdout
|
||||
assert state.device_authorizations == 1
|
||||
|
||||
result = _run_subprocess(reuse_script, issuer_url, cache_dir, target)
|
||||
assert "REUSE-OK" in result.stdout
|
||||
assert "DATABASE-UNREACHABLE-AS-EXPECTED" in result.stdout
|
||||
|
||||
# The second process refreshed exactly once and never started a new
|
||||
# interactive device flow.
|
||||
assert state.refresh_grants == 1
|
||||
assert state.device_authorizations == 1
|
||||
assert state.invalid_grant_rejections == 0
|
||||
|
||||
assert len(state.requests) == 3
|
||||
for params in state.requests:
|
||||
for key in ("resource", "audience"):
|
||||
assert params.get(key) == ([target[key]] if key in target else None)
|
||||
config = _device_config(_remote_oauth(), issuer_url, cache_dir)
|
||||
config.resource = target.get("resource")
|
||||
config.audience = target.get("audience")
|
||||
logout = asyncio.run(_remote_oauth().OAuthSession(config).logout())
|
||||
assert logout.removed is True
|
||||
finally:
|
||||
server.shutdown()
|
||||
server.server_close()
|
||||
|
||||
Generated
+8
-8
@@ -2005,7 +2005,7 @@ requires-dist = [
|
||||
{ name = "pyarrow-stubs", marker = "extra == 'tests'", specifier = ">=16.0" },
|
||||
{ name = "pydantic", specifier = ">=2.7.4,<3" },
|
||||
{ name = "pylance", marker = "extra == 'pylance'", specifier = ">=5.0.0b5" },
|
||||
{ name = "pylance", marker = "extra == 'tests'", specifier = "==9.0.0rc1" },
|
||||
{ name = "pylance", marker = "extra == 'tests'", specifier = "==9.0.0" },
|
||||
{ name = "pyright", marker = "extra == 'dev'", specifier = ">=1.1.350" },
|
||||
{ name = "pytest", marker = "extra == 'tests'", specifier = ">=7.0" },
|
||||
{ name = "pytest-asyncio", marker = "extra == 'tests'", specifier = ">=0.21" },
|
||||
@@ -3912,7 +3912,7 @@ crypto = [
|
||||
|
||||
[[package]]
|
||||
name = "pylance"
|
||||
version = "7.0.0"
|
||||
version = "9.0.0"
|
||||
source = { registry = "https://pypi.org/simple" }
|
||||
dependencies = [
|
||||
{ name = "lance-namespace" },
|
||||
@@ -3922,12 +3922,12 @@ dependencies = [
|
||||
{ name = "pyarrow" },
|
||||
]
|
||||
wheels = [
|
||||
{ url = "https://files.pythonhosted.org/packages/ac/ad/2f64921bf346e7075aef24a72595db44821724a3d89a9a92dd24e79632aa/pylance-7.0.0-cp39-abi3-macosx_11_0_arm64.whl", hash = "sha256:98422021975be76e72b1572f41b8c9abb3bee5bdc9bfa5e9ce731110a65ed4d1", size = 62134146, upload-time = "2026-05-27T21:59:37.459Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/73/1c/c5a01bee0160b55d9a98895cbd33091d038f0a0995b121ab72e629008d02/pylance-7.0.0-cp39-abi3-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:4bec86ee5b6fbd8bfc493e653f0a1fba0303cfe5492b9b46fc25ab908edc7183", size = 65373684, upload-time = "2026-05-27T22:04:01.584Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/eb/da/1fe8b8f7dbfe734d76af76acc994fc360a0d0c79a4874ef69f5a72a58fe3/pylance-7.0.0-cp39-abi3-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:881491432c53184e52f8d1db8d5f872f39a03f36fb104bec77b33d379519d8b5", size = 69458555, upload-time = "2026-05-27T22:16:50.567Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/76/f0/dd505cf3fd0226ab9d94759acd713125af1d3bfacfd80bbd52e3b9f89509/pylance-7.0.0-cp39-abi3-manylinux_2_28_aarch64.whl", hash = "sha256:18453999e7fff4f76b16d6b7882c9df0628bd142ff95e2461bd7dd5ee3fe0af3", size = 65394430, upload-time = "2026-05-27T22:05:30.923Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/17/ba/2357b81034f28eb00790e258ed140289a6a887a7468ca9df6349fd186b27/pylance-7.0.0-cp39-abi3-manylinux_2_28_x86_64.whl", hash = "sha256:04a58051d408c60fe76d41a220dcaf8fea8fb6d1aa0ca78a709b60bc3cc8d19a", size = 69473470, upload-time = "2026-05-27T22:17:18.935Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/1f/ec/5c00b6303a67d787f9475141832cbdc513d674ac3dcaeef8a7b169905e65/pylance-7.0.0-cp39-abi3-win_amd64.whl", hash = "sha256:467d4864af047eaab4e1370e2f1e88e2c6f507c079874421116cb41d78bc3629", size = 74792863, upload-time = "2026-05-27T22:19:23.875Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/e9/be/45733acd64801991852aac8e658601fd8fc12f76ceb81e57fca690896b90/pylance-9.0.0-cp310-abi3-macosx_11_0_arm64.whl", hash = "sha256:8257213501d3298c5b6a344d60938e4bbe4de9f00cd3265371a56d1dc3dd15ca", size = 68377982, upload-time = "2026-07-24T16:53:45.247Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/d0/3e/1ef707cb215cc7268c63ad84a91344ad6313d3984b343eeb19a9b708698e/pylance-9.0.0-cp310-abi3-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:804eedfa1fda2e703cca8580c76f0b44a1b849b725a8908c06f8acfca811732f", size = 71844362, upload-time = "2026-07-24T16:56:07.6Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/c8/fb/a499e5c53ddb75c7de44100fd2bb1f7cc735966200d20292eed7af9ef552/pylance-9.0.0-cp310-abi3-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:0a0b75595e3766c1d5f4c90abdc52337b4de60f1a52d123a4f8c4e5bcbbbfa8f", size = 75663088, upload-time = "2026-07-24T17:10:31.283Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/e9/80/0714e09f64a68dbdf62558955e737a7436df763b9834a6aba506861b5352/pylance-9.0.0-cp310-abi3-manylinux_2_28_aarch64.whl", hash = "sha256:d2f69c5c390ae3710c35a429905fe15f769951777fafb1b72a359e035ec121f7", size = 71866858, upload-time = "2026-07-24T16:56:35.937Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/4b/3c/78d3a6d6ca0d843b7c3ac0c30d9cd2cf4635b0c2cabe6ec66583d5bfc1a1/pylance-9.0.0-cp310-abi3-manylinux_2_28_x86_64.whl", hash = "sha256:836268a7832d62f3d5ccbe1c3fca239621971d90297bcfff14a70b3cb6842aa8", size = 75642656, upload-time = "2026-07-24T17:13:11.879Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/cf/c1/dc9c9a31e171530ec0add024d922488c046437d32a7087c29e254a7eacc7/pylance-9.0.0-cp310-abi3-win_amd64.whl", hash = "sha256:96441d27a5ed3805300388ccf8f31835cd280c98751f80b6a1a11dcd6808fc43", size = 81668288, upload-time = "2026-07-24T17:05:38.707Z" },
|
||||
]
|
||||
|
||||
[[package]]
|
||||
|
||||
Reference in New Issue
Block a user