Compare commits

..

8 Commits

Author SHA1 Message Date
Yang Cen ef752c2e3f Merge remote-tracking branch 'origin/main' into yang/udf-secret-api 2026-08-28 08:13:12 +08:00
Drew 83cff3ab93 fix(python): use one blobv2 type and coerce blob writes by metadata (#4065) 2026-08-27 17:07:29 -07:00
Yang Cen 0f6a355593 Merge remote-tracking branch 'origin/main' into yang/udf-secret-api 2026-08-28 07:46:10 +08:00
Will Jones b85776c22a fix(listing)!: page table listings from the store's own cursor (#3979)
BREAKING CHANGE: list_tables now provides tables in arbitrary order and
the page token is now completely opaque. `table_names` retains the old
behavior of lexical ordering and `start-after` semantics.

Listing the tables in a directory database cost what the database held
rather than what the page held. `ListingDatabase::list_tables`
enumerated every child directory of the base path, sorted the names,
then discarded all but the requested page — on every request, for every
page. On object storage that is one full listing per page.

This PR pages the store instead. `list_tables` asks for one page at a
time through `ObjectStore::read_dir_page`, carrying the store's own
continuation token, so a page is one request. Non-table children can
leave a page short of its limit, so the walk continues until the page is
full or the store runs out.

---------

Co-authored-by: Claude Opus 5 (1M context) <noreply@anthropic.com>
2026-08-27 15:52:41 -07:00
Yang Cen a6ec35502a fix(functions): complete secret validation parity 2026-08-28 00:09:04 +08:00
Yang Cen a0bb1f7597 fix(functions): harden Rust secret submissions 2026-08-27 23:03:27 +08:00
Yang Cen 2562e117b2 fix(functions): harden UDF secret submissions 2026-08-27 21:52:31 +08:00
Yang Cen 134a265ee2 feat(functions): support UDF secret values 2026-08-27 20:58:53 +08:00
28 changed files with 2027 additions and 665 deletions
Generated
+42 -42
View File
@@ -3455,8 +3455,8 @@ checksum = "42703706b716c37f96a77aea830392ad231f44c9e9a67872fa5548707e11b11c"
[[package]]
name = "fsst"
version = "12.0.0-beta.4"
source = "git+https://github.com/lance-format/lance.git?tag=v12.0.0-beta.4#474ae88dd75c437d5d506859b627510978e502e9"
version = "12.0.0-beta.2"
source = "git+https://github.com/lance-format/lance.git?tag=v12.0.0-beta.2#dafa4642658d996b3e31dde91e02f72db7860d7e"
dependencies = [
"arrow-array",
"rand 0.9.5",
@@ -4815,8 +4815,8 @@ checksum = "e037a2e1d8d5fdbd49b16a4ea09d5d6401c1f29eca5ff29d03d3824dba16256a"
[[package]]
name = "lance"
version = "12.0.0-beta.4"
source = "git+https://github.com/lance-format/lance.git?tag=v12.0.0-beta.4#474ae88dd75c437d5d506859b627510978e502e9"
version = "12.0.0-beta.2"
source = "git+https://github.com/lance-format/lance.git?tag=v12.0.0-beta.2#dafa4642658d996b3e31dde91e02f72db7860d7e"
dependencies = [
"arc-swap",
"arrow",
@@ -4888,8 +4888,8 @@ dependencies = [
[[package]]
name = "lance-arrow"
version = "12.0.0-beta.4"
source = "git+https://github.com/lance-format/lance.git?tag=v12.0.0-beta.4#474ae88dd75c437d5d506859b627510978e502e9"
version = "12.0.0-beta.2"
source = "git+https://github.com/lance-format/lance.git?tag=v12.0.0-beta.2#dafa4642658d996b3e31dde91e02f72db7860d7e"
dependencies = [
"arrow-array",
"arrow-buffer",
@@ -4911,7 +4911,7 @@ dependencies = [
[[package]]
name = "lance-arrow-scalar"
version = "58.0.0"
source = "git+https://github.com/lance-format/lance.git?tag=v12.0.0-beta.4#474ae88dd75c437d5d506859b627510978e502e9"
source = "git+https://github.com/lance-format/lance.git?tag=v12.0.0-beta.2#dafa4642658d996b3e31dde91e02f72db7860d7e"
dependencies = [
"arrow-array",
"arrow-buffer",
@@ -4925,7 +4925,7 @@ dependencies = [
[[package]]
name = "lance-arrow-stats"
version = "58.0.0"
source = "git+https://github.com/lance-format/lance.git?tag=v12.0.0-beta.4#474ae88dd75c437d5d506859b627510978e502e9"
source = "git+https://github.com/lance-format/lance.git?tag=v12.0.0-beta.2#dafa4642658d996b3e31dde91e02f72db7860d7e"
dependencies = [
"arrow-array",
"arrow-schema",
@@ -4934,8 +4934,8 @@ dependencies = [
[[package]]
name = "lance-bitpacking"
version = "12.0.0-beta.4"
source = "git+https://github.com/lance-format/lance.git?tag=v12.0.0-beta.4#474ae88dd75c437d5d506859b627510978e502e9"
version = "12.0.0-beta.2"
source = "git+https://github.com/lance-format/lance.git?tag=v12.0.0-beta.2#dafa4642658d996b3e31dde91e02f72db7860d7e"
dependencies = [
"arrayref",
"crunchy",
@@ -4945,8 +4945,8 @@ dependencies = [
[[package]]
name = "lance-core"
version = "12.0.0-beta.4"
source = "git+https://github.com/lance-format/lance.git?tag=v12.0.0-beta.4#474ae88dd75c437d5d506859b627510978e502e9"
version = "12.0.0-beta.2"
source = "git+https://github.com/lance-format/lance.git?tag=v12.0.0-beta.2#dafa4642658d996b3e31dde91e02f72db7860d7e"
dependencies = [
"arrow-array",
"arrow-buffer",
@@ -4983,8 +4983,8 @@ dependencies = [
[[package]]
name = "lance-datafusion"
version = "12.0.0-beta.4"
source = "git+https://github.com/lance-format/lance.git?tag=v12.0.0-beta.4#474ae88dd75c437d5d506859b627510978e502e9"
version = "12.0.0-beta.2"
source = "git+https://github.com/lance-format/lance.git?tag=v12.0.0-beta.2#dafa4642658d996b3e31dde91e02f72db7860d7e"
dependencies = [
"arrow",
"arrow-array",
@@ -5013,8 +5013,8 @@ dependencies = [
[[package]]
name = "lance-datagen"
version = "12.0.0-beta.4"
source = "git+https://github.com/lance-format/lance.git?tag=v12.0.0-beta.4#474ae88dd75c437d5d506859b627510978e502e9"
version = "12.0.0-beta.2"
source = "git+https://github.com/lance-format/lance.git?tag=v12.0.0-beta.2#dafa4642658d996b3e31dde91e02f72db7860d7e"
dependencies = [
"arrow",
"arrow-array",
@@ -5031,8 +5031,8 @@ dependencies = [
[[package]]
name = "lance-derive"
version = "12.0.0-beta.4"
source = "git+https://github.com/lance-format/lance.git?tag=v12.0.0-beta.4#474ae88dd75c437d5d506859b627510978e502e9"
version = "12.0.0-beta.2"
source = "git+https://github.com/lance-format/lance.git?tag=v12.0.0-beta.2#dafa4642658d996b3e31dde91e02f72db7860d7e"
dependencies = [
"proc-macro2",
"quote",
@@ -5041,8 +5041,8 @@ dependencies = [
[[package]]
name = "lance-encoding"
version = "12.0.0-beta.4"
source = "git+https://github.com/lance-format/lance.git?tag=v12.0.0-beta.4#474ae88dd75c437d5d506859b627510978e502e9"
version = "12.0.0-beta.2"
source = "git+https://github.com/lance-format/lance.git?tag=v12.0.0-beta.2#dafa4642658d996b3e31dde91e02f72db7860d7e"
dependencies = [
"arrow-arith",
"arrow-array",
@@ -5075,8 +5075,8 @@ dependencies = [
[[package]]
name = "lance-file"
version = "12.0.0-beta.4"
source = "git+https://github.com/lance-format/lance.git?tag=v12.0.0-beta.4#474ae88dd75c437d5d506859b627510978e502e9"
version = "12.0.0-beta.2"
source = "git+https://github.com/lance-format/lance.git?tag=v12.0.0-beta.2#dafa4642658d996b3e31dde91e02f72db7860d7e"
dependencies = [
"arrow-arith",
"arrow-array",
@@ -5107,8 +5107,8 @@ dependencies = [
[[package]]
name = "lance-index"
version = "12.0.0-beta.4"
source = "git+https://github.com/lance-format/lance.git?tag=v12.0.0-beta.4#474ae88dd75c437d5d506859b627510978e502e9"
version = "12.0.0-beta.2"
source = "git+https://github.com/lance-format/lance.git?tag=v12.0.0-beta.2#dafa4642658d996b3e31dde91e02f72db7860d7e"
dependencies = [
"arc-swap",
"arrow",
@@ -5172,8 +5172,8 @@ dependencies = [
[[package]]
name = "lance-index-core"
version = "12.0.0-beta.4"
source = "git+https://github.com/lance-format/lance.git?tag=v12.0.0-beta.4#474ae88dd75c437d5d506859b627510978e502e9"
version = "12.0.0-beta.2"
source = "git+https://github.com/lance-format/lance.git?tag=v12.0.0-beta.2#dafa4642658d996b3e31dde91e02f72db7860d7e"
dependencies = [
"arrow-array",
"arrow-schema",
@@ -5195,8 +5195,8 @@ dependencies = [
[[package]]
name = "lance-io"
version = "12.0.0-beta.4"
source = "git+https://github.com/lance-format/lance.git?tag=v12.0.0-beta.4#474ae88dd75c437d5d506859b627510978e502e9"
version = "12.0.0-beta.2"
source = "git+https://github.com/lance-format/lance.git?tag=v12.0.0-beta.2#dafa4642658d996b3e31dde91e02f72db7860d7e"
dependencies = [
"arrow",
"arrow-array",
@@ -5236,8 +5236,8 @@ dependencies = [
[[package]]
name = "lance-linalg"
version = "12.0.0-beta.4"
source = "git+https://github.com/lance-format/lance.git?tag=v12.0.0-beta.4#474ae88dd75c437d5d506859b627510978e502e9"
version = "12.0.0-beta.2"
source = "git+https://github.com/lance-format/lance.git?tag=v12.0.0-beta.2#dafa4642658d996b3e31dde91e02f72db7860d7e"
dependencies = [
"arrow-array",
"arrow-schema",
@@ -5251,8 +5251,8 @@ dependencies = [
[[package]]
name = "lance-namespace"
version = "12.0.0-beta.4"
source = "git+https://github.com/lance-format/lance.git?tag=v12.0.0-beta.4#474ae88dd75c437d5d506859b627510978e502e9"
version = "12.0.0-beta.2"
source = "git+https://github.com/lance-format/lance.git?tag=v12.0.0-beta.2#dafa4642658d996b3e31dde91e02f72db7860d7e"
dependencies = [
"arrow",
"async-trait",
@@ -5264,8 +5264,8 @@ dependencies = [
[[package]]
name = "lance-namespace-impls"
version = "12.0.0-beta.4"
source = "git+https://github.com/lance-format/lance.git?tag=v12.0.0-beta.4#474ae88dd75c437d5d506859b627510978e502e9"
version = "12.0.0-beta.2"
source = "git+https://github.com/lance-format/lance.git?tag=v12.0.0-beta.2#dafa4642658d996b3e31dde91e02f72db7860d7e"
dependencies = [
"arrow",
"arrow-ipc",
@@ -5318,8 +5318,8 @@ dependencies = [
[[package]]
name = "lance-select"
version = "12.0.0-beta.4"
source = "git+https://github.com/lance-format/lance.git?tag=v12.0.0-beta.4#474ae88dd75c437d5d506859b627510978e502e9"
version = "12.0.0-beta.2"
source = "git+https://github.com/lance-format/lance.git?tag=v12.0.0-beta.2#dafa4642658d996b3e31dde91e02f72db7860d7e"
dependencies = [
"arrow-array",
"arrow-buffer",
@@ -5333,8 +5333,8 @@ dependencies = [
[[package]]
name = "lance-table"
version = "12.0.0-beta.4"
source = "git+https://github.com/lance-format/lance.git?tag=v12.0.0-beta.4#474ae88dd75c437d5d506859b627510978e502e9"
version = "12.0.0-beta.2"
source = "git+https://github.com/lance-format/lance.git?tag=v12.0.0-beta.2#dafa4642658d996b3e31dde91e02f72db7860d7e"
dependencies = [
"arrow",
"arrow-array",
@@ -5374,8 +5374,8 @@ dependencies = [
[[package]]
name = "lance-testing"
version = "12.0.0-beta.4"
source = "git+https://github.com/lance-format/lance.git?tag=v12.0.0-beta.4#474ae88dd75c437d5d506859b627510978e502e9"
version = "12.0.0-beta.2"
source = "git+https://github.com/lance-format/lance.git?tag=v12.0.0-beta.2#dafa4642658d996b3e31dde91e02f72db7860d7e"
dependencies = [
"arrow-array",
"arrow-schema",
@@ -5388,8 +5388,8 @@ dependencies = [
[[package]]
name = "lance-tokenizer"
version = "12.0.0-beta.4"
source = "git+https://github.com/lance-format/lance.git?tag=v12.0.0-beta.4#474ae88dd75c437d5d506859b627510978e502e9"
version = "12.0.0-beta.2"
source = "git+https://github.com/lance-format/lance.git?tag=v12.0.0-beta.2#dafa4642658d996b3e31dde91e02f72db7860d7e"
dependencies = [
"frostem",
"icu_segmenter",
+14 -14
View File
@@ -13,20 +13,20 @@ categories = ["database-implementations"]
rust-version = "1.91.0"
[workspace.dependencies]
lance = { "version" = "=12.0.0-beta.4", default-features = false, "tag" = "v12.0.0-beta.4", "git" = "https://github.com/lance-format/lance.git" }
lance-core = { "version" = "=12.0.0-beta.4", "tag" = "v12.0.0-beta.4", "git" = "https://github.com/lance-format/lance.git" }
lance-datagen = { "version" = "=12.0.0-beta.4", "tag" = "v12.0.0-beta.4", "git" = "https://github.com/lance-format/lance.git" }
lance-file = { "version" = "=12.0.0-beta.4", "tag" = "v12.0.0-beta.4", "git" = "https://github.com/lance-format/lance.git" }
lance-io = { "version" = "=12.0.0-beta.4", default-features = false, "tag" = "v12.0.0-beta.4", "git" = "https://github.com/lance-format/lance.git" }
lance-index = { "version" = "=12.0.0-beta.4", "tag" = "v12.0.0-beta.4", "git" = "https://github.com/lance-format/lance.git" }
lance-linalg = { "version" = "=12.0.0-beta.4", "tag" = "v12.0.0-beta.4", "git" = "https://github.com/lance-format/lance.git" }
lance-namespace = { "version" = "=12.0.0-beta.4", "tag" = "v12.0.0-beta.4", "git" = "https://github.com/lance-format/lance.git" }
lance-namespace-impls = { "version" = "=12.0.0-beta.4", default-features = false, "tag" = "v12.0.0-beta.4", "git" = "https://github.com/lance-format/lance.git" }
lance-table = { "version" = "=12.0.0-beta.4", "tag" = "v12.0.0-beta.4", "git" = "https://github.com/lance-format/lance.git" }
lance-testing = { "version" = "=12.0.0-beta.4", "tag" = "v12.0.0-beta.4", "git" = "https://github.com/lance-format/lance.git" }
lance-datafusion = { "version" = "=12.0.0-beta.4", "tag" = "v12.0.0-beta.4", "git" = "https://github.com/lance-format/lance.git" }
lance-encoding = { "version" = "=12.0.0-beta.4", "tag" = "v12.0.0-beta.4", "git" = "https://github.com/lance-format/lance.git" }
lance-arrow = { "version" = "=12.0.0-beta.4", "tag" = "v12.0.0-beta.4", "git" = "https://github.com/lance-format/lance.git" }
lance = { "version" = "=12.0.0-beta.2", default-features = false, "tag" = "v12.0.0-beta.2", "git" = "https://github.com/lance-format/lance.git" }
lance-core = { "version" = "=12.0.0-beta.2", "tag" = "v12.0.0-beta.2", "git" = "https://github.com/lance-format/lance.git" }
lance-datagen = { "version" = "=12.0.0-beta.2", "tag" = "v12.0.0-beta.2", "git" = "https://github.com/lance-format/lance.git" }
lance-file = { "version" = "=12.0.0-beta.2", "tag" = "v12.0.0-beta.2", "git" = "https://github.com/lance-format/lance.git" }
lance-io = { "version" = "=12.0.0-beta.2", default-features = false, "tag" = "v12.0.0-beta.2", "git" = "https://github.com/lance-format/lance.git" }
lance-index = { "version" = "=12.0.0-beta.2", "tag" = "v12.0.0-beta.2", "git" = "https://github.com/lance-format/lance.git" }
lance-linalg = { "version" = "=12.0.0-beta.2", "tag" = "v12.0.0-beta.2", "git" = "https://github.com/lance-format/lance.git" }
lance-namespace = { "version" = "=12.0.0-beta.2", "tag" = "v12.0.0-beta.2", "git" = "https://github.com/lance-format/lance.git" }
lance-namespace-impls = { "version" = "=12.0.0-beta.2", default-features = false, "tag" = "v12.0.0-beta.2", "git" = "https://github.com/lance-format/lance.git" }
lance-table = { "version" = "=12.0.0-beta.2", "tag" = "v12.0.0-beta.2", "git" = "https://github.com/lance-format/lance.git" }
lance-testing = { "version" = "=12.0.0-beta.2", "tag" = "v12.0.0-beta.2", "git" = "https://github.com/lance-format/lance.git" }
lance-datafusion = { "version" = "=12.0.0-beta.2", "tag" = "v12.0.0-beta.2", "git" = "https://github.com/lance-format/lance.git" }
lance-encoding = { "version" = "=12.0.0-beta.2", "tag" = "v12.0.0-beta.2", "git" = "https://github.com/lance-format/lance.git" }
lance-arrow = { "version" = "=12.0.0-beta.2", "tag" = "v12.0.0-beta.2", "git" = "https://github.com/lance-format/lance.git" }
lancedb = { path = "rust/lancedb", default-features = false }
ahash = "0.8"
# Note that this one does not include pyarrow
+6 -2
View File
@@ -223,9 +223,13 @@ tokens = list(
Blob columns store large binary values out of line so they can be read lazily
instead of being materialized with the rest of the row.
::: lancedb.blob
`lancedb.BlobType` is `lance.blob.BlobType` when pylance is installed. Without
pylance, LanceDB uses a matching `lance.blob.v2` extension type so blob columns
still work. Queries return descriptors. Call
[`fetch_blob_files`][lancedb.table.Table.fetch_blob_files] for lazy reads or
[`fetch_blobs`][lancedb.table.Table.fetch_blobs] for eager bytes.
::: lancedb.BlobType
::: lancedb.blob
::: lancedb._blob.BlobFile
options:
+1 -1
View File
@@ -28,7 +28,7 @@
<properties>
<project.build.sourceEncoding>UTF-8</project.build.sourceEncoding>
<arrow.version>15.0.0</arrow.version>
<lance-core.version>12.0.0-beta.4</lance-core.version>
<lance-core.version>12.0.0-beta.2</lance-core.version>
<spotless.skip>false</spotless.skip>
<spotless.version>2.30.0</spotless.version>
<spotless.java.googlejavaformat.version>1.7</spotless.java.googlejavaformat.version>
+15 -2
View File
@@ -6,7 +6,7 @@ import importlib.metadata
import os
from concurrent.futures import ThreadPoolExecutor
from datetime import timedelta
from typing import Dict, Optional, Union, Any, List, Iterable
from typing import Dict, Optional, Union, Any, List, Iterable, TYPE_CHECKING
__version__ = importlib.metadata.version("lancedb")
@@ -20,7 +20,7 @@ from .db import AsyncConnection, DBConnection, LanceDBConnection
from .remote import ClientConfig
from .remote.db import RemoteDBConnection
from .expr import Expr, col, lit, func
from .schema import blob, vector, BlobType
from .schema import blob, vector
from .job import AsyncJob, Job
from .functions import (
FunctionArtifactRequest as FunctionArtifactRequest,
@@ -49,6 +49,19 @@ from .namespace import (
)
if TYPE_CHECKING:
from lance.blob import BlobType as BlobType
def __getattr__(name: str):
if name == "BlobType":
from .schema import BlobType
globals()["BlobType"] = BlobType
return BlobType
raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
def _check_s3_bucket_with_dots(
uri: str, storage_options: Optional[Dict[str, str]]
) -> None:
+5 -3
View File
@@ -12,7 +12,7 @@ from typing import TYPE_CHECKING, Optional, Union
import pyarrow as pa
from .expr import Expr
from .schema import blob_v2_column_paths
from .schema import row_addressable_blob_v2_paths
from .types import BlobMode, QueryProjection, QueryProjectionSpec
if TYPE_CHECKING:
@@ -119,7 +119,7 @@ def blob_v2_projection_sources(
schema: pa.Schema,
projection: QueryProjection,
) -> dict[str, str]:
blob_columns = blob_v2_column_paths(schema)
blob_columns = row_addressable_blob_v2_paths(schema)
if not blob_columns:
return {}
columns = set(blob_columns)
@@ -140,7 +140,9 @@ def v2_projection_needs_row_id(
) -> bool:
if with_row_id:
return False
return projection_includes_blob_column(projection, blob_v2_column_paths(schema))
return projection_includes_blob_column(
projection, row_addressable_blob_v2_paths(schema)
)
def blob_auto_row_id_for_scan(
+38 -7
View File
@@ -16,6 +16,7 @@ from typing import (
Iterable,
List,
Literal,
Mapping,
Optional,
Union,
)
@@ -687,17 +688,35 @@ class DBConnection(EnforceOverrides):
"""
raise NotImplementedError("serialize is not supported for this connection type")
def create_function(self, definition: UdfDefinition) -> FunctionVersion:
def create_function(
self,
definition: UdfDefinition,
*,
secrets: Optional[Mapping[str, str]] = None,
) -> FunctionVersion:
"""Register a scalar Python UDF and wait for its immutable version.
``secrets`` must contain exactly the names declared by
``@udf(secrets=[...])``. Values are sent in the create request and
stored server-side in the private execution artifact; returned
Function and Job metadata contain only the declared names.
This is the blocking counterpart of :meth:`create_function_async`.
Local connections raise ``NotImplementedError``.
"""
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]:
def create_function_async(
self,
definition: UdfDefinition,
*,
secrets: Optional[Mapping[str, str]] = None,
) -> Job[FunctionVersion]:
"""Register a scalar Python UDF through the remote Function catalog.
``secrets`` must contain exactly the names declared by
``@udf(secrets=[...])``. Values are sent in the create request and
stored server-side in the private execution artifact; returned
Function and Job metadata contain only the declared names.
Submission returns a typed job. The immutable Function version becomes
available only when :meth:`Job.wait` succeeds. Local connections raise
``NotImplementedError``.
@@ -1405,8 +1424,13 @@ class LanceDBConnection(DBConnection):
return Job(self._conn.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[Mapping[str, str]] = None,
) -> Job[FunctionVersion]:
job = LOOP.run(self._conn.create_function_async(definition, secrets=secrets))
return Job(job)
@override
@@ -2225,17 +2249,24 @@ class AsyncConnection(object):
return AsyncJob(self._inner.job(job_id))
async def create_function_async(
self, definition: UdfDefinition
self,
definition: UdfDefinition,
*,
secrets: Optional[Mapping[str, str]] = None,
) -> AsyncJob[FunctionVersion]:
"""Register a scalar Python UDF through the remote Function catalog.
``secrets`` must contain exactly the names declared by
``@udf(secrets=[...])``. Values are sent in the create request and
stored server-side in the private execution artifact; returned
Function and Job metadata contain only the declared names.
The returned typed job resolves to the immutable Function version.
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()
definition._submission_json(secrets)
)
return _typed_job(inner, FunctionVersion.from_json)
+106 -6
View File
@@ -4,7 +4,7 @@
"""Canonical Function values exchanged with LanceDB Enterprise services.
These immutable models contain client/wire state only. Catalog persistence,
environment bake, and execution are owned by Sophon.
environment bake, secret resolution, and execution are owned by Sophon.
``RefreshColumnResult`` is also the backend-neutral result of a local
expression-backed refresh job.
"""
@@ -229,7 +229,7 @@ class PythonEnvironmentSpec(_RemoteValue):
class PythonRuntimeSpec(_RemoteValue):
"""Remote runtime definition with environment values.
"""Remote runtime definition with non-secret environment values.
V1 supports ``kind="python"``. Newer runtime kinds remain readable, while
their unknown payload fields are intentionally not retained by the client.
@@ -268,6 +268,7 @@ class FunctionVersion(_RemoteValue):
runtime: PythonRuntimeSpec
runtime_digest: str
environment_digest: str
required_secrets: tuple[str, ...] = ()
created_at: str
def __call__(self, **inputs: Any) -> FunctionApplication:
@@ -329,12 +330,17 @@ class FunctionVersion(_RemoteValue):
class FunctionRegistrationRequest(_RemoteValue):
"""Stable remote registration envelope produced by :func:`udf`."""
"""Stable remote registration envelope produced by :func:`udf`.
Only secret names are represented. Secret values are supplied separately
when the definition is submitted and are not part of this durable value.
"""
name: str
artifact: FunctionArtifactRequest
signature: FunctionSignature
runtime: PythonRuntimeSpec
required_secrets: tuple[str, ...] = ()
class FunctionVersionRef(_OpenRemoteValue):
@@ -479,6 +485,27 @@ class RefreshColumnResult(_RemoteValue):
_FUNCTION_NAME = re.compile(r"^[A-Za-z_][A-Za-z0-9_.-]*$")
_SECRET_NAME = re.compile(r"^[A-Za-z_][A-Za-z0-9_]*$")
# Keep this byte limit aligned with Sophon's MAX_FUNCTION_SECRET_VALUE_BYTES.
_MAX_FUNCTION_SECRET_VALUE_BYTES = 64 * 1024
_MAX_FUNCTION_SECRET_VALUES_BYTES = 512 * 1024
def _validate_secret_value(name: str, value: Any) -> str:
"""Validate one secret value before building the create request."""
if not isinstance(value, str):
raise TypeError(f"Function secret {name!r} value must be a string")
if not value:
raise ValueError(f"Function secret {name!r} value must be non-empty")
if "\0" in value:
raise ValueError(f"Function secret {name!r} value must not contain NUL")
value_bytes = len(value.encode("utf-8"))
if value_bytes > _MAX_FUNCTION_SECRET_VALUE_BYTES:
raise ValueError(
f"Function secret {name!r} value exceeds the "
f"{_MAX_FUNCTION_SECRET_VALUE_BYTES}-byte limit"
)
return value
_GRAMMAR_PRIMITIVES = (
@@ -909,6 +936,7 @@ class UdfDefinition:
output_schema: Optional[pa.DataType | pa.Field | pa.Schema],
pip: tuple[str, ...],
env: Mapping[str, str],
secrets: tuple[str, ...],
python_version: Optional[str],
conda: tuple[str, ...] = (),
conda_channels: tuple[str, ...] = (),
@@ -935,6 +963,17 @@ class UdfDefinition:
for key, value in environment.items()
):
raise TypeError("Function env keys and values must be strings")
required_secrets = tuple(sorted(set(secrets)))
invalid_secrets = [
secret for secret in required_secrets if not _SECRET_NAME.fullmatch(secret)
]
if invalid_secrets:
raise ValueError(f"invalid Function secret names: {invalid_secrets!r}")
overlap = set(environment) & set(required_secrets)
if overlap:
raise ValueError(
f"Function env and secret names must be disjoint: {sorted(overlap)!r}"
)
signature = _infer_signature(function, input_schema, output_schema)
source = _package_source(function)
digest = f"sha256:{hashlib.sha256(source).hexdigest()}"
@@ -963,14 +1002,65 @@ class UdfDefinition:
),
signature=signature,
runtime=runtime,
required_secrets=required_secrets,
)
functools.update_wrapper(self, function)
@property
def registration_request(self) -> FunctionRegistrationRequest:
"""The immutable request sent by ``create_function_async``."""
"""The immutable, value-free client model for a Function submission."""
return self._request
def _submission_json(self, secrets: Optional[Mapping[str, str]]) -> str:
"""Build one registration submission without retaining values on self."""
if secrets is None:
secret_values: Mapping[str, str] = {}
elif not isinstance(secrets, Mapping):
raise TypeError("Function secrets must be a mapping of names to strings")
else:
secret_values = secrets
if any(not isinstance(name, str) for name in secret_values):
raise TypeError("Function secret names must be strings")
expected = set(self._request.required_secrets)
provided = set(secret_values)
if provided != expected:
missing = sorted(expected - provided)
unexpected = sorted(provided - expected)
details = []
if missing:
details.append(f"missing: {missing!r}")
if unexpected:
details.append(f"unexpected: {unexpected!r}")
raise ValueError(
"Function secret values must exactly match the declared secrets ("
+ "; ".join(details)
+ ")"
)
canonical_values = {}
total_bytes = 0
for name in sorted(secret_values):
value = _validate_secret_value(name, secret_values[name])
total_bytes += len(value.encode("utf-8"))
if total_bytes > _MAX_FUNCTION_SECRET_VALUES_BYTES:
raise ValueError(
"Function secret values exceed the "
f"{_MAX_FUNCTION_SECRET_VALUES_BYTES}-byte request limit"
)
canonical_values[name] = value
submission = self._request._known_dict()
if canonical_values:
submission["secret_values"] = canonical_values
return json.dumps(
submission,
ensure_ascii=False,
allow_nan=False,
sort_keys=True,
separators=(",", ":"),
)
def __call__(self, *args, **kwargs):
return self._function(*args, **kwargs)
@@ -988,6 +1078,7 @@ def udf(
output_schema: Optional[pa.DataType | pa.Field | pa.Schema] = None,
pip: tuple[str, ...] | list[str] = (),
env: Optional[Mapping[str, str]] = None,
secrets: tuple[str, ...] | list[str] = (),
python_version: Optional[str] = None,
conda: tuple[str, ...] | list[str] = (),
conda_channels: tuple[str, ...] | list[str] = (),
@@ -1002,6 +1093,7 @@ def udf(
output_schema: Optional[pa.DataType | pa.Field | pa.Schema] = None,
pip: tuple[str, ...] | list[str] = (),
env: Optional[Mapping[str, str]] = None,
secrets: tuple[str, ...] | list[str] = (),
python_version: Optional[str] = None,
conda: tuple[str, ...] | list[str] = (),
conda_channels: tuple[str, ...] | list[str] = (),
@@ -1032,7 +1124,10 @@ def udf(
conda_channels : sequence of str, optional
Conda channels in priority order; requires ``conda``.
env : mapping of str to str, optional
Environment variables included in the Function definition.
Non-secret environment variables. Use ``secrets`` for credentials.
secrets : sequence of str, optional
Names of secrets required by the callable. Supply their values separately
to ``create_function`` or ``create_function_async``.
python_version : str, optional
Remote Python major/minor version. Defaults to the client version.
@@ -1054,11 +1149,15 @@ def udf(
Examples
--------
>>> from lancedb import udf
>>> @udf(pip=["numpy==2.2.0"])
>>> @udf(pip=["numpy==2.2.0"], secrets=["MODEL_TOKEN"])
... def score(value: float) -> float:
... return value * 2
>>> score(1.5)
3.0
>>> db.create_function( # doctest: +SKIP
... score, secrets={"MODEL_TOKEN": "user-secret-value"}
... )
"""
def decorate(target: Callable[..., Any]) -> UdfDefinition:
@@ -1069,6 +1168,7 @@ def udf(
output_schema=output_schema,
pip=tuple(pip),
env={} if env is None else env,
secrets=tuple(secrets),
python_version=python_version,
conda=tuple(conda),
conda_channels=tuple(conda_channels),
+10 -3
View File
@@ -7,7 +7,7 @@ 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, Mapping, Optional, Union
from urllib.parse import urlparse
import warnings
@@ -742,8 +742,15 @@ class RemoteDBConnection(DBConnection):
return Job(self._conn.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[Mapping[str, str]] = None,
) -> Job[FunctionVersion]:
return Job(
LOOP.run(self._conn.create_function_async(definition, secrets=secrets))
)
@override
def get_function(self, name: str, *, version: str) -> FunctionVersion:
+101 -34
View File
@@ -4,30 +4,34 @@
"""Schema helpers for Lance blob columns."""
import importlib
from typing import TYPE_CHECKING
import pyarrow as pa
import pyarrow.ipc
if TYPE_CHECKING:
from lance.blob import BlobType as BlobType
_BLOB_EXTENSION_NAME = "lance.blob.v2"
_BLOB_V1_KEY = "lance-encoding:blob"
_ARROW_EXT_NAME_KEY = "ARROW:extension:name"
_BLOB_V2_STORAGE_TYPE = pa.struct(
[
pa.field("data", pa.large_binary(), nullable=True),
pa.field("uri", pa.utf8(), nullable=True),
pa.field("position", pa.uint64(), nullable=True),
pa.field("size", pa.uint64(), nullable=True),
]
)
_resolved_blob_type = None
class BlobType(pa.ExtensionType):
"""PyArrow extension type for a Lance blob v2 column.
Queries return descriptors; call :meth:`~lancedb.table.Table.fetch_blob_files`
for lazy reads or :meth:`~lancedb.table.Table.fetch_blobs` for eager bytes.
"""
class _FallbackBlobType(pa.ExtensionType):
"""lance.blob.v2 extension type used when pylance is not installed."""
def __init__(self) -> None:
storage_type = pa.struct(
[
pa.field("data", pa.large_binary(), nullable=True),
pa.field("uri", pa.utf8(), nullable=True),
pa.field("position", pa.uint64(), nullable=True),
pa.field("size", pa.uint64(), nullable=True),
]
)
super().__init__(storage_type, _BLOB_EXTENSION_NAME)
pa.ExtensionType.__init__(self, _BLOB_V2_STORAGE_TYPE, _BLOB_EXTENSION_NAME)
def __arrow_ext_serialize__(self) -> bytes:
return b""
@@ -35,23 +39,16 @@ class BlobType(pa.ExtensionType):
@classmethod
def __arrow_ext_deserialize__(
cls, storage_type: pa.DataType, serialized: bytes
) -> "BlobType":
) -> "_FallbackBlobType":
return cls()
def __reduce__(self):
# Ensure pickle round-trips on older pyarrow (apache/arrow#35599).
return type(self).__arrow_ext_deserialize__, (
self.storage_type,
self.__arrow_ext_serialize__(),
)
try:
pa.register_extension_type(BlobType()) # type: ignore[arg-type]
except pa.ArrowKeyError:
pass
def _metadata_value(metadata: dict, key: str):
return metadata.get(key.encode()) or metadata.get(key)
@@ -92,43 +89,105 @@ def is_blob_like_field(field: pa.Field) -> bool:
return is_blob_v2_field(field) or _metadata_marks_legacy_blob(field.metadata or {})
def _collect_blob_paths(schema: pa.Schema, is_blob) -> list[str]:
paths: list[str] = []
def _collect_blob_paths(schema: pa.Schema, is_blob) -> list[tuple[str, bool]]:
"""Walk the schema and return (path, has_list_ancestor) for each blob field."""
paths: list[tuple[str, bool]] = []
def walk(fields, prefix: str) -> None:
def walk(fields, prefix: str, has_list_ancestor: bool) -> None:
for field in fields:
path = f"{prefix}.{field.name}" if prefix else field.name
if is_blob(field):
paths.append(path)
paths.append((path, has_list_ancestor))
elif pa.types.is_struct(field.type):
walk(field.type, path)
walk(field.type, path, has_list_ancestor)
elif (
pa.types.is_list(field.type)
or pa.types.is_large_list(field.type)
or pa.types.is_fixed_size_list(field.type)
):
walk([field.type.value_field], path)
walk([field.type.value_field], path, True)
walk(schema, "")
walk(schema, "", False)
return paths
def blob_column_paths(schema: pa.Schema) -> list[str]:
"""Dotted paths of blob-like columns (v2 extension or legacy metadata)."""
return _collect_blob_paths(schema, is_blob_like_field)
return [path for path, _ in _collect_blob_paths(schema, is_blob_like_field)]
def blob_v2_column_paths(schema: pa.Schema) -> list[str]:
return _collect_blob_paths(schema, is_blob_v2_field)
return [path for path, _ in _collect_blob_paths(schema, is_blob_v2_field)]
def row_addressable_blob_v2_paths(schema: pa.Schema) -> list[str]:
"""Blob v2 paths with one blob addressable by table row id.
``fetch_blobs`` and the descriptor row-id ride-along address one blob per
row, so a blob inside a list container has no row-id slot and no fetch
path. Those columns still store and query as raw descriptors.
"""
return [
path
for path, has_list_ancestor in _collect_blob_paths(schema, is_blob_v2_field)
if not has_list_ancestor
]
def schema_has_blob_field(schema: pa.Schema) -> bool:
return bool(blob_column_paths(schema))
def _deserialize_registered_type(extension_type: pa.ExtensionType) -> pa.DataType:
"""Return the type Arrow reconstructs for this extension name."""
schema = pa.schema([pa.field("value", extension_type)])
restored = pa.ipc.read_schema(schema.serialize())
return restored.field("value").type
def _resolve_blob_type():
"""Return the BlobType class this process should use.
pylance's class when it owns the lance.blob.v2 registry entry,
otherwise LanceDB's fallback. A different registered class is an error.
"""
global _resolved_blob_type
if _resolved_blob_type is not None:
return _resolved_blob_type
try:
blob_module = importlib.import_module("lance.blob")
except ModuleNotFoundError as err:
if err.name not in ("lance", "lance.blob"):
raise
else:
blob_type = getattr(blob_module, "BlobType", None)
if blob_type is not None:
registered_type = _deserialize_registered_type(blob_type())
if type(registered_type) is not blob_type:
registered_cls = type(registered_type)
raise ValueError(
"lance.blob.v2 is already registered by "
f"{registered_cls.__module__}.{registered_cls.__qualname__}"
)
_resolved_blob_type = blob_type
return blob_type
try:
pa.register_extension_type(_FallbackBlobType()) # type: ignore[arg-type]
except pa.ArrowKeyError as err:
raise ValueError(
"lance.blob.v2 is already registered by another extension class"
) from err
_resolved_blob_type = _FallbackBlobType
return _resolved_blob_type
def blob(name: str, nullable: bool = True) -> pa.Field:
"""Create a Lance blob v2 column field."""
return pa.field(name, BlobType(), nullable=nullable)
"""Create a Lance blob v2 column field.
When pylance is installed this is ``lance.blob.BlobType``.
"""
blob_type = _resolve_blob_type()
return pa.field(name, blob_type(), nullable=nullable)
def vector(dimension: int, value_type: pa.DataType = pa.float32()) -> pa.DataType:
@@ -155,3 +214,11 @@ def vector(dimension: int, value_type: pa.DataType = pa.float32()) -> pa.DataTyp
... ])
"""
return pa.list_(value_type, dimension)
def __getattr__(name: str):
if name == "BlobType":
blob_type = _resolve_blob_type()
globals()["BlobType"] = blob_type
return blob_type
raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
+234 -58
View File
@@ -104,7 +104,12 @@ from .util import (
value_to_sql,
)
from .index import lang_mapping
from .schema import blob_v2_column_paths, schema_has_blob_field
from .schema import (
blob_v2_column_paths,
is_blob_v2_field,
row_addressable_blob_v2_paths,
schema_has_blob_field,
)
def _should_push_down_query_table(
@@ -426,6 +431,7 @@ def _cast_to_target_schema(
def gen():
for batch in reader:
batch = _coerce_blob_write_columns(batch, reordered_schema)
# Table but not RecordBatch has cast.
cast_batches = (
pa.Table.from_batches([batch]).cast(reordered_schema).to_batches()
@@ -438,6 +444,166 @@ def _cast_to_target_schema(
return pa.RecordBatchReader.from_batches(reordered_schema, gen())
def _coerce_blob_write_columns(
batch: pa.RecordBatch, target_schema: pa.Schema
) -> pa.RecordBatch:
"""Materialize blob storage structs before the stream leaves Python.
merge_insert requires its source reader to already match the table's
physical schema. Unlike add and insert, it does not pass through
LanceDB's Rust blob coercion, so preserving binary input here would
reach Lance as binary and fail the schema check.
"""
columns = []
fields = []
changed = False
for field, column in zip(batch.schema, batch.columns):
target_field = target_schema.field(field.name)
coerced = _coerce_blob_value(column, target_field)
if coerced is not column:
column = coerced
field = pa.field(
field.name,
coerced.type,
field.nullable,
target_field.metadata,
)
changed = True
columns.append(column)
fields.append(field)
if not changed:
return batch
return pa.RecordBatch.from_arrays(
columns, schema=pa.schema(fields, metadata=batch.schema.metadata)
)
def _coerce_blob_value(column: pa.Array, target_field: pa.Field) -> pa.Array:
if is_blob_v2_field(target_field) and _can_coerce_to_blob(column.type):
return _coerce_value_to_blob(column, target_field)
target_type = target_field.type
if pa.types.is_struct(target_type) and pa.types.is_struct(column.type):
children = []
fields = []
changed = False
for source_field in column.type:
source_column = column.field(source_field.name)
nested_target = next(
(field for field in target_type if field.name == source_field.name),
None,
)
if nested_target is None:
children.append(source_column)
fields.append(source_field)
continue
coerced = _coerce_blob_value(source_column, nested_target)
if coerced is not source_column:
changed = True
child_array, child_type = _physical_array_and_type(coerced)
children.append(child_array)
fields.append(
pa.field(
source_field.name,
child_type,
source_field.nullable,
nested_target.metadata,
)
)
if not changed:
return column
return pa.StructArray.from_arrays(
children,
fields=fields,
mask=column.is_null() if column.null_count else None,
)
if _is_list_like(target_type) and _is_list_like(column.type):
return _coerce_blob_list_values(column, target_type.value_field)
return column
def _coerce_blob_list_values(
column: pa.Array, target_value_field: pa.Field
) -> pa.Array:
"""Coerce blob values inside a list column, preserving offsets and nulls.
Works on the raw child values window instead of ``pc.list_flatten`` because
flatten drops values spanned by null slots, which would misalign offsets.
"""
mask = column.is_null() if column.null_count else None
if pa.types.is_fixed_size_list(column.type):
list_size = column.type.list_size
values = column.values.slice(column.offset * list_size, len(column) * list_size)
coerced = _coerce_blob_value(values, target_value_field)
if coerced is values:
return column
physical_values, _ = _physical_array_and_type(coerced)
return pa.FixedSizeListArray.from_arrays(physical_values, list_size, mask=mask)
offsets = column.offsets
first_offset = offsets[0].as_py()
values = column.values.slice(
first_offset,
offsets[-1].as_py() - first_offset,
)
coerced = _coerce_blob_value(values, target_value_field)
if coerced is values:
return column
physical_values, _ = _physical_array_and_type(coerced)
if first_offset:
offsets = pc.subtract(offsets, pa.scalar(first_offset, offsets.type))
if pa.types.is_large_list(column.type):
return pa.LargeListArray.from_arrays(offsets, physical_values, mask=mask)
return pa.ListArray.from_arrays(offsets, physical_values, mask=mask)
def _coerce_value_to_blob(values: pa.Array, target_field: pa.Field) -> pa.Array:
if pa.types.is_null(values.type):
data = pa.nulls(len(values), type=pa.large_binary())
elif pa.types.is_large_binary(values.type):
data = values
else:
data = values.cast(pa.large_binary())
length = len(values)
storage_type = target_field.type
if isinstance(storage_type, pa.ExtensionType):
storage_type = storage_type.storage_type
storage_fields = list(storage_type)
children = []
for storage_field in storage_fields:
if storage_field.name == "data":
children.append(data)
else:
children.append(pa.nulls(length, type=storage_field.type))
storage = pa.StructArray.from_arrays(
children,
fields=storage_fields,
mask=values.is_null() if values.null_count else None,
)
if isinstance(target_field.type, pa.ExtensionType):
return pa.ExtensionArray.from_storage(target_field.type, storage)
return storage
def _physical_array_and_type(array: pa.Array) -> tuple[pa.Array, pa.DataType]:
if isinstance(array.type, pa.ExtensionType):
return array.storage, array.type.storage_type
return array, array.type
def _can_coerce_to_blob(data_type: pa.DataType) -> bool:
return _is_binary_like(data_type) or pa.types.is_null(data_type)
def _is_binary_like(data_type: pa.DataType) -> bool:
return (
pa.types.is_binary(data_type)
or pa.types.is_large_binary(data_type)
or pa.types.is_binary_view(data_type)
)
def _field_extension_name(field: pa.Field) -> Optional[str]:
extension_name = getattr(field.type, "extension_name", None)
if extension_name is not None:
@@ -464,63 +630,71 @@ def _align_field_types(
target_field = next((f for f in target_fields if f.name == field.name), None)
if target_field is None:
raise ValueError(f"Field '{field.name}' not found in target schema")
# 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
# 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"
):
new_fields.append(field)
continue
if pa.types.is_struct(target_field.type):
if pa.types.is_struct(field.type):
new_type = pa.struct(
_align_field_types(
field.type.fields,
target_field.type.fields,
)
new_fields.append(_align_field(field, target_field))
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_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
# 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 pa.types.is_struct(target_field.type):
if pa.types.is_struct(field.type):
new_type = pa.struct(
_align_field_types(
field.type.fields,
target_field.type.fields,
)
else:
new_type = target_field.type
elif pa.types.is_list(target_field.type):
if _is_list_like(field.type):
new_type = pa.list_(
_align_field_types(
[field.type.value_field],
[target_field.type.value_field],
)[0]
)
else:
new_type = target_field.type
elif pa.types.is_large_list(target_field.type):
if _is_list_like(field.type):
new_type = pa.large_list(
_align_field_types(
[field.type.value_field],
[target_field.type.value_field],
)[0]
)
else:
new_type = target_field.type
elif pa.types.is_fixed_size_list(target_field.type):
if _is_list_like(field.type):
new_type = pa.list_(
_align_field_types(
[field.type.value_field],
[target_field.type.value_field],
)[0],
target_field.type.list_size,
)
else:
new_type = target_field.type
)
else:
new_type = target_field.type
new_fields.append(
pa.field(field.name, new_type, field.nullable, target_field.metadata)
)
return new_fields
elif pa.types.is_list(target_field.type):
if _is_list_like(field.type):
new_type = pa.list_(
_align_list_value_field(
field.type.value_field, target_field.type.value_field
)
)
else:
new_type = target_field.type
elif pa.types.is_large_list(target_field.type):
if _is_list_like(field.type):
new_type = pa.large_list(
_align_list_value_field(
field.type.value_field, target_field.type.value_field
)
)
else:
new_type = target_field.type
elif pa.types.is_fixed_size_list(target_field.type):
if _is_list_like(field.type):
new_type = pa.list_(
_align_list_value_field(
field.type.value_field, target_field.type.value_field
),
target_field.type.list_size,
)
else:
new_type = target_field.type
else:
new_type = target_field.type
return pa.field(field.name, new_type, field.nullable, target_field.metadata)
def _infer_subschema(
@@ -589,7 +763,7 @@ def sanitize_create_table(
schema = data.schema
else:
if schema is not None:
data = pa.Table.from_pylist([], schema)
data = pa.Table.from_batches([], schema=schema)
if schema is None:
if data is None:
raise ValueError("Either data or schema must be provided")
@@ -2698,7 +2872,7 @@ class LanceTable(Table):
arrow_tbl = self.to_arrow()
if blob_mode == "descriptions":
arrow_tbl = strip_auto_row_ids(
arrow_tbl, blob_v2_column_paths(self.schema)
arrow_tbl, row_addressable_blob_v2_paths(self.schema)
)
return arrow_tbl.to_pandas(**kwargs)
@@ -5102,7 +5276,9 @@ class AsyncTable:
if blob_mode == "descriptions" or not schema_has_blob_field(schema):
arrow_tbl = await self.to_arrow()
if blob_mode == "descriptions":
arrow_tbl = strip_auto_row_ids(arrow_tbl, blob_v2_column_paths(schema))
arrow_tbl = strip_auto_row_ids(
arrow_tbl, row_addressable_blob_v2_paths(schema)
)
return arrow_tbl.to_pandas(**kwargs)
if blob_mode == "lazy" and get_uri_scheme(await self.uri()) == "memory":
+480
View File
@@ -2,10 +2,15 @@
# SPDX-FileCopyrightText: Copyright The LanceDB Authors
import io
import subprocess
import sys
import textwrap
import lance
import pyarrow as pa
import pyarrow.compute as pc
import pytest
from lance.blob import BlobType as LanceBlobType
import lancedb
from lancedb._blob import (
@@ -18,6 +23,20 @@ from lancedb.index import FTS
from lancedb.schema import blob_column_paths, blob_v2_column_paths
_HIDE_LANCE_BLOB = """\
import importlib.abc
import sys
class _MissingLanceBlob(importlib.abc.MetaPathFinder):
def find_spec(self, fullname, path, target=None):
if fullname == "lance.blob" or fullname.startswith("lance.blob."):
raise ModuleNotFoundError(fullname, name="lance.blob")
sys.modules.pop("lance.blob", None)
sys.meta_path.insert(0, _MissingLanceBlob())
"""
def _blob_table(name, rows):
db = lancedb.connect("memory:///")
schema = pa.schema([pa.field("id", pa.int64()), lancedb.blob("image")])
@@ -51,6 +70,181 @@ def test_blob_factory_declares_v2_field():
field = lancedb.blob("image")
assert isinstance(field.type, pa.ExtensionType)
assert field.type.extension_name == "lance.blob.v2"
assert lancedb.BlobType is LanceBlobType
assert type(field.type) is LanceBlobType
def test_blob_type_works_without_pylance():
script = _HIDE_LANCE_BLOB + textwrap.dedent(
"""\
import lancedb
import pyarrow as pa
field = lancedb.blob("image")
if not isinstance(field.type, pa.ExtensionType):
raise SystemExit("expected an extension type")
if field.type.extension_name != "lance.blob.v2":
raise SystemExit(field.type.extension_name)
if lancedb.BlobType is not type(field.type):
raise SystemExit("BlobType is not the field type class")
if lancedb.BlobType.__module__ != "lancedb.schema":
raise SystemExit(lancedb.BlobType.__module__)
db = lancedb.connect("memory:///")
table = db.create_table(
"images",
schema=pa.schema([pa.field("id", pa.int64()), field]),
)
table.add([{"id": 1, "image": b"hello"}])
result = (
table.merge_insert("id")
.when_matched_update_all()
.when_not_matched_insert_all()
.execute([{"id": 1, "image": b"updated"}, {"id": 2, "image": b"inserted"}])
)
if result.num_updated_rows != 1 or result.num_inserted_rows != 1:
raise SystemExit(
f"merge_insert rows updated={result.num_updated_rows} "
f"inserted={result.num_inserted_rows}"
)
"""
)
result = subprocess.run(
[sys.executable, "-c", script],
capture_output=True,
text=True,
check=False,
)
assert result.returncode == 0, result.stderr
def test_blob_resolves_pylance_type_without_eager_import():
script = textwrap.dedent(
"""\
import sys
import lancedb
if "lance.blob" in sys.modules:
raise SystemExit("import lancedb imported lance.blob")
field = lancedb.blob("image")
from lance.blob import BlobType
if type(field.type) is not BlobType:
raise SystemExit(f"{type(field.type)} is not {BlobType}")
import lance
image = lance.blob_array([b"x"])
if type(image.type) is not BlobType:
raise SystemExit("blob_array used a different class")
if type(image.type) is not type(field.type):
raise SystemExit("field and array classes differ")
"""
)
result = subprocess.run(
[sys.executable, "-c", script],
capture_output=True,
text=True,
check=False,
)
assert result.returncode == 0, result.stderr
def test_blob_fallback_fails_if_name_already_registered():
script = _HIDE_LANCE_BLOB + textwrap.dedent(
"""\
import pyarrow as pa
class OtherBlobType(pa.ExtensionType):
def __init__(self):
super().__init__(
pa.struct([pa.field("data", pa.large_binary())]),
"lance.blob.v2",
)
def __arrow_ext_serialize__(self):
return b""
@classmethod
def __arrow_ext_deserialize__(cls, storage_type, serialized):
return cls()
pa.register_extension_type(OtherBlobType())
import lancedb
try:
lancedb.blob("image")
except ValueError as err:
if "already registered" not in str(err):
raise SystemExit(err)
else:
raise SystemExit("expected ValueError")
"""
)
result = subprocess.run(
[sys.executable, "-c", script],
capture_output=True,
text=True,
check=False,
)
assert result.returncode == 0, result.stderr
def test_blob_type_rejects_competing_registration_with_pylance():
script = textwrap.dedent(
"""\
import pyarrow as pa
import pyarrow.ipc
class OtherBlobType(pa.ExtensionType):
def __init__(self):
super().__init__(
pa.struct(
[
pa.field("data", pa.large_binary()),
pa.field("uri", pa.utf8()),
pa.field("position", pa.uint64()),
pa.field("size", pa.uint64()),
]
),
"lance.blob.v2",
)
def __arrow_ext_serialize__(self):
return b""
@classmethod
def __arrow_ext_deserialize__(cls, storage_type, serialized):
return cls()
pa.register_extension_type(OtherBlobType())
from lance.blob import BlobType
if BlobType is OtherBlobType:
raise SystemExit("pylance BlobType was replaced")
schema = pa.schema([pa.field("value", BlobType())])
restored = pa.ipc.read_schema(schema.serialize())
if type(restored.field("value").type) is not OtherBlobType:
raise SystemExit(type(restored.field("value").type))
import lancedb
try:
lancedb.blob("image")
except ValueError as err:
if "__main__.OtherBlobType" not in str(err):
raise SystemExit(err)
else:
raise SystemExit("expected ValueError")
"""
)
result = subprocess.run(
[sys.executable, "-c", script],
capture_output=True,
text=True,
check=False,
)
assert result.returncode == 0, result.stderr
def test_blob_v2_column_paths_include_list_children():
@@ -203,6 +397,292 @@ def test_fetch_blobs_round_trip():
assert [blobs[0].as_py(), blobs[1].as_py()] == [b"alpha", b"beta"]
def test_merge_insert_writes_python_bytes():
table = _blob_table("merge_bytes", [{"id": 1, "image": b"before"}])
result = (
table.merge_insert("id")
.when_matched_update_all()
.when_not_matched_insert_all()
.execute([{"id": 1, "image": b"updated"}, {"id": 2, "image": b"inserted"}])
)
assert result.num_updated_rows == 1
assert result.num_inserted_rows == 1
by_id = _row_ids_by_id(table)
blobs = table.fetch_blobs("image", [by_id[1], by_id[2]])
assert blobs.to_pylist() == [b"updated", b"inserted"]
def test_merge_insert_bytes_after_reopen_without_touching_blob_type(tmp_path):
db = lancedb.connect(tmp_path)
schema = pa.schema([pa.field("id", pa.int64()), lancedb.blob("image")])
table = db.create_table("images", schema=schema)
table.add([{"id": 1, "image": b"hello"}])
script = textwrap.dedent(
f"""\
import lancedb
db = lancedb.connect({str(tmp_path)!r})
table = db.open_table("images")
image_type = table.schema.field("image").type
if type(image_type).__name__ != "StructType":
raise SystemExit(f"expected StructType, got {{type(image_type)}}")
result = (
table.merge_insert("id")
.when_matched_update_all()
.when_not_matched_insert_all()
.execute(
[{{"id": 1, "image": b"updated"}}, {{"id": 2, "image": b"inserted"}}]
)
)
if result.num_updated_rows != 1 or result.num_inserted_rows != 1:
raise SystemExit(
f"rows updated={{result.num_updated_rows}} "
f"inserted={{result.num_inserted_rows}}"
)
hits = table.search().with_row_id(True).limit(10).to_arrow()
by_id = dict(zip(hits["id"].to_pylist(), hits["_rowid"].to_pylist()))
blobs = table.fetch_blobs("image", [by_id[1], by_id[2]])
if blobs.to_pylist() != [b"updated", b"inserted"]:
raise SystemExit(blobs.to_pylist())
"""
)
result = subprocess.run(
[sys.executable, "-c", script],
capture_output=True,
text=True,
check=False,
)
assert result.returncode == 0, result.stderr
def test_merge_insert_bytes_after_reopen_without_pylance(tmp_path):
db = lancedb.connect(tmp_path)
schema = pa.schema([pa.field("id", pa.int64()), lancedb.blob("image")])
table = db.create_table("images", schema=schema)
table.add([{"id": 1, "image": b"hello"}])
script = _HIDE_LANCE_BLOB + textwrap.dedent(
f"""\
import lancedb
db = lancedb.connect({str(tmp_path)!r})
table = db.open_table("images")
image_type = table.schema.field("image").type
if type(image_type).__name__ != "StructType":
raise SystemExit(f"expected StructType, got {{type(image_type)}}")
result = (
table.merge_insert("id")
.when_matched_update_all()
.when_not_matched_insert_all()
.execute(
[{{"id": 1, "image": b"updated"}}, {{"id": 2, "image": b"inserted"}}]
)
)
if result.num_updated_rows != 1 or result.num_inserted_rows != 1:
raise SystemExit(
f"rows updated={{result.num_updated_rows}} "
f"inserted={{result.num_inserted_rows}}"
)
hits = table.search().with_row_id(True).limit(10).to_arrow()
by_id = dict(zip(hits["id"].to_pylist(), hits["_rowid"].to_pylist()))
blobs = table.fetch_blobs("image", [by_id[1], by_id[2]])
if blobs.to_pylist() != [b"updated", b"inserted"]:
raise SystemExit(blobs.to_pylist())
"""
)
result = subprocess.run(
[sys.executable, "-c", script],
capture_output=True,
text=True,
check=False,
)
assert result.returncode == 0, result.stderr
def test_merge_insert_blob_array_into_reopened_unregistered_table(tmp_path):
db = lancedb.connect(tmp_path)
schema = pa.schema([pa.field("id", pa.int64()), lancedb.blob("image")])
table = db.create_table("images", schema=schema)
table.add([{"id": 1, "image": b"before"}])
script = textwrap.dedent(
f"""\
import pyarrow as pa
import lancedb
db = lancedb.connect({str(tmp_path)!r})
table = db.open_table("images")
image_type = table.schema.field("image").type
if type(image_type).__name__ != "StructType":
raise SystemExit(
f"expected StructType before lance import, got {{type(image_type)}}"
)
import lance
updates = pa.Table.from_arrays(
[
pa.array([1, 2], type=pa.int64()),
lance.blob_array([b"updated", b"inserted"]),
],
names=["id", "image"],
)
result = (
table.merge_insert("id")
.when_matched_update_all()
.when_not_matched_insert_all()
.execute(updates)
)
if result.num_updated_rows != 1 or result.num_inserted_rows != 1:
raise SystemExit(
f"rows updated={{result.num_updated_rows}} "
f"inserted={{result.num_inserted_rows}}"
)
hits = table.search().with_row_id(True).limit(10).to_arrow()
by_id = dict(zip(hits["id"].to_pylist(), hits["_rowid"].to_pylist()))
blobs = table.fetch_blobs("image", [by_id[1], by_id[2]])
if blobs.to_pylist() != [b"updated", b"inserted"]:
raise SystemExit(blobs.to_pylist())
"""
)
result = subprocess.run(
[sys.executable, "-c", script],
capture_output=True,
text=True,
check=False,
)
assert result.returncode == 0, result.stderr
def test_add_all_null_blob_column():
db = lancedb.connect("memory:///")
schema = pa.schema([pa.field("id", pa.int64()), lancedb.blob("image")])
table = db.create_table("all_null", schema=schema)
table.add([{"id": 1, "image": None}, {"id": 2, "image": None}])
by_id = _row_ids_by_id(table)
blobs = table.fetch_blobs("image", [by_id[1], by_id[2]])
assert blobs.to_pylist() == [None, None]
def test_create_table_nested_blob_schema_without_rows():
db = lancedb.connect("memory:///")
schema = pa.schema(
[
pa.field("id", pa.int64()),
pa.field("info", pa.struct([lancedb.blob("blob")])),
pa.field("images", pa.list_(lancedb.blob("image"))),
]
)
table = db.create_table("nested_empty", schema=schema)
assert table.count_rows() == 0
def test_merge_insert_nested_blob_dicts():
db = lancedb.connect("memory:///")
info = pa.StructArray.from_arrays(
[
pa.array(["first"], type=pa.string()),
_blob_array("blob", [b"before"]),
],
names=["name", "blob"],
)
data = pa.Table.from_arrays(
[pa.array([1], type=pa.int64()), info],
names=["id", "info"],
)
table = db.create_table("nested_merge", data=data)
result = (
table.merge_insert("id")
.when_matched_update_all()
.execute([{"id": 1, "info": {"name": "first", "blob": b"after"}}])
)
assert result.num_updated_rows == 1
by_id = _row_ids_by_id(table)
blobs = table.fetch_blobs("info.blob", [by_id[1]])
assert blobs.to_pylist() == [b"after"]
def _list_blob_table(name):
db = lancedb.connect("memory:///")
blob_field = lancedb.blob("image")
images = pa.ListArray.from_arrays(
pa.array([0, 1], type=pa.int32()), _blob_array("image", [b"before"])
)
data = pa.Table.from_arrays(
[pa.array([1], type=pa.int64()), images],
schema=pa.schema(
[pa.field("id", pa.int64()), pa.field("images", pa.list_(blob_field))]
),
)
return db.create_table(name, data=data)
def test_merge_insert_list_blob_dicts():
table = _list_blob_table("list_merge")
result = (
table.merge_insert("id")
.when_matched_update_all()
.when_not_matched_insert_all()
.execute([{"id": 1, "images": [b"one", b"two"]}, {"id": 2, "images": None}])
)
assert result.num_updated_rows == 1
assert result.num_inserted_rows == 1
hits = table.search().limit(10).to_arrow()
sizes = {
row["id"]: None if row["images"] is None else [d["size"] for d in row["images"]]
for row in hits.to_pylist()
}
assert sizes == {1: [3, 3], 2: None}
def test_list_blob_column_queries_as_raw_descriptors():
table = _list_blob_table("list_query")
hits = table.search().limit(10).to_arrow()
element = hits.schema.field("images").type.value_type
assert pa.types.is_struct(element)
assert "_lance_row_id" not in element.names
with pytest.raises(ValueError, match="expected struct before segment"):
table.fetch_blobs("images.image", [0])
def test_row_addressable_paths_exclude_list_children():
from lancedb.schema import row_addressable_blob_v2_paths
schema = pa.schema(
[
pa.field("id", pa.int64()),
pa.field("info", pa.struct([lancedb.blob("blob")])),
pa.field("images", pa.list_(lancedb.blob("image"))),
]
)
assert blob_v2_column_paths(schema) == ["info.blob", "images.image"]
assert row_addressable_blob_v2_paths(schema) == ["info.blob"]
def test_merge_insert_writes_pylance_blob_array():
table = _blob_table("merge_pylance", [{"id": 1, "image": b"before"}])
image = lance.blob_array([b"updated", b"inserted"])
assert type(image.type) is LanceBlobType
assert type(image.type) is type(lancedb.BlobType())
updates = pa.Table.from_arrays(
[pa.array([1, 2], type=pa.int64()), image], names=["id", "image"]
)
result = (
table.merge_insert("id")
.when_matched_update_all()
.when_not_matched_insert_all()
.execute(updates)
)
assert result.num_updated_rows == 1
assert result.num_inserted_rows == 1
by_id = _row_ids_by_id(table)
blobs = table.fetch_blobs("image", [by_id[1], by_id[2]])
assert blobs.to_pylist() == [b"updated", b"inserted"]
def test_fetch_blobs_accepts_query_result():
table = _blob_table("from_result", [{"id": 1, "image": b"gamma"}])
hits = table.search().limit(10).to_arrow()
@@ -37,6 +37,21 @@ def job_result(name: str) -> dict:
return json.loads(fixture(name))["result"]
def assert_no_secret_values(value):
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()
@@ -94,6 +109,7 @@ def test_function_version_identity_is_immutable_and_exact():
version = FunctionVersion.from_json(json.dumps(value))
assert version.name == "embed"
assert version.version == "fv_01K3EXACT"
assert version.required_secrets == ("HF_TOKEN",)
with pytest.raises((TypeError, ValueError)):
version.version = "fv_changed"
@@ -276,6 +292,15 @@ def test_refresh_result_rejects_non_u64_values(field):
RefreshColumnResult.from_json(json.dumps(value))
def test_canonical_client_values_contain_secret_names_only():
version = FunctionVersion.from_json(
json.dumps(job_result("remote_function_job.json"))
)
canonical = json.loads(version.to_canonical_json())
assert canonical["required_secrets"] == ["HF_TOKEN"]
assert_no_secret_values(canonical)
class _FunctionDeclarationInner:
def __init__(self):
self.calls = []
@@ -19,7 +19,13 @@ import pyarrow as pa
import pytest
import lancedb
from lancedb.functions import UdfDefinition, udf
from lancedb.functions import (
_MAX_FUNCTION_SECRET_VALUE_BYTES,
_MAX_FUNCTION_SECRET_VALUES_BYTES,
FunctionRegistrationRequest,
UdfDefinition,
udf,
)
THRESHOLD = 20
_CACHE = None
@@ -39,12 +45,28 @@ FIXTURES = (
@udf(
pip=["numpy>=2"],
env={"MODE": "test"},
secrets=["API_TOKEN"],
python_version="3.12",
)
def normalize_score(value: float) -> float:
return value / 100.0
def _assert_no_secret_values(value):
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_scalar_udf_matches_shared_registration_golden_and_remains_callable():
assert isinstance(normalize_score, UdfDefinition)
assert normalize_score(25.0) == 0.25
@@ -59,6 +81,8 @@ def test_scalar_udf_matches_shared_registration_golden_and_remains_callable():
"kind": "scalar_to_arrow_batch",
"version": 1,
}
assert request["required_secrets"] == ["API_TOKEN"]
_assert_no_secret_values(request)
def _run_packaged(definition, *args):
@@ -372,6 +396,7 @@ def test_udf_recursion_versus_a_rebound_module_name(tmp_path):
output_schema=None,
pip=(),
env={},
secrets=(),
python_version=None,
)
with pytest.raises(ValueError, match="binds that name to another value"):
@@ -526,13 +551,54 @@ def test_annotation_and_explicit_schema_validation_fail_closed():
return value
def test_secret_names_are_canonical_and_disjoint_from_environment():
@udf(secrets=["Z_TOKEN", "A_TOKEN", "Z_TOKEN"])
def canonical_secrets(value: int) -> int:
return value
assert canonical_secrets.registration_request.required_secrets == (
"A_TOKEN",
"Z_TOKEN",
)
with pytest.raises(ValueError, match="must be disjoint"):
@udf(env={"TOKEN": "plaintext"}, secrets=["TOKEN"])
def overlapping(value: int) -> int:
return value
def test_declared_secret_api_still_requires_explicit_create_values():
@udf(secrets=["API_TOKEN"])
def declared_secret(value: int) -> int:
return value
with pytest.raises(ValueError, match="missing"):
declared_secret._submission_json(None)
submission = json.loads(
declared_secret._submission_json({"API_TOKEN": "explicit-secret"})
)
assert submission["required_secrets"] == ["API_TOKEN"]
assert submission["secret_values"] == {"API_TOKEN": "explicit-secret"}
def test_no_secrets_preserve_canonical_registration_shape():
@udf
def no_secrets(value: int) -> int:
return value
canonical = json.loads(no_secrets.registration_request.to_canonical_json())
assert "required_secrets" not in canonical
assert json.loads(no_secrets._submission_json(None)) == canonical
def test_local_function_catalog_operations_are_not_supported(tmp_path):
db = lancedb.connect(tmp_path)
message = "Function catalog operations are not supported by this database"
with pytest.raises(NotImplementedError, match=message):
db.create_function(normalize_score)
db.create_function(normalize_score, secrets={"API_TOKEN": "value"})
with pytest.raises(NotImplementedError, match=message):
db.create_function_async(normalize_score)
db.create_function_async(normalize_score, secrets={"API_TOKEN": "value"})
with pytest.raises(NotImplementedError, match=message):
db.get_function("normalize_score", version="fv_exact")
@@ -562,6 +628,7 @@ def _mock_remote_function_catalog():
"runtime": body["runtime"],
"runtime_digest": "sha256:runtime",
"environment_digest": "sha256:environment",
"required_secrets": body.get("required_secrets", []),
"created_at": "2026-08-21T00:00:00Z",
}
response = {"job_id": "job-register"}
@@ -608,7 +675,9 @@ def test_remote_registration_job_and_exact_version_reopen_round_trip():
host_override=host,
client_config={"retry_config": {"retries": 0}},
)
registration = db.create_function_async(normalize_score)
registration = db.create_function_async(
normalize_score, secrets={"API_TOKEN": "secret-value"}
)
assert registration.id == "job-register"
created = registration.wait()
reopened = db.get_function("normalize_score", version=created.version)
@@ -617,9 +686,18 @@ def test_remote_registration_job_and_exact_version_reopen_round_trip():
assert reopened.name == "normalize_score"
assert reopened.version == "fv_exact"
create_request = state["requests"][0][1]
assert create_request == json.loads(
expected = json.loads(normalize_score.registration_request.to_canonical_json())
expected["secret_values"] = {"API_TOKEN": "secret-value"}
assert create_request == expected
durable_request = FunctionRegistrationRequest.from_json(json.dumps(create_request))
assert not hasattr(durable_request, "secret_values")
assert "secret_values" not in json.loads(durable_request.to_canonical_json())
assert "secret_values" not in json.loads(
normalize_score.registration_request.to_canonical_json()
)
assert "secret-value" not in repr(normalize_score)
assert "secret-value" not in repr(normalize_score.registration_request)
assert not hasattr(created, "secret_values")
def test_blocking_remote_registration_returns_function_version():
@@ -630,7 +708,9 @@ def test_blocking_remote_registration_returns_function_version():
host_override=host,
client_config={"retry_config": {"retries": 0}},
)
created = db.create_function(normalize_score)
created = db.create_function(
normalize_score, secrets={"API_TOKEN": "blocking-secret"}
)
assert created.name == "normalize_score"
assert created.version == "fv_exact"
@@ -638,3 +718,121 @@ def test_blocking_remote_registration_returns_function_version():
"/v1/functions/create",
"/v1/jobs/describe",
]
assert state["requests"][0][1]["secret_values"] == {"API_TOKEN": "blocking-secret"}
@pytest.mark.parametrize(
("secret_values", "error_type", "message"),
[
(None, ValueError, "missing"),
({}, ValueError, "missing"),
({"OTHER": "value"}, ValueError, "missing.*unexpected"),
({"API_TOKEN": ""}, ValueError, "non-empty"),
({"API_TOKEN": "bad\0value"}, ValueError, "NUL"),
({"API_TOKEN": 123}, TypeError, "must be a string"),
([("API_TOKEN", "value")], TypeError, "must be a mapping"),
],
)
def test_secret_values_are_validated_before_remote_request(
secret_values, error_type, message
):
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}},
)
with pytest.raises(error_type, match=message):
db.create_function_async(normalize_score, secrets=secret_values)
assert state["requests"] == []
@pytest.mark.parametrize(
"value",
[
"x" * _MAX_FUNCTION_SECRET_VALUE_BYTES,
"é" * (_MAX_FUNCTION_SECRET_VALUE_BYTES // len("é".encode("utf-8"))),
],
ids=["ascii", "multibyte"],
)
def test_secret_value_accepts_exact_utf8_byte_limit(value):
submission = json.loads(normalize_score._submission_json({"API_TOKEN": value}))
assert submission["secret_values"]["API_TOKEN"] == value
assert len(value.encode("utf-8")) == _MAX_FUNCTION_SECRET_VALUE_BYTES
@pytest.mark.parametrize(
"value",
[
"x" * (_MAX_FUNCTION_SECRET_VALUE_BYTES + 1),
"é" * (_MAX_FUNCTION_SECRET_VALUE_BYTES // len("é".encode("utf-8")) + 1),
],
ids=["ascii", "multibyte"],
)
def test_secret_value_rejects_over_utf8_byte_limit_before_json_construction(
monkeypatch, value
):
def fail_if_json_construction_starts(self):
pytest.fail("oversized secret reached JSON construction")
monkeypatch.setattr(
FunctionRegistrationRequest, "_known_dict", fail_if_json_construction_starts
)
with pytest.raises(ValueError, match=r"exceeds the 65536-byte limit"):
normalize_score._submission_json({"API_TOKEN": value})
def test_secret_values_accept_exact_aggregate_utf8_byte_limit(monkeypatch):
names = tuple(f"SECRET_{index}" for index in range(8))
value = "é" * (_MAX_FUNCTION_SECRET_VALUE_BYTES // len("é".encode("utf-8")))
values = {name: value for name in names}
monkeypatch.setattr(
normalize_score,
"_request",
normalize_score._request._copy(update={"required_secrets": names}),
)
submission = json.loads(normalize_score._submission_json(values))
assert submission["secret_values"] == values
assert sum(len(item.encode("utf-8")) for item in values.values()) == (
_MAX_FUNCTION_SECRET_VALUES_BYTES
)
def test_secret_values_reject_aggregate_over_limit_before_construction(monkeypatch):
names = tuple(f"SECRET_{index}" for index in range(9))
values = {name: "x" * _MAX_FUNCTION_SECRET_VALUE_BYTES for name in names}
monkeypatch.setattr(
normalize_score,
"_request",
normalize_score._request._copy(update={"required_secrets": names}),
)
def fail_if_json_construction_starts(self):
pytest.fail("oversized aggregate reached JSON construction")
monkeypatch.setattr(
FunctionRegistrationRequest, "_known_dict", fail_if_json_construction_starts
)
with pytest.raises(ValueError, match=r"exceed.*524288-byte request limit"):
normalize_score._submission_json(values)
@pytest.mark.asyncio
async def test_async_remote_registration_submits_secret_values_only_once():
with _mock_remote_function_catalog() as (host, state):
db = await lancedb.connect_async(
"db://dev",
api_key="fake",
host_override=host,
client_config={"retry_config": {"retries": 0}},
)
registration = await db.create_function_async(
normalize_score, secrets={"API_TOKEN": "async-secret"}
)
created = await registration.wait()
assert state["requests"][0][1]["secret_values"] == {"API_TOKEN": "async-secret"}
assert not hasattr(created, "secret_values")
+160
View File
@@ -7,6 +7,7 @@ import pathlib
from typing import Optional
import lance
from lance.blob import BlobType as LanceBlobType
from lancedb.conftest import MockTextEmbeddingFunction
from lancedb.embeddings.base import EmbeddingFunctionConfig
from lancedb.embeddings.registry import EmbeddingFunctionRegistry
@@ -907,6 +908,165 @@ def test_cast_to_target_schema():
assert output == expected
def test_cast_to_target_schema_coerces_binary_to_blob_v2():
data = pa.table({"image": pa.array([b"hello", None], type=pa.binary())})
target = pa.schema([lancedb.blob("image")])
output = _cast_to_target_schema(data.to_reader(), target).read_all()
image = output["image"].chunk(0)
assert type(image.type) is lancedb.BlobType
assert image.storage.to_pylist() == [
{"data": b"hello", "uri": None, "position": None, "size": None},
None,
]
def test_cast_to_target_schema_coerces_binary_to_metadata_blob_struct():
storage = lancedb.blob("image").type.storage_type
target = pa.schema(
[
pa.field(
"image",
storage,
metadata={
b"ARROW:extension:name": b"lance.blob.v2",
b"ARROW:extension:metadata": b"",
},
)
]
)
data = pa.table({"image": pa.array([b"hello", None], type=pa.binary())})
output = _cast_to_target_schema(data.to_reader(), target).read_all()
image = output["image"].chunk(0)
assert not isinstance(image.type, pa.ExtensionType)
assert image.to_pylist() == [
{"data": b"hello", "uri": None, "position": None, "size": None},
None,
]
def test_cast_to_target_schema_coerces_nested_binary_blob():
data = pa.table(
{
"info": pa.array(
[{"blob": b"hello"}, {"blob": None}],
type=pa.struct([pa.field("blob", pa.binary())]),
)
}
)
target = pa.schema([pa.field("info", pa.struct([lancedb.blob("blob")]))])
output = _cast_to_target_schema(data.to_reader(), target).read_all()
blob = output["info"].chunk(0).field("blob")
assert type(blob.type) is lancedb.BlobType
assert blob.storage.to_pylist() == [
{"data": b"hello", "uri": None, "position": None, "size": None},
None,
]
def test_cast_to_target_schema_coerces_list_binary_blob_with_inferred_child_name():
data = pa.table(
{"images": pa.array([[b"a", b"b"], None], type=pa.list_(pa.binary()))}
)
target = pa.schema([pa.field("images", pa.list_(lancedb.blob("image")))])
output = _cast_to_target_schema(data.to_reader(), target).read_all()
images = output["images"].chunk(0)
assert images.type.value_field.name == "image"
assert type(images.type.value_type) is lancedb.BlobType
assert images.to_pylist()[1] is None
assert images.values.storage.to_pylist() == [
{"data": b"a", "uri": None, "position": None, "size": None},
{"data": b"b", "uri": None, "position": None, "size": None},
]
def test_list_blob_coercion_preserves_null_slots_with_nonzero_extent():
child = pa.field("image", pa.binary())
source = pa.ListArray.from_arrays(
pa.array([0, 2, 4], type=pa.int32()),
pa.array([b"a", b"b", b"dead", b"beef"], type=pa.binary()),
mask=pa.array([False, True]),
).cast(pa.list_(child))
target = pa.schema([pa.field("images", pa.list_(lancedb.blob("image")))])
output = _cast_to_target_schema(
pa.table({"images": source}).to_reader(), target
).read_all()
images = output["images"].chunk(0)
assert images.to_pylist()[1] is None
assert [b["data"] for b in images.to_pylist()[0]] == [b"a", b"b"]
def test_fixed_size_list_blob_coercion_keeps_null_rows():
child = pa.field("frame", pa.binary())
source = (
pa.FixedSizeListArray.from_arrays(
pa.array([b"a", b"b", b"c", b"d"], type=pa.binary()), 2
)
.take(pa.array([0, None], type=pa.int32()))
.cast(pa.list_(child, 2))
)
target = pa.schema([pa.field("frames", pa.list_(lancedb.blob("frame"), 2))])
output = _cast_to_target_schema(
pa.table({"frames": source}).to_reader(), target
).read_all()
frames = output["frames"].chunk(0)
assert frames.to_pylist()[1] is None
assert [b["data"] for b in frames.to_pylist()[0]] == [b"a", b"b"]
def test_cast_to_target_schema_accepts_pylance_blob_v2():
target_type = lancedb.BlobType()
source = lance.blob_array([b"hello", None])
assert type(source.type) is LanceBlobType
assert type(source.type) is type(target_type)
data = pa.table({"image": source})
target = pa.schema([pa.field("image", target_type)])
output = _cast_to_target_schema(data.to_reader(), target).read_all()
image = output["image"].chunk(0)
assert type(image.type) is LanceBlobType
assert image.type == target_type
assert image.storage.to_pylist() == [
{"data": b"hello", "uri": None, "position": None, "size": None},
None,
]
def test_cast_to_target_schema_rejects_different_blob_v2_class():
class OtherBlobType(pa.ExtensionType):
def __init__(self):
super().__init__(lancedb.BlobType().storage_type, "lance.blob.v2")
def __arrow_ext_serialize__(self) -> bytes:
return b""
@classmethod
def __arrow_ext_deserialize__(
cls, storage_type: pa.DataType, serialized: bytes
) -> "OtherBlobType":
return cls()
storage = lance.blob_array([b"hello"]).storage
source = pa.ExtensionArray.from_storage(OtherBlobType(), storage)
data = pa.table({"image": source})
target = pa.schema([lancedb.blob("image")])
with pytest.raises(pa.ArrowTypeError, match="different extension type"):
_cast_to_target_schema(data.to_reader(), target).read_all()
def test_sanitize_data_stream():
# Make sure we don't collect the whole stream when running sanitize_data
schema = pa.schema({"a": pa.int32()})
+309 -4
View File
@@ -5,9 +5,9 @@
//! backend-neutral terminal result of a computed-column refresh.
//!
//! This module contains client/wire values only. Catalog persistence,
//! environment bake, and execution are owned by Sophon.
//! environment bake, secret resolution, and execution are owned by Sophon.
use std::collections::BTreeMap;
use std::collections::{BTreeMap, BTreeSet};
use serde::de::{self, DeserializeOwned};
use serde::{Deserialize, Deserializer, Serialize, Serializer};
@@ -15,6 +15,16 @@ use serde_json::Value;
use crate::{Error, Result};
// Keep these byte limits aligned with Sophon's Function submission validation.
pub(crate) const MAX_FUNCTION_SECRET_VALUE_BYTES: usize = 64 * 1024;
const MAX_FUNCTION_SECRET_VALUES_BYTES: usize = 512 * 1024;
fn is_portable_environment_name(name: &str) -> bool {
let mut bytes = name.bytes();
matches!(bytes.next(), Some(b'A'..=b'Z' | b'a'..=b'z' | b'_'))
&& bytes.all(|byte| matches!(byte, b'A'..=b'Z' | b'a'..=b'z' | b'0'..=b'9' | b'_'))
}
fn invalid_json(error: impl std::fmt::Display) -> Error {
Error::InvalidInput {
message: format!("invalid remote Function JSON: {error}"),
@@ -198,6 +208,11 @@ pub struct PythonEnvironmentSpec {
}
/// Reproducible Python runtime definition understood by Sophon.
///
/// `env` contains non-secret values. Secret values are submission-only in the
/// client model and do not become part of this public runtime identity;
/// [`FunctionVersion::required_secrets`] contains names only. Sophon persists
/// submitted values separately in the private execution artifact.
#[derive(Debug, Clone, PartialEq, Eq)]
#[non_exhaustive]
pub enum PythonRuntimeSpec {
@@ -239,7 +254,7 @@ impl PythonRuntimeSpec {
}
}
/// Environment variables, or `None` for an unknown kind.
/// Non-secret environment variables, or `None` for an unknown kind.
pub fn env(&self) -> Option<&BTreeMap<String, String>> {
match self {
Self::Python { env, .. } => Some(env),
@@ -324,6 +339,8 @@ pub struct FunctionVersion {
runtime: PythonRuntimeSpec,
runtime_digest: String,
environment_digest: String,
#[serde(default, skip_serializing_if = "Vec::is_empty")]
required_secrets: Vec<String>,
created_at: String,
}
@@ -356,6 +373,12 @@ impl FunctionVersion {
&self.environment_digest
}
/// Required secret names. Resolved values exist only in Sophon's private
/// execution artifact and worker launch path.
pub fn required_secrets(&self) -> &[String] {
&self.required_secrets
}
pub fn created_at(&self) -> &str {
&self.created_at
}
@@ -397,12 +420,115 @@ pub struct FunctionArtifactRequest {
}
/// Stable request envelope for remote immutable Function registration.
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
///
/// Secret values are submission-only in the client model. Sophon persists them
/// in the database-scoped private execution artifact; returned
/// [`FunctionVersion`] and Job metadata contain only
/// [`Self::required_secrets`] names. Debug formatting always redacts values.
#[derive(Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct FunctionRegistrationRequest {
pub name: String,
pub artifact: FunctionArtifactRequest,
pub signature: FunctionSignature,
pub runtime: PythonRuntimeSpec,
#[serde(default, skip_serializing_if = "Vec::is_empty")]
pub required_secrets: Vec<String>,
#[serde(default, skip_serializing_if = "BTreeMap::is_empty")]
pub secret_values: BTreeMap<String, String>,
}
impl FunctionRegistrationRequest {
pub(crate) fn validate_secret_values(&self) -> Result<()> {
let mut required = BTreeSet::new();
for name in &self.required_secrets {
if !is_portable_environment_name(name) {
return Err(Error::InvalidInput {
message: format!(
"Function secret name {name:?} must be a portable environment variable name"
),
});
}
if !required.insert(name) {
return Err(Error::InvalidInput {
message: format!("Function required_secrets contains duplicate name {name:?}"),
});
}
}
if let PythonRuntimeSpec::Python { env, .. } = &self.runtime
&& let Some(name) = required.iter().find(|name| env.contains_key(**name))
{
return Err(Error::InvalidInput {
message: format!(
"Function runtime env and secret names must be disjoint: {name:?}"
),
});
}
let provided = self.secret_values.keys().collect::<BTreeSet<_>>();
if required != provided {
return Err(Error::InvalidInput {
message: "Function secret_values keys must exactly match required_secrets"
.to_string(),
});
}
let mut total_bytes = 0usize;
for (name, value) in &self.secret_values {
if value.is_empty() {
return Err(Error::InvalidInput {
message: format!("Function secret {name:?} value must be non-empty"),
});
}
if value.contains('\0') {
return Err(Error::InvalidInput {
message: format!("Function secret {name:?} value must not contain NUL"),
});
}
if value.len() > MAX_FUNCTION_SECRET_VALUE_BYTES {
return Err(Error::InvalidInput {
message: format!(
"Function secret {name:?} value exceeds the \
{MAX_FUNCTION_SECRET_VALUE_BYTES}-byte limit"
),
});
}
total_bytes =
total_bytes
.checked_add(value.len())
.ok_or_else(|| Error::InvalidInput {
message: "Function secret values exceed the request byte limit".to_string(),
})?;
}
if total_bytes > MAX_FUNCTION_SECRET_VALUES_BYTES {
return Err(Error::InvalidInput {
message: format!(
"Function secret values exceed the \
{MAX_FUNCTION_SECRET_VALUES_BYTES}-byte request limit"
),
});
}
Ok(())
}
}
impl std::fmt::Debug for FunctionRegistrationRequest {
fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
let secret_values = self
.secret_values
.keys()
.map(|name| (name, "[REDACTED]"))
.collect::<BTreeMap<_, _>>();
formatter
.debug_struct("FunctionRegistrationRequest")
.field("name", &self.name)
.field("artifact", &self.artifact)
.field("signature", &self.signature)
.field("runtime", &self.runtime)
.field("required_secrets", &self.required_secrets)
.field("secret_values", &secret_values)
.finish()
}
}
impl_json!(FunctionRegistrationRequest);
@@ -587,6 +713,185 @@ impl RefreshColumnResult {
impl_json!(RefreshColumnResult);
#[cfg(test)]
mod secret_value_tests {
use super::{
FunctionRegistrationRequest, MAX_FUNCTION_SECRET_VALUE_BYTES,
MAX_FUNCTION_SECRET_VALUES_BYTES, PythonRuntimeSpec,
};
use crate::Error;
fn request() -> FunctionRegistrationRequest {
FunctionRegistrationRequest::from_json(include_str!(
"../tests/fixtures/first_class_functions/v1/remote_function_registration_request.json"
))
.unwrap()
}
#[test]
fn validates_secret_name_and_value_invariants() {
let missing = request();
assert!(matches!(
missing.validate_secret_values(),
Err(Error::InvalidInput { message }) if message.contains("exactly match")
));
let mut empty = request();
empty
.secret_values
.insert("API_TOKEN".to_string(), String::new());
assert!(matches!(
empty.validate_secret_values(),
Err(Error::InvalidInput { message }) if message.contains("non-empty")
));
let mut nul = request();
nul.secret_values
.insert("API_TOKEN".to_string(), "before\0after".to_string());
assert!(matches!(
nul.validate_secret_values(),
Err(Error::InvalidInput { message }) if message.contains("NUL")
));
let mut unexpected = request();
unexpected
.secret_values
.insert("OTHER".to_string(), "value".to_string());
assert!(matches!(
unexpected.validate_secret_values(),
Err(Error::InvalidInput { message }) if message.contains("exactly match")
));
}
#[test]
fn rejects_invalid_duplicate_and_overlapping_secret_declarations() {
let mut invalid_name = request();
invalid_name.required_secrets = vec!["BAD=NAME".to_string()];
invalid_name
.secret_values
.insert("BAD=NAME".to_string(), "secret".to_string());
let mut duplicate = request();
duplicate.required_secrets = vec!["API_TOKEN".to_string(), "API_TOKEN".to_string()];
duplicate
.secret_values
.insert("API_TOKEN".to_string(), "secret".to_string());
let mut overlap = request();
overlap
.secret_values
.insert("API_TOKEN".to_string(), "secret".to_string());
if let PythonRuntimeSpec::Python { env, .. } = &mut overlap.runtime {
env.insert("API_TOKEN".to_string(), "public".to_string());
}
assert!(matches!(
invalid_name.validate_secret_values(),
Err(Error::InvalidInput { message }) if message.contains("portable environment variable")
));
assert!(matches!(
duplicate.validate_secret_values(),
Err(Error::InvalidInput { message }) if message.contains("duplicate")
));
assert!(matches!(
overlap.validate_secret_values(),
Err(Error::InvalidInput { message }) if message.contains("must be disjoint")
));
}
#[test]
fn enforces_portable_secret_name_boundaries() {
for name in ["A", "_", "A0_"] {
let mut request = request();
request.required_secrets = vec![name.to_string()];
request
.secret_values
.insert(name.to_string(), "secret".to_string());
request.validate_secret_values().unwrap();
}
for name in ["", "0TOKEN", "BAD-NAME", "TÖKEN"] {
let mut request = request();
request.required_secrets = vec![name.to_string()];
request
.secret_values
.insert(name.to_string(), "secret".to_string());
assert!(matches!(
request.validate_secret_values(),
Err(Error::InvalidInput { message })
if message.contains("portable environment variable")
));
}
}
#[test]
fn accepts_exact_secret_value_utf8_byte_limit() {
for value in [
"x".repeat(MAX_FUNCTION_SECRET_VALUE_BYTES),
"é".repeat(MAX_FUNCTION_SECRET_VALUE_BYTES / "é".len()),
] {
assert_eq!(value.len(), MAX_FUNCTION_SECRET_VALUE_BYTES);
let mut request = request();
request.secret_values.insert("API_TOKEN".to_string(), value);
request.validate_secret_values().unwrap();
}
}
#[test]
fn rejects_secret_value_over_utf8_byte_limit() {
for value in [
"x".repeat(MAX_FUNCTION_SECRET_VALUE_BYTES + 1),
"é".repeat(MAX_FUNCTION_SECRET_VALUE_BYTES / "é".len() + 1),
] {
assert!(value.len() > MAX_FUNCTION_SECRET_VALUE_BYTES);
let mut request = request();
request.secret_values.insert("API_TOKEN".to_string(), value);
assert!(matches!(
request.validate_secret_values(),
Err(Error::InvalidInput { message }) if message.contains("65536-byte limit")
));
}
}
#[test]
fn rejects_aggregate_secret_value_bytes_over_server_limit() {
let mut request = request();
request.required_secrets = (0..9).map(|index| format!("SECRET_{index}")).collect();
request.secret_values = request
.required_secrets
.iter()
.map(|name| (name.clone(), "x".repeat(MAX_FUNCTION_SECRET_VALUE_BYTES)))
.collect();
assert!(matches!(
request.validate_secret_values(),
Err(Error::InvalidInput { message })
if message.contains(&format!("{MAX_FUNCTION_SECRET_VALUES_BYTES}-byte request limit"))
));
}
#[test]
fn accepts_exact_aggregate_secret_value_byte_limit() {
let mut request = request();
request.required_secrets = (0..8).map(|index| format!("SECRET_{index}")).collect();
request.secret_values = request
.required_secrets
.iter()
.map(|name| (name.clone(), "x".repeat(MAX_FUNCTION_SECRET_VALUE_BYTES)))
.collect();
assert_eq!(
request
.secret_values
.values()
.map(String::len)
.sum::<usize>(),
MAX_FUNCTION_SECRET_VALUES_BYTES
);
request.validate_secret_values().unwrap();
}
}
#[cfg(test)]
mod conda_environment_tests {
use super::PythonEnvironmentSpec;
+100 -15
View File
@@ -7,6 +7,7 @@ use reqwest::{
Body, Request, RequestBuilder, Response,
header::{HeaderMap, HeaderValue},
};
use serde_json::Value;
use std::{collections::HashMap, future::Future, str::FromStr, sync::Arc, time::Duration};
use crate::error::{Error, Result};
@@ -14,6 +15,60 @@ use crate::remote::db::RemoteOptions;
use crate::remote::retry::{ResolvedRetryConfig, RetryCounter};
const REQUEST_ID_HEADER: HeaderName = HeaderName::from_static("x-request-id");
const REDACTED_JSON_VALUE: &str = "[REDACTED]";
const SUPPRESSED_JSON_BODY: &str = "[JSON BODY SUPPRESSED]";
fn is_sensitive_json_field(name: &str) -> bool {
name.to_ascii_lowercase().contains("secret")
}
fn redact_sensitive_json_fields(value: &mut Value) {
match value {
Value::Object(fields) => {
for (name, child) in fields {
if is_sensitive_json_field(name) {
*child = Value::String(REDACTED_JSON_VALUE.to_string());
} else {
redact_sensitive_json_fields(child);
}
}
}
Value::Array(values) => values.iter_mut().for_each(redact_sensitive_json_fields),
_ => {}
}
}
fn redacted_json_body(request: &Request) -> Option<String> {
let body = request.body()?.as_bytes()?;
let mut value = serde_json::from_slice(body).ok()?;
redact_sensitive_json_fields(&mut value);
serde_json::to_string(&value).ok()
}
fn request_log_message(request: &Request, request_id: &str) -> String {
let prefix = format!(
"Sending request_id={}: {} {}",
request_id,
request.method(),
request.url()
);
let content_type = request
.headers()
.get("content-type")
.and_then(|value| value.to_str().ok())
.and_then(|value| value.split(';').next());
if content_type.is_some_and(|value| value.eq_ignore_ascii_case("application/json")) {
// Never format the raw Request here: its Debug representation is not a
// redaction boundary and may include the original body. If the JSON body
// cannot be structurally parsed, suppress it instead of logging raw bytes.
let body = redacted_json_body(request).unwrap_or_else(|| SUPPRESSED_JSON_BODY.to_string());
format!("{prefix} with body {body}")
} else {
// Method and URL are sufficient request context. Raw Request formatting
// may expose headers or a non-JSON body, so it is never a logging fallback.
prefix
}
}
/// Configuration for TLS/mTLS settings.
#[derive(Clone, Debug)]
@@ -839,22 +894,9 @@ impl<S: HttpSend> RestfulLanceDbClient<S> {
}
}
pub(crate) fn log_request(&self, request: &Request, request_id: &String) {
pub(crate) fn log_request(&self, request: &Request, request_id: &str) {
if log::log_enabled!(log::Level::Debug) {
let content_type = request
.headers()
.get("content-type")
.map(|v| v.to_str().unwrap());
if content_type == Some("application/json") {
let body = request.body().as_ref().unwrap().as_bytes().unwrap();
let body = String::from_utf8_lossy(body);
debug!(
"Sending request_id={}: {:?} with body {}",
request_id, request, body
);
} else {
debug!("Sending request_id={}: {:?}", request_id, request);
}
debug!("{}", request_log_message(request, request_id));
}
}
@@ -1077,6 +1119,49 @@ mod tests {
ENV_MUTEX.lock().unwrap_or_else(|e| e.into_inner())
}
#[test]
fn test_request_log_message_redacts_secrets_and_never_formats_raw_requests() {
const SECRET_SENTINEL: &str = "udf-secret-log-sentinel-7e4e";
const MALFORMED_SENTINEL: &str = "malformed-secret-log-sentinel-b652";
const NON_JSON_SENTINEL: &str = "non-json-secret-log-sentinel-7fd1";
let request = reqwest::Client::new()
.post("https://example.com/v1/functions/create")
.json(&serde_json::json!({
"name": "uses_secret",
"nested": {
"secret_values": {"OPENAI_API_KEY": SECRET_SENTINEL},
"safe": "visible-value"
}
}))
.build()
.unwrap();
let log_message = request_log_message(&request, "valid-json");
let malformed_request = reqwest::Client::new()
.post("https://example.com/v1/functions/create")
.header("content-type", "application/json; charset=utf-8")
.body(format!(r#"{{"secret_values":"{MALFORMED_SENTINEL}""#))
.build()
.unwrap();
let malformed_log_message = request_log_message(&malformed_request, "malformed-json");
let non_json_request = reqwest::Client::new()
.post("https://example.com/v1/functions/create")
.header("content-type", "text/plain")
.body(NON_JSON_SENTINEL)
.build()
.unwrap();
let non_json_log_message = request_log_message(&non_json_request, "non-json");
assert!(log_message.contains("visible-value"));
assert!(log_message.contains(REDACTED_JSON_VALUE));
assert!(!log_message.contains(SECRET_SENTINEL));
assert!(malformed_log_message.contains(SUPPRESSED_JSON_BODY));
assert!(!malformed_log_message.contains(MALFORMED_SENTINEL));
assert!(!non_json_log_message.contains(NON_JSON_SENTINEL));
}
#[test]
fn test_timeout_config_default() {
let config = TimeoutConfig::default();
+33 -2
View File
@@ -554,6 +554,7 @@ impl<S: HttpSend> Database for RemoteDatabase<S> {
&self,
request: FunctionRegistrationRequest,
) -> Result<Job<FunctionVersion>> {
request.validate_secret_values()?;
let req = self.client.post("/v1/functions/create").json(&request);
let (request_id, response) = self.client.send(req).await?;
let response = self.client.check_response(&request_id, response).await?;
@@ -2642,7 +2643,8 @@ mod tests {
);
const FUNCTION_JOB: &str =
include_str!("../../tests/fixtures/first_class_functions/v1/remote_function_job.json");
let expected: serde_json::Value = serde_json::from_str(REQUEST).unwrap();
let mut expected: serde_json::Value = serde_json::from_str(REQUEST).unwrap();
expected["secret_values"] = serde_json::json!({"API_TOKEN": "secret-value"});
let conn = Connection::new_with_handler(move |request| match request.url().path() {
"/v1/functions/create" => {
assert_eq!(request.method(), &reqwest::Method::POST);
@@ -2660,7 +2662,10 @@ mod tests {
.unwrap(),
path => panic!("unexpected path: {path}"),
});
let request = crate::function::FunctionRegistrationRequest::from_json(REQUEST).unwrap();
let mut request = crate::function::FunctionRegistrationRequest::from_json(REQUEST).unwrap();
request
.secret_values
.insert("API_TOKEN".to_string(), "secret-value".to_string());
let job = conn.create_function_async(request).await.unwrap();
assert_eq!(job.id(), Some("job-function-1"));
let version = job.wait().await.unwrap();
@@ -2668,6 +2673,32 @@ mod tests {
assert_eq!(version.version(), "fv_01K3EXACT");
}
#[tokio::test]
async fn test_create_function_async_validates_secrets_before_serialization_and_send() {
const REQUEST: &str = include_str!(
"../../tests/fixtures/first_class_functions/v1/remote_function_registration_request.json"
);
let sends = Arc::new(AtomicUsize::new(0));
let sends_ref = sends.clone();
let conn = Connection::new_with_handler(move |_| {
sends_ref.fetch_add(1, Ordering::SeqCst);
http::Response::builder().status(500).body("").unwrap()
});
let mut request = crate::function::FunctionRegistrationRequest::from_json(REQUEST).unwrap();
request.secret_values.insert(
"API_TOKEN".to_string(),
"x".repeat(crate::function::MAX_FUNCTION_SECRET_VALUE_BYTES + 1),
);
let error = conn.create_function_async(request).await.unwrap_err();
assert!(matches!(
error,
Error::InvalidInput { message } if message.contains("65536-byte limit")
));
assert_eq!(sends.load(Ordering::SeqCst), 0);
}
#[tokio::test]
async fn test_get_function_requires_and_sends_exact_version() {
const VERSION: &str = include_str!(
+3 -2
View File
@@ -5734,8 +5734,9 @@ mod tests {
))
.execute()
.await;
let Err(err) = result else {
panic!("legacy remote query unexpectedly succeeded")
let err = match result {
Ok(_) => panic!("legacy remote query unexpectedly succeeded"),
Err(err) => err,
};
assert!(
+8 -97
View File
@@ -22,7 +22,7 @@
use std::collections::{BTreeSet, HashMap};
use std::sync::Arc;
use arrow_schema::{DataType, Field as ArrowField, Fields, Schema as ArrowSchema, SchemaRef};
use arrow_schema::{DataType, Field as ArrowField, Schema as ArrowSchema, SchemaRef};
use datafusion_common::tree_node::TreeNode;
use datafusion_physical_plan::PhysicalExpr;
use lance::dataset::NewColumnTransform;
@@ -1273,11 +1273,6 @@ pub(crate) fn bind(schema: SchemaRef, column: &str, expression: &str) -> Result<
/// refresh time: that the expression parses, that every column it reads
/// exists, and that the target name is free. A declaration that survives this
/// is one a refresh can always act on.
///
/// Each accepted column joins the schema the next one resolves against, so a
/// batch may declare `a` and then `b = a + 1` in one commit. Refresh fills a
/// column's computed inputs before the column, so the order of refresh calls
/// does not matter.
pub(crate) fn plan(schema: SchemaRef, columns: &[(String, String)]) -> Result<Vec<ArrowField>> {
if columns.is_empty() {
return Err(Error::InvalidInput {
@@ -1285,11 +1280,11 @@ pub(crate) fn plan(schema: SchemaRef, columns: &[(String, String)]) -> Result<Ve
});
}
let mut schema = schema;
let mut fields = Vec::with_capacity(columns.len());
let mut declared: Vec<&str> = Vec::with_capacity(columns.len());
for (name, expression) in columns {
if schema.field_with_name(name).is_ok() {
if schema.field_with_name(name).is_ok() || declared.contains(&name.as_str()) {
return Err(Error::ColumnAlreadyExists { name: name.clone() });
}
@@ -1297,50 +1292,16 @@ pub(crate) fn plan(schema: SchemaRef, columns: &[(String, String)]) -> Result<Ve
// Declared columns start entirely null, so nullability is a property
// of the declaration rather than of what the expression yields.
let field = ArrowField::new(name, bound.data_type, true)
.with_metadata(computed_column_metadata(expression, &bound.inputs));
schema = Arc::new(ArrowSchema::new_with_metadata(
schema
.fields()
.iter()
.cloned()
.chain(std::iter::once(Arc::new(field.clone())))
.collect::<Fields>(),
schema.metadata().clone(),
));
fields.push(field);
fields.push(
ArrowField::new(name, bound.data_type, true)
.with_metadata(computed_column_metadata(expression, &bound.inputs)),
);
declared.push(name);
}
Ok(fields)
}
/// Run the schema-level checks of
/// [`AddColumnsBuilder::computed`](super::AddColumnsBuilder::computed) against
/// `schema` without committing: the Function-binding guard and the planning of
/// every declaration. For callers that stage declarations behind other work
/// and need those rejections before any of it lands.
///
/// Only the schema is consulted. Declaring also refuses a table with an LSM
/// write spec or retained SSTables; that is table state, checked at commit.
///
/// ```
/// # use std::sync::Arc;
/// # use arrow_schema::{DataType, Field, Schema};
/// use lancedb::table::computed_columns::validate_declarations;
///
/// let schema = Arc::new(Schema::new(vec![Field::new("x", DataType::Int32, false)]));
/// let declarations = vec![
/// ("a".to_string(), "x + 1".to_string()),
/// ("b".to_string(), "a * 2".to_string()),
/// ];
/// assert!(validate_declarations(schema.clone(), &declarations).is_ok());
/// assert!(validate_declarations(schema, &[("c".into(), "random()".into())]).is_err());
/// ```
pub fn validate_declarations(schema: SchemaRef, columns: &[(String, String)]) -> Result<()> {
ensure_no_function_bindings_for_mutation(schema.as_ref(), "schema evolution")?;
plan(schema, columns).map(drop)
}
/// Build the transform that declares `columns` against `schema`.
///
/// An all-null column is how a binding with no values yet is carried into a
@@ -1379,22 +1340,6 @@ pub(super) async fn add_foreign_kind(table: &crate::Table, name: &str, kind: &st
#[cfg(test)]
mod tests {
/// The gate's reproducer: the validator applies the same schema-level
/// guard declaring does, so a staging caller is refused before it commits
/// anything else.
#[test]
fn test_validate_declarations_matches_schema_admission_barriers() {
let schema = Arc::new(ArrowSchema::new_with_metadata(
vec![ArrowField::new("x", DataType::Int32, true)],
HashMap::from([(
FUNCTION_BINDINGS_META_KEY.to_string(),
"not valid binding metadata".to_string(),
)]),
));
let declarations = vec![("a".to_string(), "x + 1".to_string())];
assert!(super::validate_declarations(schema, &declarations).is_err());
}
#[test]
fn output_arrow_type_grammar_matches_the_shared_golden() {
let golden: serde_json::Value = serde_json::from_str(include_str!(
@@ -1637,40 +1582,6 @@ mod tests {
assert!(declared(&table).await.is_empty());
}
/// A batch may build on itself: one commit, and the later entry's inputs
/// name the earlier one.
#[tokio::test]
async fn test_a_declaration_may_read_one_declared_before_it() {
let table = table_with_ints("chain").await;
let before = table.version().await.unwrap();
add_computed(
&table,
&[("a".into(), "x + 1".into()), ("b".into(), "a * 2".into())],
)
.await
.unwrap();
assert_eq!(table.version().await.unwrap(), before + 1);
let declared = declared(&table).await;
assert_eq!(declared[1].name, "b");
assert_eq!(declared[1].inputs, vec!["a".to_string()]);
// Order is the dependency order; reading ahead is still unknown.
let err = add_computed(
&table,
&[("c".into(), "d + 1".into()), ("d".into(), "x + 1".into())],
)
.await
.unwrap_err();
assert!(matches!(err, Error::InvalidExpression { column, .. } if column == "c"));
assert!(
validate_declarations(
table.schema().await.unwrap(),
&[("e".into(), "random()".into())]
)
.is_err()
);
}
/// A column added by an ordinary transform is materialized, not bound, so
/// it carries no declaration to report.
#[tokio::test]
@@ -36,6 +36,14 @@ pub(super) fn coerce_blob_expr(
};
let input_shape = match input_field.data_type() {
DataType::Null => {
let expr: Arc<dyn PhysicalExpr> = Arc::new(CastExpr::new(
input_expr,
table_field.data_type().clone(),
None,
));
return Ok((expr, table_field.clone()));
}
DataType::Binary | DataType::LargeBinary | DataType::BinaryView => BlobInputShape::Bytes,
DataType::Utf8 | DataType::LargeUtf8 | DataType::Utf8View => BlobInputShape::String,
DataType::Struct(children) => {
@@ -155,7 +163,7 @@ mod tests {
use crate::blob::blob;
use arrow_array::{
Array, ArrayRef, BinaryArray, BinaryViewArray, Int32Array, Int64Array, LargeBinaryArray,
RecordBatch, StringArray, StringViewArray, StructArray, UInt8Array, UInt64Array,
NullArray, RecordBatch, StringArray, StringViewArray, StructArray, UInt8Array, UInt64Array,
};
use arrow_schema::Schema;
use datafusion::prelude::SessionContext;
@@ -279,6 +287,18 @@ mod tests {
assert_eq!(data.value(0), b"view");
}
#[tokio::test]
async fn null_column_coerces_to_all_null_blob_struct() {
let batch = batch_with_image(
Field::new("image", DataType::Null, true),
Arc::new(NullArray::new(2)),
);
let coerced = coerce(batch, &blob_table_schema()).await;
let image = image_struct(&coerced);
assert!(image.is_null(0));
assert!(image.is_null(1));
}
#[tokio::test]
async fn binary_nulls_stay_null_after_coercion() {
let batch = batch_with_image(
+24 -363
View File
@@ -7,16 +7,6 @@
//! therefore idempotent and does not observe input mutation -- once a row is
//! filled, changing what the expression reads leaves the stored result alone.
//!
//! A column's computed inputs are filled first -- the dependency graph is
//! walked once, each reachable column filled once in dependency order, each
//! fill its own commit. Every fill in the pass, the requested column's
//! included, covers only the fragments of the snapshot the pass started
//! from: a commit may rebase over a concurrent append, and the fragment that
//! admits carries placeholder nulls no earlier fill covered, so it waits for
//! a later refresh rather than being read as values. Two concurrent fills of
//! one input collide on its field in lance's conflict check, so a dependent
//! fill can only commit over inputs that were durable when it read them.
//!
//! Two passes per fragment. The first scans only the unfilled live rows and
//! evaluates the expression over them, which yields the exact fill count and
//! decides whether the fragment is staged at all -- a fragment where nothing
@@ -29,7 +19,6 @@
//! inputs masked to null first, so a poison value in a row nobody is filling
//! cannot fail the refresh.
use std::collections::HashSet;
use std::sync::Arc;
use arrow_array::{ArrayRef, BooleanArray, RecordBatch, RecordBatchOptions};
@@ -52,8 +41,7 @@ use crate::{Error, Result};
/// The result of refreshing a computed column.
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, Default)]
pub struct RefreshColumnResult {
/// Rows that had a value computed, in the requested column only; inputs
/// filled on its behalf are not counted.
/// Rows that had a value computed.
#[serde(default)]
pub rows_filled: u64,
/// The commit version associated with the operation.
@@ -63,79 +51,7 @@ pub struct RefreshColumnResult {
struct RefreshExecution {
result: RefreshColumnResult,
/// The snapshot the requested column was evaluated against, after its
/// inputs were filled.
source_version: u64,
/// The last version any fill in the pass committed, inputs included.
published_version: Option<u64>,
/// Unfilled live rows of the requested column in fragments the pass did
/// not cover, counted on the snapshot the result reports -- the published
/// one, or the source when nothing was published.
rows_deferred: u64,
}
/// One column's fill against one snapshot.
struct Fill {
rows_filled: u64,
/// The dataset the commit produced, if anything was filled.
committed: Option<Arc<Dataset>>,
}
/// Test-only one-shot pauses before a named column's fill:
/// `(column, reached, resume)`, one per column, consumed when hit.
#[cfg(test)]
type Pause = (String, Arc<tokio::sync::Notify>, Arc<tokio::sync::Notify>);
#[cfg(test)]
static PAUSES_BEFORE_FILL: std::sync::Mutex<Vec<Pause>> = std::sync::Mutex::new(Vec::new());
/// The SQL-computed columns `column` reads, transitively, each once, in an
/// order that fills every column after the columns it reads.
///
/// Declarations are acyclic by construction: a column can only read what
/// existed when it was declared.
fn dependency_order(schema: &ArrowSchema, column: &str) -> Result<Vec<String>> {
fn visit(
schema: &ArrowSchema,
column: &str,
visited: &mut HashSet<String>,
order: &mut Vec<String>,
) -> Result<()> {
let Some(declaration) = schema
.field_with_name(column)
.ok()
.and_then(computed_column_from_field)
else {
return Ok(());
};
for input in &declaration.inputs {
let input = super::computed_columns::root(input);
if visited.contains(input) {
continue;
}
let Some(input_declaration) = schema
.field_with_name(input)
.ok()
.and_then(computed_column_from_field)
else {
continue;
};
if !matches!(input_declaration.kind, ComputedColumnKind::Sql { .. }) {
return Err(Error::NotSupported {
message: format!(
"computed column '{column}' reads '{input}', which this refresh \
cannot fill first; refresh '{input}' before '{column}'"
),
});
}
visit(schema, input, visited, order)?;
visited.insert(input.to_string());
order.push(input.to_string());
}
Ok(())
}
let mut order = Vec::new();
visit(schema, column, &mut HashSet::new(), &mut order)?;
Ok(order)
}
/// Internal implementation of the refresh logic.
@@ -154,96 +70,9 @@ async fn execute_refresh_column_with_source(
) -> Result<RefreshExecution> {
table.dataset.ensure_mutable()?;
ensure_no_lsm_write_spec(table).await?;
let mut dataset = table.dataset.get().await?;
declared_expression(&dataset, column)?;
let dataset = table.dataset.get().await?;
// The pass covers exactly these fragments. A commit below may rebase over
// a concurrent append, and the fragment that admits was scanned by no
// earlier fill, so it is excluded from every later one.
let eligible: HashSet<u64> = dataset
.get_fragments()
.iter()
.map(|f| f.id() as u64)
.collect();
let mut published_version = None;
for input in dependency_order(&ArrowSchema::from(dataset.schema()), column)? {
pause_before_fill(&input).await;
if let Some(committed) = fill_column(table, &dataset, &input, &eligible)
.await?
.committed
{
published_version = Some(committed.version().version);
dataset = committed;
}
}
pause_before_fill(column).await;
let source_version = dataset.version().version;
let fill = fill_column(table, &dataset, column, &eligible).await?;
if let Some(committed) = &fill.committed {
published_version = Some(committed.version().version);
}
// Counted after the commit: it may have rebased over an append, and the
// fragment that admits is deferred but part of the published version.
let terminal = fill.committed.as_ref().unwrap_or(&dataset);
let rows_deferred = count_deferred(terminal, column, &eligible).await?;
Ok(RefreshExecution {
result: RefreshColumnResult {
rows_filled: fill.rows_filled,
version: published_version.unwrap_or(source_version),
},
source_version,
published_version,
rows_deferred,
})
}
/// Unfilled live rows of `column` in `dataset`'s fragments outside `eligible`.
async fn count_deferred(dataset: &Dataset, column: &str, eligible: &HashSet<u64>) -> Result<u64> {
let deferred: Vec<_> = dataset
.get_fragments()
.into_iter()
.filter(|fragment| !eligible.contains(&(fragment.id() as u64)))
.map(|fragment| fragment.metadata().clone())
.collect();
if deferred.is_empty() {
return Ok(0);
}
let mut scanner = dataset.scan();
scanner
.with_fragments(deferred)
.filter(&format!("{} IS NULL", quote_identifier(column)))?;
Ok(scanner.count_rows().await?)
}
#[cfg(test)]
async fn pause_before_fill(column: &str) {
let pause = {
let mut pauses = PAUSES_BEFORE_FILL.lock().unwrap();
pauses
.iter()
.position(|(paused, _, _)| paused == column)
.map(|index| pauses.remove(index))
};
if let Some((_, reached, resume)) = pause {
reached.notify_one();
resume.notified().await;
}
}
#[cfg(not(test))]
async fn pause_before_fill(_column: &str) {}
/// Fill `column`'s unfilled live rows in the `eligible` fragments as they
/// stand in `dataset`, committing against that snapshot.
async fn fill_column(
table: &NativeTable,
dataset: &Arc<Dataset>,
column: &str,
eligible: &HashSet<u64>,
) -> Result<Fill> {
let expression = declared_expression(dataset, column)?;
let expression = declared_expression(&dataset, column)?;
let schema = Arc::new(ArrowSchema::from(dataset.schema()));
let bound = Arc::new(super::computed_columns::bind(schema, column, &expression)?);
let field = dataset
@@ -262,22 +91,23 @@ async fn fill_column(
let mut rows_filled = 0u64;
let mut replacements = Vec::new();
for fragment in dataset.get_fragments() {
if !eligible.contains(&(fragment.id() as u64)) {
continue;
}
let gained = count_fragment_gains(dataset, &fragment, &bound, column).await?;
let gained = count_fragment_gains(&dataset, &fragment, &bound, column).await?;
if gained == 0 {
continue;
}
rows_filled += gained;
let values = fill_stream(dataset, &fragment, bound.clone(), column).await?;
let values = fill_stream(&dataset, &fragment, bound.clone(), column).await?;
replacements.push(fragment.write_columns(values, &column_schema).await?);
}
if replacements.is_empty() {
return Ok(Fill {
rows_filled: 0,
committed: None,
let source_version = dataset.version().version;
return Ok(RefreshExecution {
result: RefreshColumnResult {
rows_filled: 0,
version: source_version,
},
source_version,
});
}
@@ -295,10 +125,15 @@ async fn fill_column(
false,
)
.await?;
table.dataset.update(new_dataset.clone());
Ok(Fill {
rows_filled,
committed: Some(Arc::new(new_dataset)),
let version = new_dataset.version().version;
table.dataset.update(new_dataset);
Ok(RefreshExecution {
result: RefreshColumnResult {
rows_filled,
version,
},
source_version: read_version,
})
}
@@ -323,9 +158,10 @@ pub(crate) async fn execute_refresh_column_async(
Ok(crate::function::RefreshColumnResult {
rows_assigned: execution.result.rows_filled,
rows_failed: 0,
rows_remaining: execution.rows_deferred,
rows_remaining: 0,
source_version: execution.source_version,
published_version: execution.published_version,
published_version: (execution.result.rows_filled > 0)
.then_some(execution.result.version),
})
})))
}
@@ -578,181 +414,6 @@ mod tests {
table.add(batch).execute().await.unwrap();
}
/// The gate's reproducer: `b = coalesce(a, 0)` refreshed before `a`
/// must not bake zeros from `a`'s placeholder null.
#[tokio::test]
async fn test_dependent_refresh_cannot_fill_from_placeholder_null() {
let table = table_with("dependent_refresh_order", vec![1, 2, 3]).await;
table
.add_columns()
.computed("a", "x + 1")
.computed("b", "coalesce(a, 0)")
.execute()
.await
.unwrap();
let result = table.refresh_column("b").await.unwrap();
assert_eq!(result.rows_filled, 3);
assert_eq!(read(&table, "a").await, vec![Some(2), Some(3), Some(4)]);
assert_eq!(
table.count_rows(Some("b = a".to_string())).await.unwrap(),
3
);
assert_eq!(table.refresh_column("a").await.unwrap().rows_filled, 0);
// Appended rows: the input is filled in the new fragment first too.
append(&table, vec![10]).await;
assert_eq!(table.refresh_column("b").await.unwrap().rows_filled, 1);
assert_eq!(
table.count_rows(Some("b = 0".to_string())).await.unwrap(),
0
);
}
/// The gate's reproducer: a fragment appended between the input fill and
/// the requested column's fill has inputs the pass never covered, so it
/// is left unfilled rather than read as null.
#[tokio::test]
async fn test_dependent_refresh_fences_an_append_after_input_fill() {
let table = table_with("dependent_refresh_append_gap", vec![1, 2, 3]).await;
table
.add_columns()
.computed("a_gap", "x + 1")
.computed("b_gap", "coalesce(a_gap, 0)")
.execute()
.await
.unwrap();
let reached = Arc::new(tokio::sync::Notify::new());
let resume = Arc::new(tokio::sync::Notify::new());
super::PAUSES_BEFORE_FILL.lock().unwrap().push((
"b_gap".to_string(),
reached.clone(),
resume.clone(),
));
let refresh = {
let table = table.clone();
tokio::spawn(async move {
table
.refresh_column_async("b_gap")
.await
.unwrap()
.wait()
.await
})
};
reached.notified().await;
append(&table, vec![10]).await;
resume.notify_one();
let result = refresh.await.unwrap().unwrap();
assert_eq!(result.rows_assigned, 3);
// The target commit rebased over the append: the published version
// holds the deferred row, and the count is taken there.
assert_eq!(result.rows_remaining, 1);
table.refresh_column("a_gap").await.unwrap();
assert_eq!(table.count_rows(Some("b_gap = 0".into())).await.unwrap(), 0);
assert_eq!(table.refresh_column("b_gap").await.unwrap().rows_filled, 1);
}
/// The gate's reproducer: an append landing between two input fills is
/// rebased into the second's snapshot, but the first never covered it.
#[tokio::test]
async fn test_dependent_refresh_fences_an_append_between_input_fills() {
let table = table_with("dependent_refresh_between_inputs", vec![1, 2, 3]).await;
table
.add_columns()
.computed("a_mid", "x + 1")
.computed("c_mid", "x + 2")
.computed("b_mid", "coalesce(a_mid, 0) + coalesce(c_mid, 0)")
.execute()
.await
.unwrap();
let reached = Arc::new(tokio::sync::Notify::new());
let resume = Arc::new(tokio::sync::Notify::new());
super::PAUSES_BEFORE_FILL.lock().unwrap().push((
"c_mid".to_string(),
reached.clone(),
resume.clone(),
));
let refresh = {
let table = table.clone();
tokio::spawn(async move {
table
.refresh_column_async("b_mid")
.await
.unwrap()
.wait()
.await
})
};
reached.notified().await;
append(&table, vec![10]).await;
resume.notify_one();
let result = refresh.await.unwrap().unwrap();
assert_eq!(result.rows_assigned, 3);
// The appended row is in the reported source version but was deferred.
assert_eq!(result.rows_remaining, 1);
table.refresh_column("a_mid").await.unwrap();
table.refresh_column("c_mid").await.unwrap();
assert_eq!(table.count_rows(Some("b_mid = 0".into())).await.unwrap(), 0);
assert_eq!(table.refresh_column("b_mid").await.unwrap().rows_filled, 1);
assert_eq!(
table.count_rows(Some("b_mid = 23".into())).await.unwrap(),
1
);
}
/// A dependency commit is a publication even when the requested column
/// itself fills nothing.
#[tokio::test]
async fn test_refresh_async_reports_a_dependency_publication() {
let table = table_with("refresh_async_dependency_publication", vec![1, 2, 3]).await;
table
.add_columns()
.computed("a", "x + 1")
.computed("b", "nullif(a, a)")
.execute()
.await
.unwrap();
let result = table
.refresh_column_async("b")
.await
.unwrap()
.wait()
.await
.unwrap();
assert_eq!(result.rows_assigned, 0);
assert_eq!(result.published_version, Some(result.source_version));
}
/// The gate's reproducer: a dense graph is walked once, not once per path.
#[tokio::test]
async fn test_dependency_refresh_visits_each_column_once() {
let table = table_with("dependency_refresh_deduplicates", vec![1]).await;
let mut declaration = table.add_columns();
for i in 0..=6 {
let expression = if i == 0 {
"x + 1".to_string()
} else {
(0..i)
.map(|j| format!("a{j}"))
.collect::<Vec<_>>()
.join(" + ")
};
declaration = declaration.computed(format!("a{i}"), expression);
}
declaration.execute().await.unwrap();
let order = super::dependency_order(table.schema().await.unwrap().as_ref(), "a6").unwrap();
assert_eq!(order, (0..6).map(|i| format!("a{i}")).collect::<Vec<_>>());
// One commit per column, the requested one included.
let before = table.version().await.unwrap();
table.refresh_column("a6").await.unwrap();
assert_eq!(table.version().await.unwrap(), before + 7);
}
#[tokio::test]
async fn test_refresh_fills_a_declared_column() {
let table = table_with("refresh_fills", vec![1, 2, 3]).await;
@@ -20,6 +20,25 @@ fn job_result(name: &str) -> Value {
serde_json::from_str::<Value>(&fixture(name)).expect("remote Job fixture")["result"].clone()
}
fn assert_no_secret_values(value: &Value) {
match value {
Value::Object(values) => {
for (key, value) in values {
assert!(
!matches!(
key.as_str(),
"secret_value" | "secret_values" | "resolved_secret" | "resolved_secrets"
),
"client canonical value must not model resolved secret material"
);
assert_no_secret_values(value);
}
}
Value::Array(values) => values.iter().for_each(assert_no_secret_values),
_ => {}
}
}
#[test]
fn function_version_job_result_matches_shared_canonical_golden() {
let result = job_result("remote_function_job.json");
@@ -28,6 +47,7 @@ fn function_version_job_result_matches_shared_canonical_golden() {
assert_eq!(version.name(), "embed");
assert_eq!(version.version(), "fv_01K3EXACT");
assert_eq!(version.runtime_digest(), "sha256:runtime");
assert_eq!(version.required_secrets(), &["HF_TOKEN"]);
assert_eq!(
version.to_canonical_json().expect("canonical JSON"),
fixture("remote_function_version.canonical.json").trim()
@@ -142,3 +162,21 @@ fn floating_point_application_literals_are_rejected_consistently() {
.contains("floating-point Function literals")
);
}
#[test]
fn canonical_client_values_contain_secret_names_only() {
let result = job_result("remote_function_job.json");
let version = FunctionVersion::from_json(&result.to_string()).expect("FunctionVersion result");
let canonical: Value = serde_json::from_str(
&version
.to_canonical_json()
.expect("canonical FunctionVersion"),
)
.expect("canonical JSON");
assert_eq!(
canonical["required_secrets"],
serde_json::json!(["HF_TOKEN"])
);
assert_no_secret_values(&canonical);
}
@@ -6,6 +6,7 @@ use std::path::PathBuf;
use lancedb::Error;
use lancedb::function::FunctionRegistrationRequest;
use serde_json::Value;
fn fixture(name: &str) -> String {
let path = PathBuf::from(env!("CARGO_MANIFEST_DIR"))
@@ -14,6 +15,25 @@ fn fixture(name: &str) -> String {
fs::read_to_string(path).expect("fixture must be readable")
}
fn assert_no_secret_values(value: &Value) {
match value {
Value::Object(values) => {
for (key, value) in values {
assert!(
!matches!(
key.as_str(),
"secret_value" | "secret_values" | "resolved_secret" | "resolved_secrets"
),
"registration requests must not model resolved secret material"
);
assert_no_secret_values(value);
}
}
Value::Array(values) => values.iter().for_each(assert_no_secret_values),
_ => {}
}
}
#[test]
fn registration_request_matches_shared_canonical_golden() {
let request = FunctionRegistrationRequest::from_json(&fixture(
@@ -22,10 +42,33 @@ fn registration_request_matches_shared_canonical_golden() {
.expect("registration request");
assert_eq!(request.name, "normalize_score");
assert_eq!(request.artifact.adapter.kind, "scalar_to_arrow_batch");
assert_eq!(request.required_secrets, ["API_TOKEN"]);
assert!(request.secret_values.is_empty());
assert_eq!(
request.to_canonical_json().expect("canonical request"),
fixture("remote_function_registration_request.canonical.json").trim()
);
let value: Value =
serde_json::from_str(&request.to_canonical_json().expect("canonical request"))
.expect("request JSON");
assert_no_secret_values(&value);
}
#[test]
fn registration_request_serializes_secret_values_but_redacts_debug_output() {
let mut value: Value =
serde_json::from_str(&fixture("remote_function_registration_request.json")).unwrap();
value["secret_values"] = serde_json::json!({"API_TOKEN": "secret-plaintext"});
let request = FunctionRegistrationRequest::from_json(&value.to_string()).unwrap();
assert_eq!(request.secret_values["API_TOKEN"], "secret-plaintext");
let canonical = request.to_canonical_json().unwrap();
assert!(canonical.contains("secret-plaintext"));
let debug = format!("{request:?}");
assert!(debug.contains("API_TOKEN"));
assert!(debug.contains("[REDACTED]"));
assert!(!debug.contains("secret-plaintext"));
}
#[tokio::test]
@@ -24,6 +24,7 @@
},
"runtime_digest": "sha256:runtime",
"environment_digest": "sha256:environment",
"required_secrets": ["HF_TOKEN"],
"created_at": "2026-08-21T00:00:00Z"
},
"future_job": {"trace_id": "trace-1"}
@@ -1 +1 @@
{"artifact":{"adapter":{"kind":"scalar_to_arrow_batch","version":1},"content":{"data":"ZnJvbSBfX2Z1dHVyZV9fIGltcG9ydCBhbm5vdGF0aW9ucwoKZGVmIG5vcm1hbGl6ZV9zY29yZSh2YWx1ZTogZmxvYXQpIC0+IGZsb2F0OgogICAgcmV0dXJuIHZhbHVlIC8gMTAwLjAK","encoding":"base64"},"digest":"sha256:760784bdcef57b802f389b97804cc0b618aae39e86044733451bcd0b13089a7f","entrypoint":"normalize_score","kind":"python_callable"},"name":"normalize_score","runtime":{"env":{"MODE":"test"},"environment":{"kind":"pip","packages":["numpy>=2"]},"kind":"python","python_version":"3.12"},"signature":{"inputs":[{"arrow_type":"float64","name":"value","nullable":false}],"output":{"arrow_type":"float64","kind":"scalar","nullable":false}}}
{"artifact":{"adapter":{"kind":"scalar_to_arrow_batch","version":1},"content":{"data":"ZnJvbSBfX2Z1dHVyZV9fIGltcG9ydCBhbm5vdGF0aW9ucwoKZGVmIG5vcm1hbGl6ZV9zY29yZSh2YWx1ZTogZmxvYXQpIC0+IGZsb2F0OgogICAgcmV0dXJuIHZhbHVlIC8gMTAwLjAK","encoding":"base64"},"digest":"sha256:760784bdcef57b802f389b97804cc0b618aae39e86044733451bcd0b13089a7f","entrypoint":"normalize_score","kind":"python_callable"},"name":"normalize_score","required_secrets":["API_TOKEN"],"runtime":{"env":{"MODE":"test"},"environment":{"kind":"pip","packages":["numpy>=2"]},"kind":"python","python_version":"3.12"},"signature":{"inputs":[{"arrow_type":"float64","name":"value","nullable":false}],"output":{"arrow_type":"float64","kind":"scalar","nullable":false}}}
@@ -39,5 +39,8 @@
"env": {
"MODE": "test"
}
}
},
"required_secrets": [
"API_TOKEN"
]
}
@@ -1 +1 @@
{"artifact":{"digest":"sha256:code","entrypoint":"embed","kind":"python_callable"},"created_at":"2026-08-21T00:00:00Z","environment_digest":"sha256:environment","name":"embed","runtime":{"env":{"TOKENIZERS_PARALLELISM":"false"},"environment":{"kind":"pip","packages":["sentence-transformers>=3"]},"kind":"python","python_version":"3.12"},"runtime_digest":"sha256:runtime","signature":{"inputs":[{"arrow_type":"utf8","name":"text","nullable":true}],"output":{"arrow_type":"list<float32>","kind":"scalar","nullable":false}},"version":"fv_01K3EXACT"}
{"artifact":{"digest":"sha256:code","entrypoint":"embed","kind":"python_callable"},"created_at":"2026-08-21T00:00:00Z","environment_digest":"sha256:environment","name":"embed","required_secrets":["HF_TOKEN"],"runtime":{"env":{"TOKENIZERS_PARALLELISM":"false"},"environment":{"kind":"pip","packages":["sentence-transformers>=3"]},"kind":"python","python_version":"3.12"},"runtime_digest":"sha256:runtime","signature":{"inputs":[{"arrow_type":"utf8","name":"text","nullable":true}],"output":{"arrow_type":"list<float32>","kind":"scalar","nullable":false}},"version":"fv_01K3EXACT"}