From f8bb90405ff70fa442b217d2a124e9965d0b11dd Mon Sep 17 00:00:00 2001 From: Xuanwo Date: Wed, 12 Aug 2026 03:57:52 +0800 Subject: [PATCH] feat: declare Python function capabilities --- python/python/lancedb/__init__.py | 3 +- python/python/lancedb/_udf.py | 128 ++++- .../tests/test_first_class_udf_capability.py | 453 ++++++++++++++++++ 3 files changed, 582 insertions(+), 2 deletions(-) create mode 100644 python/python/tests/test_first_class_udf_capability.py diff --git a/python/python/lancedb/__init__.py b/python/python/lancedb/__init__.py index 31c0bc377..60d8c212f 100644 --- a/python/python/lancedb/__init__.py +++ b/python/python/lancedb/__init__.py @@ -24,7 +24,7 @@ from .schema import blob, vector, BlobType from .job import AsyncJob, Job from .table import AsyncTable, Table from .types import BaseTokenizerType -from ._udf import udf +from ._udf import FunctionCapability, udf from ._lancedb import Session from .namespace import ( connect_namespace, @@ -510,6 +510,7 @@ __all__ = [ "col", "Expr", "Function", + "FunctionCapability", "func", "lit", "URI", diff --git a/python/python/lancedb/_udf.py b/python/python/lancedb/_udf.py index 1034fb524..181bfe250 100644 --- a/python/python/lancedb/_udf.py +++ b/python/python/lancedb/_udf.py @@ -23,7 +23,7 @@ from typing import NoReturn, ParamSpec, TypeVar import pyarrow as pa -__all__ = ["udf"] +__all__ = ["FunctionCapability", "udf"] _P = ParamSpec("_P") _R = TypeVar("_R") @@ -37,6 +37,111 @@ _ALLOWED_PARAM_KINDS = ( ) +class FunctionCapability: + """Local capability declaration for a first-class UDF. + + Construct via :meth:`network` or :meth:`secret`. Direct construction is + rejected so callers cannot create an uninitialized capability. + """ + + __slots__ = ("_kind", "_origin", "_reference", "_environment_variable") + + def __new__(cls, *args: object, **kwargs: object) -> FunctionCapability: + raise TypeError( + "FunctionCapability cannot be constructed directly; " + "use FunctionCapability.network() or FunctionCapability.secret()" + ) + + @classmethod + def _create( + cls, + kind: str, + origin: str | None, + reference: str | None, + environment_variable: str | None, + ) -> FunctionCapability: + obj = object.__new__(cls) + object.__setattr__(obj, "_kind", kind) + object.__setattr__(obj, "_origin", origin) + object.__setattr__(obj, "_reference", reference) + object.__setattr__(obj, "_environment_variable", environment_variable) + return obj + + @classmethod + def network(cls, origin: str) -> FunctionCapability: + if not isinstance(origin, str): + raise TypeError("origin must be a string") + if origin == "": + raise ValueError("origin must be non-empty") + return cls._create("network", origin, None, None) + + @classmethod + def secret(cls, reference: str, *, environment_variable: str) -> FunctionCapability: + if not isinstance(reference, str): + raise TypeError("reference must be a string") + if not isinstance(environment_variable, str): + raise TypeError("environment_variable must be a string") + if reference == "": + raise ValueError("reference must be non-empty") + if environment_variable == "": + raise ValueError("environment_variable must be non-empty") + return cls._create("secret", None, reference, environment_variable) + + @property + def kind(self) -> str: + return self._kind + + @property + def origin(self) -> str | None: + return self._origin + + @property + def reference(self) -> str | None: + return self._reference + + @property + def environment_variable(self) -> str | None: + return self._environment_variable + + def __setattr__(self, name: str, value: object) -> None: + raise AttributeError( + f"{type(self).__name__!r} object attribute {name!r} is read-only" + ) + + def __delattr__(self, name: str) -> None: + raise AttributeError( + f"{type(self).__name__!r} object attribute {name!r} is read-only" + ) + + def __eq__(self, other: object) -> bool: + if not isinstance(other, FunctionCapability): + return NotImplemented + return ( + self._kind == other._kind + and self._origin == other._origin + and self._reference == other._reference + and self._environment_variable == other._environment_variable + ) + + def __hash__(self) -> int: + return hash( + ( + self._kind, + self._origin, + self._reference, + self._environment_variable, + ) + ) + + def __repr__(self) -> str: + if self._kind == "network": + return f"FunctionCapability.network({self._origin!r})" + return ( + "FunctionCapability.secret(" + f"environment_variable={self._environment_variable!r})" + ) + + @dataclass(frozen=True, slots=True) class _UdfConfig: """Private frozen snapshot of a ``@udf`` declaration.""" @@ -46,6 +151,7 @@ class _UdfConfig: output_nullable: bool python: str packages: tuple[str, ...] + capabilities: tuple[FunctionCapability, ...] @dataclass(frozen=True, slots=True) @@ -100,6 +206,23 @@ def _validate_packages(packages: object) -> tuple[str, ...]: return tuple(snapshot) +def _validate_capabilities(capabilities: object) -> tuple[FunctionCapability, ...]: + if isinstance(capabilities, (str, bytes, bytearray)): + raise TypeError( + "udf capabilities must be a sequence of FunctionCapability, not a string" + ) + if not isinstance(capabilities, Sequence): + raise TypeError("udf capabilities must be a sequence of FunctionCapability") + snapshot: list[FunctionCapability] = [] + for capability in capabilities: + if not isinstance(capability, FunctionCapability): + raise TypeError( + "udf capabilities must contain only FunctionCapability values" + ) + snapshot.append(capability) + return tuple(snapshot) + + def udf( *, inputs: Mapping[str, pa.DataType], @@ -107,6 +230,7 @@ def udf( python: str, packages: Sequence[str] = (), output_nullable: bool = True, + capabilities: Sequence[FunctionCapability] = (), ) -> Callable[[Callable[_P, _R]], Callable[_P, _R]]: """Declare a local UDF without packaging or registration. @@ -123,6 +247,7 @@ def udf( package_snapshot = _validate_packages(packages) if not isinstance(output_nullable, bool): raise TypeError("udf output_nullable must be a bool") + capability_snapshot = _validate_capabilities(capabilities) config = _UdfConfig( inputs=input_snapshot, @@ -130,6 +255,7 @@ def udf( output_nullable=output_nullable, python=python, packages=package_snapshot, + capabilities=capability_snapshot, ) def decorator(fn: Callable[_P, _R]) -> Callable[_P, _R]: diff --git a/python/python/tests/test_first_class_udf_capability.py b/python/python/tests/test_first_class_udf_capability.py new file mode 100644 index 000000000..168f840f2 --- /dev/null +++ b/python/python/tests/test_first_class_udf_capability.py @@ -0,0 +1,453 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright The LanceDB Authors + +"""RED contract tests for local FunctionCapability authoring and @udf capabilities.""" + +from __future__ import annotations + +import inspect + +import pyarrow as pa +import pytest + +import lancedb +from lancedb import Function, FunctionCapability, Job, udf +from lancedb._udf import _get_udf_config, _package_udf + +_SECRET_REFERENCE = "secret://team/capability-redact-token-xyz" +_SECRET_ENV = "API_TOKEN" +_NETWORK_ORIGIN = "https://api.example.com" + +_OVERDESIGN_ATTRS = ( + "id", + "function_id", + "FunctionVersion", + "function_version", + "artifact", + "digest", + "authorization", + "authorized", + "value", + "plaintext", + "plaintext_secret", + "secret_value", + "job", + "job_id", + "catalog", + "retry_key", + "user_version", + "idempotency_key", + "deterministic", + "null_handling", + "null_policy", + "on_error", + "error_policy", + "geneva", +) + + +def _decorate(fn, **overrides): + kwargs = { + "inputs": {"x": pa.int32()}, + "output": pa.int64(), + "python": "3.12", + } + kwargs.update(overrides) + return udf(**kwargs)(fn) + + +def _network(origin: str = _NETWORK_ORIGIN) -> FunctionCapability: + return FunctionCapability.network(origin) + + +def _secret( + reference: str = _SECRET_REFERENCE, + *, + environment_variable: str = _SECRET_ENV, +) -> FunctionCapability: + return FunctionCapability.secret( + reference, + environment_variable=environment_variable, + ) + + +@udf( + inputs={"x": pa.int32()}, + output=pa.int64(), + python="3.12", + packages=["pkg-a==1"], + output_nullable=False, +) +def packable_without_capabilities(x): + return x + 1 + + +@udf( + inputs={"x": pa.int32()}, + output=pa.int64(), + python="3.12", + packages=["pkg-a==1"], + output_nullable=False, + capabilities=[ + FunctionCapability.network(_NETWORK_ORIGIN), + FunctionCapability.secret( + _SECRET_REFERENCE, + environment_variable=_SECRET_ENV, + ), + ], +) +def packable_with_capabilities(x): + return x + 1 + + +def test_function_capability_export_factories_projection_equality_immutability(): + assert "FunctionCapability" in lancedb.__all__ + assert FunctionCapability is lancedb.FunctionCapability + + network = _network() + secret = _secret() + + assert network.kind == "network" + assert network.origin == _NETWORK_ORIGIN + assert network.reference is None + assert network.environment_variable is None + + assert secret.kind == "secret" + assert secret.reference == _SECRET_REFERENCE + assert secret.environment_variable == _SECRET_ENV + assert secret.origin is None + + assert network == FunctionCapability.network(_NETWORK_ORIGIN) + assert secret == FunctionCapability.secret( + _SECRET_REFERENCE, + environment_variable=_SECRET_ENV, + ) + assert network != secret + assert network != FunctionCapability.network("https://other.example.com") + assert secret != FunctionCapability.secret( + _SECRET_REFERENCE, + environment_variable="OTHER_TOKEN", + ) + + public_attrs = ("kind", "origin", "reference", "environment_variable") + internal_slots = ("_kind", "_origin", "_reference", "_environment_variable") + immutable_attrs = public_attrs + internal_slots + + for attr in public_attrs: + with pytest.raises(AttributeError): + setattr(network, attr, None) + with pytest.raises(AttributeError): + setattr(secret, attr, None) + + for attr in immutable_attrs: + # Fresh instances per attempt so a RED slot mutation cannot corrupt + # shared fixtures used by later assertions in this test. + fresh_network = _network("https://fresh-immutability.example.com") + fresh_secret = _secret( + "secret://team/fresh-immutability-token", + environment_variable="FRESH_IMMUTABILITY_TOKEN", + ) + with pytest.raises(AttributeError): + setattr(fresh_network, attr, None) + with pytest.raises(AttributeError): + setattr(fresh_secret, attr, None) + with pytest.raises(AttributeError): + delattr(fresh_network, attr) + with pytest.raises(AttributeError): + delattr(fresh_secret, attr) + + retained_origin = "https://config-retain.example.com" + retained_reference = "secret://team/config-retain-token" + retained_env = "CONFIG_RETAIN_TOKEN" + retained_network = FunctionCapability.network(retained_origin) + retained_secret = FunctionCapability.secret( + retained_reference, + environment_variable=retained_env, + ) + expected_capabilities = ( + FunctionCapability.network(retained_origin), + FunctionCapability.secret( + retained_reference, + environment_variable=retained_env, + ), + ) + + def retain_target(x): + return x + + retained = udf( + inputs={"x": pa.int32()}, + output=pa.int32(), + python="3.12", + capabilities=[retained_network, retained_secret], + )(retain_target) + retained_config = _get_udf_config(retained) + assert retained_config.capabilities == expected_capabilities + + for attr in immutable_attrs: + with pytest.raises(AttributeError): + setattr(retained_network, attr, "mutated") + with pytest.raises(AttributeError): + setattr(retained_secret, attr, "mutated") + with pytest.raises(AttributeError): + delattr(retained_network, attr) + with pytest.raises(AttributeError): + delattr(retained_secret, attr) + + assert retained_config.capabilities == expected_capabilities + assert retained_config.capabilities[0] is retained_network + assert retained_config.capabilities[1] is retained_secret + assert retained_config.capabilities[0].kind == "network" + assert retained_config.capabilities[0].origin == retained_origin + assert retained_config.capabilities[0].reference is None + assert retained_config.capabilities[0].environment_variable is None + assert retained_config.capabilities[1].kind == "secret" + assert retained_config.capabilities[1].reference == retained_reference + assert retained_config.capabilities[1].environment_variable == retained_env + assert retained_config.capabilities[1].origin is None + + with pytest.raises(TypeError): + FunctionCapability() + with pytest.raises(TypeError): + FunctionCapability( # type: ignore[call-arg] + kind="network", + origin=_NETWORK_ORIGIN, + ) + + assert not isinstance(network, Function) + assert not isinstance(secret, Function) + assert not isinstance(network, Job) + assert not isinstance(secret, Job) + for attr in _OVERDESIGN_ATTRS: + assert not hasattr(network, attr) + assert not hasattr(secret, attr) + + +def test_function_capability_validation_and_secret_redaction(): + with pytest.raises(TypeError): + FunctionCapability.network(None) # type: ignore[arg-type] + with pytest.raises(TypeError): + FunctionCapability.network(123) # type: ignore[arg-type] + with pytest.raises(ValueError): + FunctionCapability.network("") + + # Backend authorization owns URL/scheme policy; non-empty is enough here. + loose = FunctionCapability.network("example.com") + assert loose.kind == "network" + assert loose.origin == "example.com" + + with pytest.raises(TypeError): + FunctionCapability.secret( # type: ignore[misc] + _SECRET_REFERENCE, + _SECRET_ENV, + ) + with pytest.raises(TypeError): + FunctionCapability.secret(None, environment_variable=_SECRET_ENV) # type: ignore[arg-type] + with pytest.raises(TypeError): + FunctionCapability.secret(123, environment_variable=_SECRET_ENV) # type: ignore[arg-type] + with pytest.raises(TypeError): + FunctionCapability.secret(_SECRET_REFERENCE, environment_variable=None) # type: ignore[arg-type] + with pytest.raises(TypeError): + FunctionCapability.secret(_SECRET_REFERENCE, environment_variable=1) # type: ignore[arg-type] + + with pytest.raises(ValueError) as empty_ref: + FunctionCapability.secret("", environment_variable=_SECRET_ENV) + assert _SECRET_REFERENCE not in str(empty_ref.value) + assert _SECRET_REFERENCE not in repr(empty_ref.value) + + with pytest.raises(ValueError) as empty_env: + FunctionCapability.secret(_SECRET_REFERENCE, environment_variable="") + assert _SECRET_REFERENCE not in str(empty_env.value) + assert _SECRET_REFERENCE not in repr(empty_env.value) + + with pytest.raises(TypeError): + FunctionCapability.secret( # type: ignore[call-arg] + _SECRET_REFERENCE, + environment_variable=_SECRET_ENV, + value="super-secret", + ) + with pytest.raises(TypeError): + FunctionCapability.secret( # type: ignore[call-arg] + _SECRET_REFERENCE, + environment_variable=_SECRET_ENV, + plaintext_secret="super-secret", + ) + with pytest.raises(TypeError): + FunctionCapability.secret( # type: ignore[call-arg] + _SECRET_REFERENCE, + environment_variable=_SECRET_ENV, + environment={_SECRET_ENV: "super-secret"}, + ) + with pytest.raises(TypeError): + FunctionCapability.secret( # type: ignore[call-arg] + _SECRET_REFERENCE, + environment_variable=_SECRET_ENV, + headers={"Authorization": "Bearer super-secret"}, + ) + with pytest.raises(TypeError): + FunctionCapability.network( # type: ignore[call-arg] + _NETWORK_ORIGIN, + headers={"X-Trace": "1"}, + ) + + secret = _secret() + assert not hasattr(secret, "value") + assert not hasattr(secret, "plaintext") + assert not hasattr(secret, "plaintext_secret") + assert not hasattr(secret, "secret_value") + + secret_text = repr(secret) + assert "secret" in secret_text.lower() + assert _SECRET_ENV in secret_text + assert _SECRET_REFERENCE not in secret_text + assert "super-secret" not in secret_text + + network_text = repr(_network()) + assert "network" in network_text.lower() + assert _NETWORK_ORIGIN in network_text + + +def test_udf_capabilities_ordered_immutable_config_default_and_validation(): + params = inspect.signature(udf).parameters + assert "capabilities" in params + assert params["capabilities"].kind is inspect.Parameter.KEYWORD_ONLY + assert params["capabilities"].default == () + + def identity_target(x): + """capabilities identity marker.""" + return x + 1 + + original = identity_target + decorated = _decorate(identity_target) + assert decorated is original + assert decorated.__name__ == "identity_target" + assert decorated.__doc__ == "capabilities identity marker." + assert decorated(2) == 3 + assert _get_udf_config(decorated).capabilities == () + + first = _network("https://b.example.com") + second = _network("https://a.example.com") + third = _network("https://b.example.com") + secret = _secret() + capabilities = [first, second, third, secret] + + def combine(x): + return x + + with_caps = udf( + inputs={"x": pa.int32()}, + output=pa.int32(), + python="3.12", + packages=["pkg-b==2", "pkg-a==1"], + capabilities=capabilities, + )(combine) + config = _get_udf_config(with_caps) + assert config.capabilities == (first, second, third, secret) + assert isinstance(config.capabilities, tuple) + assert config.packages == ("pkg-b==2", "pkg-a==1") + assert config.inputs == (("x", pa.int32()),) + + capabilities.append(_network("https://mutated.example.com")) + capabilities[0] = _network("https://replaced.example.com") + assert config.capabilities == (first, second, third, secret) + + with pytest.raises(AttributeError): + setattr(config, "capabilities", ()) + + def target(x): + return x + + with pytest.raises(TypeError): + _decorate(target, capabilities="https://api.example.com") + + with pytest.raises(TypeError): + _decorate(target, capabilities=b"https://api.example.com") + + class _BadCapability: + def __repr__(self) -> str: + return "unique-bad-capability-repr-xyz" + + with pytest.raises(TypeError) as bad_item: + _decorate(target, capabilities=[_BadCapability()]) + assert "unique-bad-capability-repr-xyz" not in str(bad_item.value) + assert "unique-bad-capability-repr-xyz" not in repr(bad_item.value) + + with pytest.raises(TypeError) as bad_mixed: + _decorate( + target, + capabilities=[_network(), "unique-bad-capability-string-xyz"], + ) + assert "unique-bad-capability-string-xyz" not in str(bad_mixed.value) + assert "unique-bad-capability-string-xyz" not in repr(bad_mixed.value) + + +def test_package_udf_preserves_capabilities_and_redacts_secret_reference(): + packaged = _package_udf(packable_with_capabilities) + config = packaged.config + + assert packaged.config is _get_udf_config(packable_with_capabilities) + assert config.capabilities == ( + FunctionCapability.network(_NETWORK_ORIGIN), + FunctionCapability.secret( + _SECRET_REFERENCE, + environment_variable=_SECRET_ENV, + ), + ) + assert config.capabilities[0].kind == "network" + assert config.capabilities[0].origin == _NETWORK_ORIGIN + assert config.capabilities[1].kind == "secret" + assert config.capabilities[1].reference == _SECRET_REFERENCE + assert config.capabilities[1].environment_variable == _SECRET_ENV + assert config.packages == ("pkg-a==1",) + assert config.python == "3.12" + assert config.output_nullable is False + + nested = ( + f"{packaged!r}\n{config!r}\n{config.capabilities!r}\n{config.capabilities[1]!r}" + ) + assert _SECRET_REFERENCE not in nested + assert _SECRET_ENV in repr(config.capabilities[1]) + + +def test_capabilities_are_additive_to_existing_declaration_and_packaging(): + def score(x): + return x + + decorated = _decorate( + score, + packages=["score==1.0"], + output_nullable=True, + ) + config = _get_udf_config(decorated) + assert config.inputs == (("x", pa.int32()),) + assert config.output == pa.int64() + assert config.output_nullable is True + assert config.python == "3.12" + assert config.packages == ("score==1.0",) + assert config.capabilities == () + assert decorated is score + assert decorated(4) == 4 + + packaged = _package_udf(packable_without_capabilities) + assert packaged.config is _get_udf_config(packable_without_capabilities) + assert packaged.callable_name == "packable_without_capabilities" + assert packaged.config.capabilities == () + assert packaged.config.packages == ("pkg-a==1",) + assert packaged.config.output_nullable is False + assert packable_without_capabilities(1) == 2 + + params = inspect.signature(udf).parameters + for name in ( + "user_version", + "idempotency_key", + "deterministic", + "null_handling", + "null_policy", + "on_error", + "error_policy", + "FunctionVersion", + "artifact", + "digest", + "geneva", + ): + assert name not in params