feat: declare Python function capabilities

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

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

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