Merge remote-tracking branch 'origin/main' into feature/wal-hybrid-pk

This commit is contained in:
Daniel Rammer
2026-09-23 12:29:50 -05:00
175 changed files with 32431 additions and 2151 deletions
-1
View File
@@ -1 +0,0 @@
AGENTS.md
+2 -1
View File
@@ -1,6 +1,6 @@
[package]
name = "lancedb-python"
version = "0.39.0-beta.6"
version = "0.40.0-beta.8"
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
+2 -1
View File
@@ -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",
+20
View File
@@ -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,11 +40,14 @@ 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,
MaterializedViewDefinition,
)
from .view import ViewDescription as ViewDescription
from .table import AsyncTable, Table
from .types import BaseTokenizerType
from ._lancedb import Session
@@ -52,6 +58,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,9 +575,15 @@ async def connect_async(
__all__ = [
"Catalog",
"AsyncCatalog",
"ListDatabasesResponse",
"connect_catalog",
"connect_catalog_async",
"AsyncMaterializedView",
"MaterializedView",
"MaterializedViewDefinition",
"ViewDescription",
"connect",
"connect_async",
"tokenize",
+136
View File
@@ -153,6 +153,36 @@ 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 drop_function_async(
self, name: str, version: str
) -> Tuple[bool, Job]: ...
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 create_view(
self, name: str, query: str, namespace_path: Optional[List[str]] = None
) -> Tuple[str, List[str], str, str, List[str], pa.Schema]: ...
async def describe_view(
self, name: str, namespace_path: Optional[List[str]] = None
) -> Tuple[str, List[str], str, str, List[str], pa.Schema]: ...
async def drop_view(
self, name: str, namespace_path: Optional[List[str]] = None
) -> None: ...
async def list_views(
self, namespace_path: Optional[List[str]] = None
) -> List[str]: ...
async def list_jobs(self) -> List[JobInfo]: ...
async def cancel_job(self, job_id: str) -> bool: ...
async def execute_query_async(
@@ -211,8 +241,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 +313,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 +476,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]]
@@ -756,7 +843,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
@@ -805,3 +917,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: ...
+196
View File
@@ -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,
)
+564 -32
View File
@@ -17,6 +17,8 @@ from typing import (
List,
Literal,
Optional,
Sequence,
Tuple,
Union,
)
from uuid import UUID
@@ -57,6 +59,12 @@ from .materialized_view import (
SelectArg,
normalize_select,
)
from .secrets import (
EnvVarSecret,
SecretInfo,
validate_namespace_path,
validate_secret_name,
)
from .table import (
AsyncTable,
LanceTable,
@@ -67,6 +75,7 @@ from .util import (
get_uri_scheme,
validate_table_name,
)
from .view import ViewDescription
import deprecation
@@ -88,6 +97,28 @@ from .namespace_utils import (
)
def _view_description(
described: Tuple[str, List[str], str, str, List[str], "pa.Schema"],
) -> ViewDescription:
"""Name the fields the binding returns positionally."""
(
name,
namespace_path,
query,
default_database,
default_namespace_path,
schema,
) = described
return ViewDescription(
name=name,
query=query,
default_database=default_database,
schema=schema,
namespace_path=namespace_path,
default_namespace_path=default_namespace_path,
)
class DBConnection(EnforceOverrides):
"""An active LanceDB connection interface."""
@@ -529,13 +560,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 +588,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 +599,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 +640,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 +773,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 +852,157 @@ class DBConnection(EnforceOverrides):
)
def drop_function(self, name: str, *, version: str) -> bool:
"""Drop one exact immutable Function version from the remote catalog.
"""Drop a Function name and the object it was bound to.
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. Returns
True when the name was removed and False when it was absent. The
object's content is deleted with it, which may finish after this
returns; use :meth:`drop_function_async` to wait for that. Local
connections raise NotImplementedError.
"""
raise NotImplementedError(
"Function catalog operations are not supported for this connection type"
)
def drop_function_async(self, name: str, *, version: str) -> Tuple[bool, Job]:
"""Drop a Function name and return its cleanup job.
The name is unbound before this returns; the object's content may still
be being deleted. Call :meth:`Job.wait` to wait for that to finish. When
the server deletes inline, or when nothing was bound, the returned job
is already finished and has no id. 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 create_view(
self, name: str, query: str, *, namespace_path: Optional[List[str]] = None
) -> ViewDescription:
"""Create a view: a named query the database plans on every read.
The query is planned once, at creation, so one that cannot be planned
is refused now rather than at the first read. A view holds no rows, and
its readers see its sources as they are at read time.
There is no replace: a name already taken is an error, and changing a
view is a drop followed by a create. Local connections raise
``NotImplementedError``.
>>> import lancedb
>>> db = lancedb.connect("db://my_database") # doctest: +SKIP
>>> view = db.create_view(
... "adults", "SELECT name FROM people WHERE age >= 18"
... ) # doctest: +SKIP
>>> view.schema # doctest: +SKIP
name: string
"""
raise NotImplementedError(
"View operations are not supported for this connection type"
)
def describe_view(
self, name: str, *, namespace_path: Optional[List[str]] = None
) -> ViewDescription:
"""What this database records about a view: its defining query and the
schema that query resolved to.
The schema is the one recorded at creation; a source altered since then
shows up when the view is read. Local connections raise
``NotImplementedError``.
"""
raise NotImplementedError(
"View operations are not supported for this connection type"
)
def drop_view(
self, name: str, *, namespace_path: Optional[List[str]] = None
) -> None:
"""Drop a view.
The tables it reads are untouched: a view holds no rows of its own.
Local connections raise ``NotImplementedError``.
"""
raise NotImplementedError(
"View operations are not supported for this connection type"
)
def list_views(self, *, namespace_path: Optional[List[str]] = None) -> List[str]:
"""The names of the views in one namespace.
Names only; a definition comes from
[describe_view][lancedb.db.DBConnection.describe_view]. Local
connections raise ``NotImplementedError``.
"""
raise NotImplementedError(
"View 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 +1528,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 +1549,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 +1599,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 +1762,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 +1783,63 @@ class LanceDBConnection(DBConnection):
def drop_function(self, name: str, *, version: str) -> bool:
return LOOP.run(self._conn.drop_function(name, version=version))
@override
def drop_function_async(self, name: str, *, version: str) -> Tuple[bool, Job]:
dropped, job = LOOP.run(self._conn.drop_function_async(name, version=version))
return dropped, Job(job)
@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 create_view(
self, name: str, query: str, *, namespace_path: Optional[List[str]] = None
) -> ViewDescription:
return LOOP.run(
self._conn.create_view(name, query, namespace_path=namespace_path)
)
@override
def describe_view(
self, name: str, *, namespace_path: Optional[List[str]] = None
) -> ViewDescription:
return LOOP.run(self._conn.describe_view(name, namespace_path=namespace_path))
@override
def drop_view(
self, name: str, *, namespace_path: Optional[List[str]] = None
) -> None:
LOOP.run(self._conn.drop_view(name, namespace_path=namespace_path))
@override
def list_views(self, *, namespace_path: Optional[List[str]] = None) -> List[str]:
return LOOP.run(self._conn.list_views(namespace_path=namespace_path))
@override
def list_jobs(self) -> List[JobInfo]:
"""List server-side jobs across the database's tables."""
@@ -2080,6 +2447,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 +2459,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 +2505,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 +2689,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 +2725,111 @@ class AsyncConnection(object):
]
async def drop_function(self, name: str, *, version: str) -> bool:
"""Drop one exact immutable Function version from the remote catalog."""
"""Drop a Function name and the object it was bound to."""
return await self._inner.drop_function(name, version)
async def drop_function_async(
self, name: str, *, version: str
) -> Tuple[bool, AsyncJob]:
"""Drop a Function name and return its cleanup job.
The name is unbound before this returns; the object's content may still
be being deleted. Await :meth:`AsyncJob.wait` to wait for that to
finish.
"""
dropped, job = await self._inner.drop_function_async(name, version)
return dropped, AsyncJob(job)
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 create_view(
self, name: str, query: str, *, namespace_path: Optional[List[str]] = None
) -> ViewDescription:
"""Create a view: a named query the database plans on every read.
See
[DBConnection.create_view][lancedb.DBConnection.create_view].
"""
return _view_description(
await self._inner.create_view(name, query, list(namespace_path or []))
)
async def describe_view(
self, name: str, *, namespace_path: Optional[List[str]] = None
) -> ViewDescription:
"""What this database records about a view: query and schema."""
return _view_description(
await self._inner.describe_view(name, list(namespace_path or []))
)
async def drop_view(
self, name: str, *, namespace_path: Optional[List[str]] = None
) -> None:
"""Drop a view. The tables it reads are untouched."""
await self._inner.drop_view(name, list(namespace_path or []))
async def list_views(
self, *, namespace_path: Optional[List[str]] = None
) -> List[str]:
"""The names of the views in one namespace."""
return await self._inner.list_views(list(namespace_path or []))
async def list_jobs(self) -> List[JobInfo]:
"""List server-side jobs across the database's tables."""
return await self._inner.list_jobs()
@@ -7,6 +7,7 @@ import io
import os
import urllib.parse as urlparse
from typing import TYPE_CHECKING, List, Union
from urllib.request import url2pathname
import numpy as np
import pyarrow as pa
@@ -154,7 +155,7 @@ class OpenClipEmbeddings(EmbeddingFunction):
parsed = urlparse.urlparse(image)
# TODO handle drive letter on windows.
if parsed.scheme == "file":
return PIL_Image.open(parsed.path)
return PIL_Image.open(url2pathname(parsed.path))
elif parsed.scheme == "":
return PIL_Image.open(image if os.name == "nt" else parsed.path)
elif parsed.scheme.startswith("http"):
+1 -1
View File
@@ -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.
File diff suppressed because it is too large Load Diff
+108 -31
View File
@@ -7,10 +7,11 @@ maintained by refresh. See ``DBConnection.create_materialized_view``."""
from __future__ import annotations
import json
from dataclasses import dataclass, field
from dataclasses import dataclass
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
@@ -28,22 +29,25 @@ SelectArg = Union[
]
DEFINITION_FORMAT = 2
"""The newest stored layout this version reads: ``{"format": N, "query": "<SQL>"}``,
format 2 being a query with ``GROUP BY``. A ``kind`` key beside it is for
readers older than the format number."""
@dataclass
class MaterializedViewDefinition:
"""The query that defines a materialized view."""
"""The query that defines a materialized view, as stored::
source_table: str
"""Name of the source table, in the same database as the view."""
projections: List[Tuple[str, str]]
"""``(output column, SQL expression)`` pairs, in view schema order."""
filter: Optional[str] = None
"""SQL predicate selecting the source rows the view holds."""
limit: Optional[int] = None
"""Cap on the number of rows the view holds."""
inputs: List[str] = field(default_factory=list)
"""Source columns the projections and filter read."""
source_namespace: List[str] = field(default_factory=list)
"""Namespace holding the source table; empty is the root namespace."""
SELECT columns
FROM [ns.]table [, function(args) AS alias | , UNNEST(column) AS alias]
[WHERE predicate] [GROUP BY expr, ...] [LIMIT n]
A Function in ``FROM`` position yields one row per element it returns.
"""
query: str
"""The defining query, in the canonical spelling the server stores."""
def _definition_from_schema(
@@ -53,24 +57,64 @@ def _definition_from_schema(
raw = metadata.get(DEFINITION_META_KEY)
if raw is None:
raise ValueError(f"Table '{name}' is not a materialized view")
value = json.loads(raw)
return _definition_from_value(json.loads(raw), name)
def _definition_from_json(raw: str, name: str = "") -> MaterializedViewDefinition:
"""Parse the definition native code hands over, in its stored layout."""
return _definition_from_value(json.loads(raw), name)
def _definition_from_value(value: dict, name: str) -> MaterializedViewDefinition:
fmt = value.get("format")
if fmt is not None:
# A newer writer's layout is reported, never guessed at.
if not isinstance(fmt, int) or fmt > DEFINITION_FORMAT:
raise NotImplementedError(
f"materialized view '{name}' is stored in format {fmt}, which "
"this version of lancedb cannot refresh"
)
return MaterializedViewDefinition(query=value["query"])
# The structured layout written before the format number.
kind = value.get("kind")
# "namespaced_select" keeps older readers from resolving the source at root.
if kind not in ("select", "namespaced_select"):
raise NotImplementedError(
f"materialized view '{name}' is defined by '{kind}', which this "
"version of lancedb cannot refresh"
f"materialized view '{name}' is stored in format kind '{kind}', "
"which this version of lancedb cannot refresh"
)
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", []),
return MaterializedViewDefinition(query=_legacy_query(value))
def _legacy_ident(name: str) -> str:
if name and all(c == "_" or c.islower() or c.isdigit() for c in name):
return name
return _quote_identifier(name)
def _legacy_query(value: dict) -> str:
"""Render the pre-format structured layout as the query it described."""
projections = value.get("projections", [])
if projections:
items = []
for p in projections:
output, expression = p["output"], p["expression"]
if expression in (output, _quote_identifier(output)):
items.append(expression)
else:
items.append(f"{expression} AS {_legacy_ident(output)}")
columns = ", ".join(items)
else:
columns = "*"
table = ".".join(
_legacy_ident(part)
for part in [*value.get("source_namespace", []), value["source_table"]]
)
query = f"SELECT {columns} FROM {table}"
if value.get("filter") is not None:
query += f" WHERE {value['filter']}"
if value.get("limit") is not None:
query += f" LIMIT {value['limit']}"
return query
def _quote_identifier(name: str) -> str:
@@ -126,8 +170,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 +193,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 +234,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 +243,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)
)
)
+99 -2
View File
@@ -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:
+14 -2
View File
@@ -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]
+184 -10
View File
@@ -2,12 +2,23 @@
# 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,
Tuple,
Union,
)
from urllib.parse import urlparse
from uuid import UUID
import warnings
@@ -29,6 +40,8 @@ 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
from ..view import ViewDescription
if TYPE_CHECKING:
from .._lancedb import JobInfo
@@ -187,11 +200,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 +718,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 +857,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 +878,63 @@ class RemoteDBConnection(DBConnection):
def drop_function(self, name: str, *, version: str) -> bool:
return LOOP.run(self._conn.drop_function(name, version=version))
@override
def drop_function_async(self, name: str, *, version: str) -> Tuple[bool, Job]:
dropped, job = LOOP.run(self._conn.drop_function_async(name, version=version))
return dropped, Job(job)
@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 create_view(
self, name: str, query: str, *, namespace_path: Optional[List[str]] = None
) -> ViewDescription:
return LOOP.run(
self._conn.create_view(name, query, namespace_path=namespace_path)
)
@override
def describe_view(
self, name: str, *, namespace_path: Optional[List[str]] = None
) -> ViewDescription:
return LOOP.run(self._conn.describe_view(name, namespace_path=namespace_path))
@override
def drop_view(
self, name: str, *, namespace_path: Optional[List[str]] = None
) -> None:
LOOP.run(self._conn.drop_view(name, namespace_path=namespace_path))
@override
def list_views(self, *, namespace_path: Optional[List[str]] = None) -> List[str]:
return LOOP.run(self._conn.list_views(namespace_path=namespace_path))
@override
def list_jobs(self) -> List["JobInfo"]:
"""List server-side jobs across the database's tables."""
+185
View File
@@ -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()
+13
View File
@@ -982,6 +982,9 @@ class RemoteTable(Table):
self,
transforms: Dict[str, str | FunctionApplication]
| FunctionApplication
| pa.Field
| List[pa.Field]
| pa.Schema
| None = None,
*,
computed: Dict[str, str] | None = None,
@@ -994,6 +997,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",
+159
View File
@@ -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
+226
View File
@@ -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",
]
+273 -37
View File
@@ -5,6 +5,7 @@ from __future__ import annotations
import asyncio
import inspect
import json
import deprecation
import warnings
from abc import ABC, abstractmethod
@@ -186,6 +187,7 @@ if TYPE_CHECKING:
CompactionStats,
Tag,
AddColumnsResult,
FunctionErrors,
RefreshColumnResult,
AddResult,
AlterColumnsResult,
@@ -265,12 +267,14 @@ def _into_pyarrow_reader(
# convert to list of dict if data is a bunch of LanceModels
if isinstance(data[0], LanceModel):
schema = data[0].__class__.to_arrow_schema()
model_schema = data[0].__class__.to_arrow_schema()
data = [model_to_dict(d) for d in data]
return pa.Table.from_pylist(data, schema=schema).to_reader()
data = _serialize_json_values(data, schema or model_schema)
return pa.Table.from_pylist(data, schema=model_schema).to_reader()
elif isinstance(data[0], pa.RecordBatch):
return pa.Table.from_batches(data).to_reader()
else:
data = _serialize_json_values(data, schema)
return pa.Table.from_pylist(data).to_reader()
elif _check_for_pandas(data) and isinstance(data, pd.DataFrame):
table = pa.Table.from_pandas(data, preserve_index=False)
@@ -618,6 +622,108 @@ def _field_extension_name(field: pa.Field) -> Optional[str]:
return extension_name
def _is_json_field(field: pa.Field) -> bool:
return _field_extension_name(field) in ("arrow.json", "lance.json")
@dataclass(frozen=True)
class _JsonSerializationPlan:
arrow_field: pa.Field
children: Optional[Dict[str, "_JsonSerializationPlan"]] = None
item: Optional["_JsonSerializationPlan"] = None
def _json_serialization_plan(field: pa.Field) -> Optional[_JsonSerializationPlan]:
if _is_json_field(field):
return _JsonSerializationPlan(field)
if pa.types.is_struct(field.type):
children: Dict[str, _JsonSerializationPlan] = {}
for child_field in field.type:
child_plan = _json_serialization_plan(child_field)
if child_plan is not None:
children[child_field.name] = child_plan
if children:
return _JsonSerializationPlan(field, children=children)
if _is_list_like(field.type):
item_plan = _json_serialization_plan(field.type.value_field)
if item_plan is not None:
return _JsonSerializationPlan(field, item=item_plan)
return None
def _json_serialization_plans(
schema: pa.Schema,
) -> Dict[str, _JsonSerializationPlan]:
plans: Dict[str, _JsonSerializationPlan] = {}
for field in schema:
plan = _json_serialization_plan(field)
if plan is not None:
plans[field.name] = plan
return plans
def _serialize_json_value(value: Any, plan: _JsonSerializationPlan) -> Any:
if value is None or isinstance(value, str):
return value
if _is_json_field(plan.arrow_field):
if isinstance(value, (dict, list)):
return json.dumps(value)
return value
if plan.children is not None and isinstance(value, dict):
serialized = None
for child_name, child_plan in plan.children.items():
if child_name not in value:
continue
child_value = _serialize_json_value(value[child_name], child_plan)
if child_value is not value[child_name]:
if serialized is None:
serialized = dict(value)
serialized[child_name] = child_value
return serialized if serialized is not None else value
if plan.item is not None and isinstance(value, list):
serialized = None
for index, item in enumerate(value):
serialized_item = _serialize_json_value(item, plan.item)
if serialized_item is not item:
if serialized is None:
serialized = list(value)
serialized[index] = serialized_item
return serialized if serialized is not None else value
return value
def _serialize_json_values(data: Any, target_schema: Optional[pa.Schema]) -> Any:
if target_schema is None or not isinstance(data, list):
return data
plans = _json_serialization_plans(target_schema)
if not plans:
return data
serialized_rows = []
for row in data:
if not isinstance(row, dict):
serialized_rows.append(row)
continue
serialized_row = None
for field_name, plan in plans.items():
if field_name not in row:
continue
value = _serialize_json_value(row[field_name], plan)
if value is not row[field_name]:
if serialized_row is None:
serialized_row = dict(row)
serialized_row[field_name] = value
serialized_rows.append(serialized_row if serialized_row is not None else row)
return serialized_rows
def _align_field_types(
fields: List[pa.Field],
target_fields: List[pa.Field],
@@ -634,26 +740,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 +790,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 +799,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 +808,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 +2326,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 +2349,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 +2363,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 +2414,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]]):
"""
@@ -4329,17 +4508,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:
@@ -5659,6 +5851,7 @@ class AsyncTable:
"""
schema = await self.schema()
data = _serialize_json_values(data, schema)
if on_bad_vectors is None:
on_bad_vectors = "error"
if fill_value is None:
@@ -6338,10 +6531,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``.
@@ -6403,8 +6596,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].
@@ -6421,6 +6614,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]:
+40
View File
@@ -0,0 +1,40 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright The LanceDB Authors
"""Views: a named query a database stores and plans on every read.
A view holds no rows. What it stores is the statement that defines it and the
schema that statement resolved to, so a reader sees the sources as they are
now. See ``DBConnection.create_view``.
"""
from __future__ import annotations
from dataclasses import dataclass, field
from typing import TYPE_CHECKING, List
if TYPE_CHECKING:
import pyarrow as pa
@dataclass(frozen=True)
class ViewDescription:
"""What a database records about one view."""
name: str
"""The view's name within its namespace."""
query: str
"""The defining query, as the database stores it."""
default_database: str
"""The database that unqualified table names in ``query`` resolve against."""
schema: "pa.Schema"
"""The schema the defining query resolved to when the view was created."""
namespace_path: List[str] = field(default_factory=list)
"""The namespace holding the view; empty is the root namespace."""
default_namespace_path: List[str] = field(default_factory=list)
"""The namespace path those unqualified names resolve against.
Recorded with the view because it outlives the session that declared it: a
reader resolving the query against its own default namespace could read a
different table than the view was defined over.
"""
+74
View File
@@ -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",
+140
View File
@@ -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)
+12
View File
@@ -633,6 +633,18 @@ def test_url_retrieve_downloads_image():
assert img.size[0] > 0 and img.size[1] > 0
def test_open_clip_opens_percent_encoded_file_uri(tmp_path):
"""OpenCLIP should decode local file URIs before opening them."""
Image = pytest.importorskip("PIL.Image")
from lancedb.embeddings.open_clip import OpenClipEmbeddings
image_path = tmp_path / "test image.png"
Image.new("RGB", (4, 4), color="red").save(image_path, format="PNG")
with OpenClipEmbeddings._to_pil(None, image_path.as_uri()) as image:
assert image.size == (4, 4)
def test_jina_generate_image_input_dict_local_path(tmp_path):
"""
JinaEmbeddings._generate_image_input_dict must accept a local image path
@@ -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()
@@ -65,6 +83,18 @@ def test_public_function_values_are_in_api_reference():
FunctionBinding,
False,
),
(
"remote_initialized_function_application.json",
"remote_initialized_function_application.canonical.json",
FunctionApplication,
False,
),
(
"remote_initialized_function_binding.json",
"remote_initialized_function_binding.canonical.json",
FunctionBinding,
False,
),
(
"remote_refresh_job.json",
"remote_refresh_result.canonical.json",
@@ -93,16 +123,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():
@@ -128,7 +169,7 @@ def test_function_version_binding_validates_names_and_direct_columns():
with pytest.raises(TypeError, match=r"missing inputs: \['text'\]"):
version()
with pytest.raises(TypeError, match=r"unknown inputs: \['body'\]"):
with pytest.raises(TypeError, match=r"unknown arguments: \['body'\]"):
version(text=col("text"), body=col("body"))
with pytest.raises(TypeError, match="direct col"):
version(text=col("text").lower())
@@ -137,7 +178,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 +223,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 +263,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
@@ -240,6 +281,35 @@ def test_binding_and_refresh_result_keep_stable_remote_fields():
assert RefreshColumnResult.from_json(result.to_canonical_json()) == result
def test_initialization_travels_on_the_application_and_binding():
application = FunctionApplication.from_json(
fixture("remote_initialized_function_application.json")
)
assert application.initialization["temperature"] == 0.25
assert application.initialization["retry"]["backoff"] == (1, 2, 4)
application._ensure_declarable()
binding = FunctionBinding.from_json(
fixture("remote_initialized_function_binding.json")
)
assert binding.initialization_row() == {
"model": "text-embedding-3-small",
"dimensions": 512,
}
assert (
FunctionBinding.from_json(
fixture("remote_function_binding.json")
).initialization_row()
is None
)
signature = FunctionVersion.from_json(
json.dumps(job_result("remote_function_job.json"))
).signature
assert signature.initialization == ()
assert "initialization" not in json.loads(signature.to_canonical_json())
def test_function_literal_numeric_domain_matches_rust():
with pytest.raises(ValueError, match="floating-point Function literals"):
FunctionApplication.from_json(fixture("remote_function_application_float.json"))
@@ -276,6 +346,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 +434,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,10 +11,12 @@ import types
from datetime import date
import http.server
import json
import os
from pathlib import Path
import subprocess
import sys
import threading
import urllib.parse
from typing import Optional
import pyarrow as pa
@@ -23,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
@@ -44,6 +49,11 @@ FIXTURES = (
)
FUNCTION_VERSION = json.loads(
(FIXTURES / "remote_function_version.canonical.json").read_text()
)["version"]
@udf(
pip=["numpy>=2"],
env={"MODE": "test"},
@@ -53,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
@@ -69,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:
@@ -1198,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
@@ -1213,23 +1512,47 @@ def _mock_remote_function_catalog():
def log_message(self, *args):
pass
def _write_response(self, status, response):
encoded = json.dumps(response).encode()
self.send_response(status)
self.send_header("Content-Type", "application/json")
self.send_header("Content-Length", str(len(encoded)))
self.end_headers()
self.wfile.write(encoded)
def do_POST(self):
length = int(self.headers.get("Content-Length", "0"))
body = json.loads(self.rfile.read(length) or b"{}")
state["requests"].append((self.path, body))
status = 200
if self.path == "/v1/functions/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": body["name"],
"version": "fv_exact",
"artifact": {
key: body["artifact"][key]
for key in ("kind", "digest", "entrypoint")
},
"name": "normalize_score",
"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"}
@@ -1242,43 +1565,62 @@ def _mock_remote_function_catalog():
"job_state": "DONE",
"result": state["version"],
}
elif self.path == "/v1/functions/describe":
assert body == {
"name": "normalize_score",
"version": "fv_exact",
}
elif self.path == "/v1/function/normalize_score/describe":
assert body == {"version": FUNCTION_VERSION}
response = state["version"]
elif self.path == "/v1/functions/list":
assert body["include_definition"] is True
if "page_token" not in body:
response = {
"functions": [
{
"name": "normalize_score",
"version": "fv_exact",
"definition": state["version"],
}
],
"page_token": "next",
}
else:
assert body["page_token"] == "next"
response = {"functions": []}
elif self.path == "/v1/functions/drop":
assert body == {
"name": "normalize_score",
"version": "fv_exact",
}
elif self.path == "/v1/function/normalize_score/drop":
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"}
encoded = json.dumps(response).encode()
self.send_response(status)
self.send_header("Content-Type", "application/json")
self.send_header("Content-Length", str(len(encoded)))
self.end_headers()
self.wfile.write(encoded)
self._write_response(status, response)
def do_GET(self):
url = urllib.parse.urlsplit(self.path)
query = {
key: values[-1]
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
assert query["include_definition"] == "true"
if "page_token" not in query:
response = {
"functions": [
{
"name": "normalize_score",
"version": FUNCTION_VERSION,
"definition": state["version"],
}
],
"page_token": "next",
}
else:
assert query["page_token"] == "next"
response = {"functions": []}
self._write_response(200, response)
with http.server.HTTPServer(("localhost", 0), Handler) as server:
thread = threading.Thread(target=server.serve_forever)
@@ -1305,11 +1647,99 @@ 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]
assert create_request == json.loads(
expected_request = json.loads(
normalize_score.registration_request.to_canonical_json()
)
expected_request.pop("name")
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():
@@ -1323,9 +1753,9 @@ 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/functions/create",
"/v1/function/normalize_score/create",
"/v1/jobs/describe",
]
@@ -1344,10 +1774,10 @@ def test_remote_list_functions_paginates_and_returns_typed_versions():
assert functions == [created]
assert state["requests"] == [
("/v1/functions/list", {"include_definition": True}),
("/v1/namespace/$/function/list", {"include_definition": "true"}),
(
"/v1/functions/list",
{"include_definition": True, "page_token": "next"},
"/v1/namespace/$/function/list",
{"include_definition": "true", "page_token": "next"},
),
]
@@ -1368,8 +1798,8 @@ async def test_async_remote_list_functions_returns_typed_versions():
assert functions == [created]
assert [path for path, _ in state["requests"]] == [
"/v1/functions/list",
"/v1/functions/list",
"/v1/namespace/$/function/list",
"/v1/namespace/$/function/list",
]
@@ -1381,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/functions/drop",
{"name": "normalize_score", "version": "fv_exact"},
"/v1/function/normalize_score/drop",
{"version": FUNCTION_VERSION},
)
]
@@ -1400,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/functions/drop",
{"name": "normalize_score", "version": "fv_exact"},
"/v1/function/normalize_score/drop",
{"version": FUNCTION_VERSION},
)
]
@@ -0,0 +1,521 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright The LanceDB Authors
"""Class Functions, initialization, and the source that travels with a Function."""
from __future__ import annotations
import base64
import importlib
import json
import subprocess
import sys
import textwrap
from pathlib import Path
from typing import Annotated, Optional
import pyarrow as pa
import pytest
from lancedb import col
from lancedb.functions import FunctionVersion, udf
FIXTURES = (
Path(__file__).parents[3]
/ "rust"
/ "lancedb"
/ "tests"
/ "fixtures"
/ "first_class_functions"
/ "v1"
)
def artifact(definition) -> bytes:
return base64.b64decode(definition.registration_request.artifact.content.data)
def bundle_files(definition) -> dict[str, str]:
assert definition.registration_request.artifact.kind == "python_bundle"
return json.loads(artifact(definition))["files"]
@pytest.fixture
def helper_package(tmp_path, monkeypatch):
"""A local package on sys.path, the way a project's own helpers are."""
def write(name: str, files: dict[str, str]):
root = tmp_path / name
for path, source in files.items():
target = root / path
target.parent.mkdir(parents=True, exist_ok=True)
target.write_text(textwrap.dedent(source))
monkeypatch.syspath_prepend(str(tmp_path))
for module in [
module for module in sys.modules if module.split(".")[0] == name
]:
del sys.modules[module]
importlib.invalidate_caches()
return importlib.import_module(name)
return write
@udf
class scale:
def __init__(
self,
factor: float,
label: str,
offset: float = 0.0,
tags: Optional[list[str]] = None,
retry: Annotated[dict, pa.struct([("attempts", pa.int32())])] = None,
):
self.factor, self.offset = factor, offset
def __call__(self, value: float) -> float:
return value * self.factor + self.offset
class NoCall:
pass
class BadClose:
def __call__(self, value: int) -> int:
return value
def close(self, force):
pass
def test_class_function_declares_inputs_output_and_initialization():
signature = scale.registration_request.signature
assert [(p.name, p.arrow_type) for p in signature.inputs] == [("value", "float64")]
assert [
(field.name, field.arrow_type, field.nullable)
for field in signature.initialization
] == [
("factor", "float64", False),
("label", "utf8", False),
("offset", "float64", True),
("tags", "list<utf8>", True),
(
"retry",
'{"fields":[{"name":"attempts","nullable":true,"type":{"type":"int32"}}],'
'"type":"struct"}',
True,
),
]
# Calling the definition constructs the class, so local tests are ordinary.
assert scale(factor=2.0, label="x")(1.5) == 3.0
assert scale.registration_request.artifact.entrypoint == "scale"
source = artifact(scale).decode()
assert source.startswith("from __future__ import annotations\n")
assert "class scale:" in source and "@udf" not in source
def test_function_without_initialization_keeps_its_wire_form():
@udf
def double(value: int) -> int:
return value * 2
request = json.loads(double.registration_request.to_canonical_json())
assert "initialization" not in request["signature"]
assert request["artifact"]["kind"] == "python_callable"
@pytest.mark.parametrize(
("init", "message"),
[
("def __init__(self, blob: bytes): pass", "initialization takes booleans"),
("def __init__(self, *args: int): pass", "named and non-variadic"),
("def __init__(self, factor): pass", "missing Function initialization"),
("def __init__(self, value: int): pass", "share names \\['value'\\]"),
],
)
def test_class_function_rejects_unsupported_initialization(init, message):
namespace = {}
exec(
textwrap.dedent(
f"""
class bad:
{init}
def __call__(self, value: int) -> int:
return value
"""
),
namespace,
)
with pytest.raises((TypeError, ValueError), match=message):
udf(namespace["bad"])
def test_class_function_requires_call_and_a_plain_close():
with pytest.raises(TypeError, match="synchronous __call__"):
udf(NoCall)
with pytest.raises(TypeError, match="close must be a method taking no"):
udf(BadClose)
def test_nested_class_and_closure_are_rejected_with_the_reason():
def build():
class nested:
def __call__(self, value: int) -> int:
return value
return nested
with pytest.raises(ValueError, match="defined at module level"):
udf(build())
factor = 3
with pytest.raises(ValueError, match="capture closure values"):
@udf
def closes_over(value: int) -> int:
return value * factor
def test_binding_passes_initialization_as_constants():
version = FunctionVersion.from_json(
json.dumps(
{
**json.loads(
(FIXTURES / "remote_function_version.canonical.json").read_text()
),
"signature": {
"inputs": [
{"name": "text", "arrow_type": "utf8", "nullable": True}
],
"output": {
"kind": "scalar",
"arrow_type": "list<float32>",
"nullable": False,
},
"initialization": [
{"name": "model", "arrow_type": "utf8", "nullable": False},
{"name": "dimensions", "arrow_type": "int32", "nullable": True},
],
},
}
)
)
application = version(text=col("body"), model="small", dimensions=(512))
assert dict(application.initialization) == {"model": "small", "dimensions": 512}
assert json.loads(application.to_canonical_json())["initialization"] == {
"dimensions": 512,
"model": "small",
}
assert version(text=col("body"), model="small").initialization == {"model": "small"}
with pytest.raises(TypeError, match=r"missing initialization: \['model'\]"):
version(text=col("body"))
with pytest.raises(TypeError, match="takes a constant"):
version(text=col("body"), model=col("model"))
with pytest.raises(ValueError, match="must be finite"):
version(text=col("body"), model="small", dimensions=float("nan"))
with pytest.raises(TypeError, match="takes booleans"):
version(text=col("body"), model=object())
def test_code_packages_ship_as_a_bundle_and_version_with_their_source(
tmp_path, helper_package
):
files = {
"__init__.py": "",
"rate.py": """
class TokenBucket:
def __init__(self, tokens_per_minute):
self.tokens_per_minute = tokens_per_minute
""",
"__pycache__/rate.cpython-312.pyc": "",
"fixtures/data.json": "{}",
"not-a-module.py": "",
}
helper_package("fx_helpers", files)
(tmp_path / "fx_helpers_functions.py").write_text(
textwrap.dedent(
"""
import fx_helpers
from fx_helpers.rate import TokenBucket
from lancedb import udf
@udf(code=[fx_helpers])
class limited:
def __init__(self, tpm: int):
self.limiter = TokenBucket(tpm)
def __call__(self, value: int) -> int:
return value
"""
)
)
def define():
sys.modules.pop("fx_helpers_functions", None)
return importlib.import_module("fx_helpers_functions").limited
first = define()
shipped = bundle_files(first)
assert sorted(shipped) == [
"fx_helpers/__init__.py",
"fx_helpers/rate.py",
"user_function.py",
]
assert (
"from fx_helpers.rate import TokenBucket as TokenBucket"
in (shipped["user_function.py"])
)
# Only the helper changes: a new artifact digest, the same environment.
helper_package(
"fx_helpers", {**files, "rate.py": textwrap.dedent(files["rate.py"]) + "# v2\n"}
)
second = define()
assert (
first.registration_request.artifact.digest
!= second.registration_request.artifact.digest
)
assert first.registration_request.runtime == second.registration_request.runtime
def test_code_rejects_submodules_duplicates_and_non_modules(helper_package):
helpers = helper_package("fx_rejects", {"__init__.py": "", "sub.py": ""})
import fx_rejects.sub
for code, message in [
([fx_rejects.sub], "top-level modules"),
([helpers, helpers], "more than once"),
(["fx_rejects"], "imported modules"),
([json], "cannot ship module 'json'"),
]:
with pytest.raises((TypeError, ValueError), match=message):
@udf(code=code)
def f(value: int) -> int:
return value
def test_local_module_outside_code_is_rejected_at_registration(
tmp_path, helper_package
):
helper_package(
"fx_local", {"__init__.py": "def normalize(value):\n return value\n"}
)
def define(name, options):
(tmp_path / f"{name}.py").write_text(
textwrap.dedent(
f"""
from fx_local import normalize
from lancedb import udf
@udf({options})
def {name}(value: int) -> int:
return normalize(value)
"""
)
)
return getattr(importlib.import_module(name), name)
with pytest.raises(ValueError, match=r"local module.*code=\[fx_local\]"):
define("fx_uses_local", "")
# A declared package that provides the module is trusted to install it.
installed = define("fx_uses_installed", "pip=['fx-local==1.0']")
assert "from fx_local import normalize" in artifact(installed).decode()
MMLB_HELPERS = {
"__init__.py": "",
"rate.py": """
class TokenBucket:
constructed = 0
def __init__(self, tokens_per_minute):
TokenBucket.constructed += 1
self.tokens_per_minute = tokens_per_minute
self.acquired = 0
def acquire(self, tokens):
self.acquired += tokens
""",
"embedding.py": """
from fx_mmlb.rate import TokenBucket
class OpenAIEmbedder:
def __init__(self, model: str, tpm_limit: int, max_attempts: int = 6):
self.model = model
self.limiter = TokenBucket(tpm_limit)
self.max_attempts = max_attempts
def embed(self, texts):
self.limiter.acquire(sum(len(text) for text in texts))
return [[float(len(text))] for text in texts]
""",
}
MMLB_FUNCTIONS = """
import pyarrow as pa
import fx_mmlb
from fx_mmlb.embedding import OpenAIEmbedder
from lancedb import udf
TEXT = pa.schema([pa.field("text", pa.string())])
EMBEDDING = pa.list_(pa.field("item", pa.float64(), nullable=False))
@udf(code=[fx_mmlb], input_schema=TEXT, output_schema=EMBEDDING)
class openai_embedding_batch(OpenAIEmbedder):
def __call__(self, text: pa.Array) -> pa.Array:
item = pa.field("item", pa.float64(), nullable=False)
return pa.array(self.embed(text.to_pylist()), pa.list_(item))
@udf(code=[fx_mmlb])
class openai_embedding(OpenAIEmbedder):
def __call__(self, text: str) -> list[float]:
return self.embed([text])[0]
"""
def test_shared_helpers_construct_state_on_every_path(tmp_path, helper_package):
"""The MMLB rate limiter, shared by a row and a batch Function: both are
initialized through the same helper, so neither path can skip it."""
helper_package("fx_mmlb", MMLB_HELPERS)
(tmp_path / "fx_mmlb_functions.py").write_text(MMLB_FUNCTIONS)
functions = importlib.import_module("fx_mmlb_functions")
signatures = [
functions.openai_embedding_batch.registration_request.signature,
functions.openai_embedding.registration_request.signature,
]
for signature in signatures:
assert [field.name for field in signature.initialization] == [
"model",
"tpm_limit",
"max_attempts",
]
# Each bundle imports on its own, in a process that has never seen the
# client's modules, the way the worker's image code layer does.
for definition, call in [
(
functions.openai_embedding_batch,
"instance(pa.array(['ab', 'abc'])).to_pylist()",
),
(functions.openai_embedding, "[instance('ab'), instance('abc')]"),
]:
worker = tmp_path / definition.__name__
for path, source in bundle_files(definition).items():
(worker / path).parent.mkdir(parents=True, exist_ok=True)
(worker / path).write_text(source)
probe = textwrap.dedent(
f"""
import sys
sys.path.insert(0, {str(worker)!r})
import pyarrow as pa
import user_function
from fx_mmlb.rate import TokenBucket
cls = getattr(user_function, {definition.__name__!r})
instance = cls(model="small", tpm_limit=1000)
print(repr(({call}, TokenBucket.constructed, instance.limiter.acquired)))
"""
)
result = subprocess.run(
[sys.executable, "-I", "-c", probe],
capture_output=True,
text=True,
cwd=worker,
)
assert result.returncode == 0, result.stderr
assert result.stdout.strip() == "([[2.0], [3.0]], 1, 5)"
def test_main_definitions_travel_by_source(tmp_path):
"""Notebook cells and scripts define helpers in __main__, which a worker
cannot import; they are packaged by source, dependencies first."""
script = tmp_path / "notebook.py"
script.write_text(
textwrap.dedent(
"""
import json
import math
from dataclasses import dataclass
from lancedb import udf
SCALE = 2
@dataclass
class Settings:
factor: int
def helper(value):
return math.floor(value) * SCALE
class Base:
def run(self, value):
return helper(value) * self.settings.factor
@udf
class scaled(Base):
def __init__(self, factor: int):
self.settings = Settings(factor)
def __call__(self, value: float) -> int:
return self.run(value)
print(json.dumps(scaled.registration_request.to_canonical_json()))
"""
)
)
result = subprocess.run(
[sys.executable, str(script)], capture_output=True, text=True
)
assert result.returncode == 0, result.stderr
request = json.loads(json.loads(result.stdout))
assert request["artifact"]["kind"] == "python_callable"
source = base64.b64decode(request["artifact"]["content"]["data"]).decode()
assert source.index("def helper") < source.index("class Base")
assert source.index("class Base") < source.index("class scaled")
assert source.index("class Settings") < source.index("class scaled")
assert "@dataclass" in source
assert "SCALE = 2" in source
namespace = {}
exec(compile(source, "user_function.py", "exec"), namespace)
assert namespace["scaled"](factor=3)(2.7) == 12
def test_main_lambdas_are_rejected(tmp_path):
script = tmp_path / "lambda.py"
script.write_text(
textwrap.dedent(
"""
from lancedb import udf
double = lambda value: value * 2
@udf
def uses_lambda(value: int) -> int:
return double(value)
"""
)
)
result = subprocess.run(
[sys.executable, str(script)], capture_output=True, text=True
)
assert result.returncode != 0
assert "not lambdas, closures, or nested definitions" in result.stderr
+274 -48
View File
@@ -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,32 +183,47 @@ 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")
view = db.open_materialized_view("adults")
assert view.definition == MaterializedViewDefinition(
source_table="people",
projections=[("name", "`name`"), ("age", "`age`")],
filter="age >= 18",
inputs=["age", "name"],
query="SELECT name, age FROM people WHERE age >= 18"
)
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 +237,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 +254,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 +271,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 +291,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"
@@ -127,15 +302,39 @@ async def test_async_create_refresh_and_open(tmp_path):
reopened = await db.open_materialized_view("shouts")
definition = await reopened.definition()
assert definition.projections == [("shout", "upper(name)")]
assert definition.query == "SELECT upper(name) AS shout FROM people"
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 +356,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 +368,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")
@@ -231,10 +420,21 @@ def test_namespace_connection_materialized_views(tmp_path):
assert db.list_materialized_views() == ["adults"]
reopened = db.open_materialized_view("adults")
assert reopened.definition.source_table == "people"
assert reopened.definition.query.startswith("SELECT name, age FROM ")
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):
@@ -255,7 +455,7 @@ async def test_async_namespace_connection_materialized_views(tmp_path):
assert await db.list_materialized_views() == ["adults"]
reopened = await db.open_materialized_view("adults")
assert (await reopened.definition()).source_table == "people"
assert (await reopened.definition()).query.startswith("SELECT name, age FROM ")
# The view's table came through the namespace, not straight from the
# inner connection: a bare inner table carries no namespace context, so
@@ -267,37 +467,63 @@ 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
def test_namespaced_select_kind_is_read_and_unknown_kinds_are_refused():
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_stored_queries_and_legacy_layouts_are_read():
import json
import pyarrow as pa
from lancedb.materialized_view import _definition_from_schema
def schema_with(definition: dict) -> pa.Schema:
return pa.schema([pa.field("id", pa.int32())]).with_metadata(
def read(definition: dict) -> MaterializedViewDefinition:
schema = pa.schema([pa.field("id", pa.int32())]).with_metadata(
{b"mv.definition": json.dumps(definition).encode()}
)
return _definition_from_schema(schema, "v")
# "namespaced_select" is the namespaced form of "select": same shape,
# a separate kind so readers that predate it refuse instead of
# resolving the source at the root.
definition = _definition_from_schema(
schema_with(
query = "SELECT id, c.chunk FROM ns.docs, UNNEST(chunks) AS c WHERE id > 1"
assert read({"format": 1, "query": query}).query == query
# The structured layout written before the format number reads as the
# query it described, under either of its kind tags.
assert (
read(
{
"kind": "namespaced_select",
"source_table": "people",
"source_namespace": ["ns"],
"projections": [{"output": "name", "expression": "name"}],
"projections": [
{"output": "name", "expression": "`name`"},
{"output": "Shout", "expression": "upper(name)"},
],
"filter": "age >= 18",
"limit": 10,
}
),
"v",
).query
== "SELECT `name`, upper(name) AS `Shout` FROM ns.people "
"WHERE age >= 18 LIMIT 10"
)
assert read({"kind": "select", "source_table": "people"}).query == (
"SELECT * FROM people"
)
assert definition.source_table == "people"
assert definition.source_namespace == ["ns"]
with pytest.raises(NotImplementedError, match="cannot refresh"):
_definition_from_schema(
schema_with({"kind": "select_v3", "source_table": "people"}), "v"
)
# A newer writer's layout is reported, never guessed at.
for newer in (
{"format": 3, "query": query},
{"kind": "select_v3", "source_table": "people"},
):
with pytest.raises(NotImplementedError, match="cannot refresh"):
read(newer)
@@ -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()
@@ -9,6 +9,15 @@ from pathlib import Path
import pytest
@pytest.mark.parametrize("name", ["JobInfo", "JobDescription", "JobFailureInfo"])
def test_job_metadata_types_have_resolvable_modules(name):
"""Documentation tools resolve re-exports through each type's module."""
public_type = getattr(importlib.import_module("lancedb.job"), name)
defining_module = importlib.import_module(public_type.__module__)
assert getattr(defining_module, public_type.__name__, None) is public_type
def test_pyo3_abi_matches_minimum_supported_python():
project_dir = Path(__file__).parents[2]
pyproject = (project_dir / "pyproject.toml").read_text()
+119
View File
@@ -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
@@ -2682,3 +2747,57 @@ def test_remote_job_handle_reports_its_own_detail():
"limit": 500,
"filter": "state = 'claim_complete'",
}
def test_view_crud_addresses_its_own_routes():
# The view verbs are their own routes, and the schema comes back in the
# namespace spec's JSON encoding, decoded into a pyarrow schema.
paths = []
def handler(request):
paths.append((request.command, request.path))
if request.path.endswith("/view/list"):
body = {"views": ["adults"]}
elif request.path.endswith("/drop"):
body = {}
else:
body = {
"name": "adults",
"namespace": ["analytics"],
"query": "SELECT name FROM people",
"default_database": "dev",
"default_namespace": ["analytics"],
"schema": {
"fields": [
{"name": "name", "nullable": True, "type": {"type": "utf8"}}
]
},
}
request.send_response(200)
request.send_header("Content-Type", "application/json")
request.end_headers()
request.wfile.write(json.dumps(body).encode())
with mock_lancedb_connection(handler) as db:
view = db.create_view(
"adults", "SELECT name FROM people", namespace_path=["analytics"]
)
assert view.name == "adults"
assert view.namespace_path == ["analytics"]
assert view.query == "SELECT name FROM people"
assert view.default_database == "dev"
assert view.default_namespace_path == ["analytics"]
assert view.schema == pa.schema([pa.field("name", pa.utf8(), nullable=True)])
described = db.describe_view("adults", namespace_path=["analytics"])
assert described.schema == view.schema
assert db.list_views(namespace_path=["analytics"]) == ["adults"]
db.drop_view("adults", namespace_path=["analytics"])
assert paths == [
("POST", "/v1/view/analytics$adults/create"),
("POST", "/v1/view/analytics$adults/describe"),
("GET", "/v1/namespace/analytics/view/list"),
("POST", "/v1/view/analytics$adults/drop"),
]
+104 -3
View File
@@ -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://")
+420 -23
View File
@@ -4,6 +4,7 @@
import ctypes
import gc
import json
import os
import sys
import threading
@@ -17,6 +18,7 @@ from typing import List
from unittest.mock import patch
import lancedb
from lancedb import table as table_module
from lancedb.dependencies import _PANDAS_AVAILABLE
from lancedb.index import BTree, FTS, HnswFlat, HnswPq, HnswSq, IvfPq
import numpy as np
@@ -2979,28 +2981,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 +3018,406 @@ 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")
def test_add_python_objects_to_json_column(mem_db: DBConnection):
schema = pa.schema(
[
pa.field("id", pa.int64()),
pa.field("payload", pa.json_()),
pa.field("metadata", pa.struct([("label", pa.string())])),
pa.field("tags", pa.list_(pa.string())),
]
)
table = mem_db.create_table("json_python_objects_add", schema=schema)
table.add(
[
{
"id": 1,
"payload": {"foo": "bar", "count": 2},
"metadata": {"label": "dict"},
"tags": ["a", "b"],
},
{
"id": 2,
"payload": {
"name": "alice",
"tags": ["x", "y"],
"nested": {"enabled": True},
},
"metadata": {"label": "nested"},
"tags": ["c"],
},
{
"id": 3,
"payload": ["x", "y", "z"],
"metadata": {"label": "list"},
"tags": ["d", "e"],
},
{
"id": 4,
"payload": '{"foo": "bar"}',
"metadata": {"label": "string"},
"tags": ["f"],
},
{
"id": 5,
"payload": None,
"metadata": {"label": "null"},
"tags": ["g"],
},
]
)
rows = {row["id"]: row for row in table.to_arrow().to_pylist()}
assert json.loads(rows[1]["payload"]) == {"foo": "bar", "count": 2}
assert json.loads(rows[2]["payload"]) == {
"name": "alice",
"tags": ["x", "y"],
"nested": {"enabled": True},
}
assert json.loads(rows[3]["payload"]) == ["x", "y", "z"]
assert json.loads(rows[4]["payload"]) == {"foo": "bar"}
assert rows[5]["payload"] is None
assert rows[1]["metadata"] == {"label": "dict"}
assert rows[1]["tags"] == ["a", "b"]
matched = table.search().where("json_extract(payload, '$.count') = '2'").to_list()
assert [row["id"] for row in matched] == [1]
@pytest.mark.skipif(not hasattr(pa, "json_"), reason="requires PyArrow JSON type")
def test_merge_insert_python_objects_to_json_column(mem_db: DBConnection):
schema = pa.schema(
[
pa.field("id", pa.int64()),
pa.field("payload", pa.json_()),
pa.field("metadata", pa.struct([("label", pa.string())])),
pa.field("tags", pa.list_(pa.string())),
]
)
table = mem_db.create_table("json_python_objects_merge", schema=schema)
table.add(
[
{
"id": 1,
"payload": {"old": True},
"metadata": {"label": "old"},
"tags": ["old"],
}
]
)
table.merge_insert(
"id"
).when_matched_update_all().when_not_matched_insert_all().execute(
[
{
"id": 1,
"payload": {"foo": "bar", "count": 2},
"metadata": {"label": "dict"},
"tags": ["a", "b"],
},
{
"id": 2,
"payload": {
"name": "alice",
"tags": ["x", "y"],
"nested": {"enabled": True},
},
"metadata": {"label": "nested"},
"tags": ["c"],
},
{
"id": 3,
"payload": ["x", "y", "z"],
"metadata": {"label": "list"},
"tags": ["d", "e"],
},
{
"id": 4,
"payload": '{"foo": "bar"}',
"metadata": {"label": "string"},
"tags": ["f"],
},
{
"id": 5,
"payload": None,
"metadata": {"label": "null"},
"tags": ["g"],
},
]
)
rows = {row["id"]: row for row in table.to_arrow().to_pylist()}
assert json.loads(rows[1]["payload"]) == {"foo": "bar", "count": 2}
assert json.loads(rows[2]["payload"]) == {
"name": "alice",
"tags": ["x", "y"],
"nested": {"enabled": True},
}
assert json.loads(rows[3]["payload"]) == ["x", "y", "z"]
assert json.loads(rows[4]["payload"]) == {"foo": "bar"}
assert rows[5]["payload"] is None
assert rows[1]["metadata"] == {"label": "dict"}
assert rows[1]["tags"] == ["a", "b"]
matched = table.search().where("json_extract(payload, '$.count') = '2'").to_list()
assert [row["id"] for row in matched] == [1]
@pytest.mark.skipif(not hasattr(pa, "json_"), reason="requires PyArrow JSON type")
def test_add_python_objects_to_nested_json_fields(mem_db: DBConnection):
schema = pa.schema(
[
pa.field("id", pa.int64()),
pa.field(
"info",
pa.struct([pa.field("payload", pa.json_())]),
),
pa.field("documents", pa.list_(pa.field("item", pa.json_()))),
]
)
table = mem_db.create_table("nested_json_python_objects_add", schema=schema)
table.add(
[
{
"id": 1,
"info": {"payload": {"kind": "struct", "value": 1}},
"documents": [{"kind": "list", "value": 2}],
}
]
)
row = table.to_arrow().to_pylist()[0]
assert json.loads(row["info"]["payload"]) == {"kind": "struct", "value": 1}
assert json.loads(row["documents"][0]) == {"kind": "list", "value": 2}
assert [
row["id"]
for row in table.search()
.where("json_extract(info.payload, '$.value') = '1'")
.to_list()
] == [1]
assert [
row["id"]
for row in table.search()
.where("json_extract(documents[1], '$.value') = '2'")
.to_list()
] == [1]
@pytest.mark.skipif(not hasattr(pa, "json_"), reason="requires PyArrow JSON type")
def test_json_serialization_plan_skips_non_json_branches():
schema = pa.schema(
[
pa.field("id", pa.int64()),
pa.field("vector", pa.list_(pa.float32(), 3)),
pa.field("payload", pa.json_()),
pa.field(
"metadata",
pa.struct(
[
pa.field("label", pa.string()),
pa.field("payload", pa.json_()),
]
),
),
pa.field(
"documents",
pa.list_(
pa.struct(
[
pa.field("title", pa.string()),
pa.field("payload", pa.json_()),
]
)
),
),
]
)
plans = table_module._json_serialization_plans(schema)
assert set(plans) == {"payload", "metadata", "documents"}
metadata_children = plans["metadata"].children
assert metadata_children is not None
assert set(metadata_children) == {"payload"}
documents_item = plans["documents"].item
assert documents_item is not None
assert documents_item.children is not None
assert set(documents_item.children) == {"payload"}
@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 +4559,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 +4579,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 +4605,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]
+125
View File
@@ -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()?,
})
})
}
+248 -4
View File
@@ -13,7 +13,11 @@ use crate::{
runtime::future_into_py,
table::Table,
};
use arrow::{datatypes::Schema, ffi_stream::ArrowArrayStreamReader, pyarrow::FromPyArrow};
use arrow::{
datatypes::Schema,
ffi_stream::ArrowArrayStreamReader,
pyarrow::{FromPyArrow, ToPyArrow},
};
use lancedb::{
connection::Connection as LanceConnection,
connection::NamespaceClientPushdownOperation,
@@ -100,6 +104,28 @@ fn parse_default_namespace_path(path: Option<Bound<'_, PyAny>>) -> PyResult<Vec<
}
}
/// A view description on its way to Python: name, namespace, query, default
/// database, and the schema as pyarrow renders it.
type PyViewDescription = (String, Vec<String>, String, String, Vec<String>, Py<PyAny>);
/// A view description as a plain tuple, with the schema converted to the
/// pyarrow schema the caller would get from any other lancedb API. The Python
/// layer names the fields; this keeps the binding free of a class that would
/// have to be kept in step with the Rust struct.
fn view_description_to_py(view: lancedb::view::ViewDescription) -> PyResult<PyViewDescription> {
Python::attach(|py| {
let schema = view.schema.to_pyarrow(py)?.unbind();
Ok((
view.name,
view.namespace_path,
view.query,
view.default_database,
view.default_namespace_path,
schema,
))
})
}
#[pymethods]
impl Connection {
fn __repr__(&self) -> String {
@@ -381,7 +407,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 +415,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 +429,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 +794,160 @@ impl Connection {
})
}
pub fn drop_function_async(
self_: PyRef<'_, Self>,
name: String,
version: String,
) -> PyResult<Bound<'_, PyAny>> {
let inner = self_.get_inner()?.clone();
future_into_py(self_.py(), async move {
inner
.drop_function_async(name, version)
.await
.infer_error()
.map(|(dropped, job)| (dropped, crate::job::Job::new(job)))
})
}
#[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))
})
}
#[pyo3(signature = (name, query, namespace_path=None))]
pub fn create_view(
self_: PyRef<'_, Self>,
name: String,
query: 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 view = inner
.create_view(name, query, &namespace_path)
.await
.infer_error()?;
view_description_to_py(view)
})
}
#[pyo3(signature = (name, namespace_path=None))]
pub fn describe_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 {
let view = inner
.describe_view(name, &namespace_path)
.await
.infer_error()?;
view_description_to_py(view)
})
}
#[pyo3(signature = (name, namespace_path=None))]
pub fn drop_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_view(name, &namespace_path).await.infer_error()
})
}
#[pyo3(signature = (namespace_path=None))]
pub fn list_views(
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_views(&namespace_path).await.infer_error()
})
}
pub fn list_jobs(self_: PyRef<'_, Self>) -> PyResult<Bound<'_, PyAny>> {
let inner = self_.get_inner()?.clone();
future_into_py(self_.py(), async move {
+3
View File
@@ -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())),
+3 -3
View File
@@ -151,7 +151,7 @@ impl Job {
}
/// A row from `Connection.list_jobs`: one server-side job.
#[pyclass(get_all, skip_from_py_object)]
#[pyclass(module = "lancedb._lancedb", get_all, skip_from_py_object)]
#[derive(Clone)]
pub struct JobInfo {
job_id: String,
@@ -184,7 +184,7 @@ impl From<lancedb::database::JobInfo> for JobInfo {
}
/// The server's account of why a job failed.
#[pyclass(get_all, skip_from_py_object)]
#[pyclass(module = "lancedb._lancedb", get_all, skip_from_py_object)]
#[derive(Clone)]
pub struct JobFailureInfo {
phase: Option<String>,
@@ -203,7 +203,7 @@ impl JobFailureInfo {
}
/// The server-side record behind a `Job` handle.
#[pyclass(get_all, skip_from_py_object)]
#[pyclass(module = "lancedb._lancedb", get_all, skip_from_py_object)]
#[derive(Clone)]
pub struct JobDescription {
job_id: String,
+35 -5
View File
@@ -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
View File
@@ -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
View File
@@ -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(&current)), 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);
}
}
+175
View File
@@ -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={})",
@@ -1640,6 +1762,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>,
@@ -1660,6 +1801,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()?;
view.definition().to_json().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>,
+347
View File
@@ -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()
+8 -8
View File
@@ -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]]