From 76aac96749b56a864baceb3e68ab149cb0c0eb4b Mon Sep 17 00:00:00 2001 From: Xuanwo Date: Wed, 12 Aug 2026 03:47:17 +0800 Subject: [PATCH] feat: validate Python UDF source packages --- python/python/lancedb/_udf.py | 244 ++++++++- .../tests/test_first_class_udf_packaging.py | 486 ++++++++++++++++++ 2 files changed, 727 insertions(+), 3 deletions(-) create mode 100644 python/python/tests/test_first_class_udf_packaging.py diff --git a/python/python/lancedb/_udf.py b/python/python/lancedb/_udf.py index f6b8bcdca..1034fb524 100644 --- a/python/python/lancedb/_udf.py +++ b/python/python/lancedb/_udf.py @@ -3,16 +3,23 @@ """Local authoring declaration surface for first-class UDFs. -This module only snapshots declaration metadata onto a Python function. It does -not package source, mint durable identity, or register anything with a database. +This module snapshots declaration metadata onto a Python function and privately +validates packagable callables into a source snapshot. It does not mint durable +identity or register anything with a database. """ from __future__ import annotations +import ast import inspect +import stat +import symtable +import sys from collections.abc import Callable, Mapping, Sequence from dataclasses import dataclass -from typing import ParamSpec, TypeVar +from pathlib import Path +from types import CodeType, FunctionType +from typing import NoReturn, ParamSpec, TypeVar import pyarrow as pa @@ -22,6 +29,12 @@ _P = ParamSpec("_P") _R = TypeVar("_R") _CONFIG_ATTR = "__lancedb_udf_config__" +_SYNTHETIC_SOURCE_FILENAME = "" +_PACKAGING_ERROR = "udf is not packagable" +_ALLOWED_PARAM_KINDS = ( + inspect.Parameter.POSITIONAL_OR_KEYWORD, + inspect.Parameter.KEYWORD_ONLY, +) @dataclass(frozen=True, slots=True) @@ -35,6 +48,22 @@ class _UdfConfig: packages: tuple[str, ...] +@dataclass(frozen=True, slots=True) +class _PackagedUdf: + """Private frozen snapshot of a validated packagable UDF.""" + + source: str + module: str + callable_name: str + config: _UdfConfig + + def __repr__(self) -> str: + return ( + f"_PackagedUdf(source=, module={self.module!r}, " + f"callable_name={self.callable_name!r}, config={self.config!r})" + ) + + def _validate_inputs( inputs: object, ) -> tuple[tuple[str, pa.DataType], ...]: @@ -122,3 +151,212 @@ def _get_udf_config(fn: object) -> _UdfConfig: if not isinstance(config, _UdfConfig): raise TypeError("function is not decorated with @udf") return config + + +def _packaging_reject() -> NoReturn: + raise ValueError(_PACKAGING_ERROR) from None + + +def _is_ordinary_function(fn: FunctionType) -> bool: + if fn.__name__ == "": + return False + if fn.__qualname__ != fn.__name__: + return False + if inspect.iscoroutinefunction(fn) or inspect.isasyncgenfunction(fn): + return False + if inspect.isgeneratorfunction(fn): + return False + return True + + +def _resolve_source_path(fn: FunctionType, module: object) -> Path: + try: + fn_source: str | None = inspect.getsourcefile(fn) + except TypeError: + fn_source = None + source_lookup_failed = True + else: + source_lookup_failed = False + if source_lookup_failed: + _packaging_reject() + module_file = vars(module).get("__file__") + if not fn_source or not isinstance(module_file, str) or module_file == "": + _packaging_reject() + try: + resolved_paths: tuple[Path, Path] | None = ( + Path(fn_source).resolve(), + Path(module_file).resolve(), + ) + except (OSError, RuntimeError): + resolved_paths = None + if resolved_paths is None: + _packaging_reject() + fn_path, module_path = resolved_paths + if fn_path != module_path: + _packaging_reject() + if fn_path.suffix != ".py": + _packaging_reject() + try: + mode: int | None = fn_path.stat().st_mode + except OSError: + mode = None + if mode is None: + _packaging_reject() + if not stat.S_ISREG(mode): + _packaging_reject() + return fn_path + + +def _validate_source( + source: str, callable_name: str +) -> tuple[CodeType, symtable.SymbolTable]: + try: + module_code = compile( + source, + _SYNTHETIC_SOURCE_FILENAME, + "exec", + optimize=sys.flags.optimize, + ) + ast.parse(source, filename=_SYNTHETIC_SOURCE_FILENAME, mode="exec") + table = symtable.symtable(source, _SYNTHETIC_SOURCE_FILENAME, "exec") + parsed: tuple[CodeType, symtable.SymbolTable] | None = (module_code, table) + except Exception: + parsed = None + if parsed is None: + _packaging_reject() + module_code, table = parsed + + for child in table.get_children(): + if child.get_name() == callable_name and child.get_type() == "function": + return module_code, table + _packaging_reject() + + +def _source_bound_names(table: symtable.SymbolTable) -> set[str]: + names: set[str] = set() + for symbol in table.get_symbols(): + if symbol.is_imported() or symbol.is_assigned() or symbol.is_namespace(): + names.add(symbol.get_name()) + return names + + +def _code_fingerprint(code: CodeType) -> tuple[object, ...]: + """Structural fingerprint ignoring only location/debug fields.""" + constants = tuple( + _code_fingerprint(constant) if isinstance(constant, CodeType) else constant + for constant in code.co_consts + ) + return ( + code.co_name, + getattr(code, "co_qualname", code.co_name), + code.co_argcount, + code.co_posonlyargcount, + code.co_kwonlyargcount, + code.co_flags, + code.co_code, + code.co_names, + code.co_varnames, + code.co_freevars, + code.co_cellvars, + getattr(code, "co_exceptiontable", b""), + constants, + ) + + +def _toplevel_code_candidates( + module_code: CodeType, callable_name: str +) -> list[CodeType]: + candidates: list[CodeType] = [] + for constant in module_code.co_consts: + if not isinstance(constant, CodeType): + continue + if constant.co_name != callable_name: + continue + if getattr(constant, "co_qualname", callable_name) != callable_name: + continue + candidates.append(constant) + return candidates + + +def _validate_loaded_code_matches_source( + fn: FunctionType, module_code: CodeType +) -> None: + candidates = _toplevel_code_candidates(module_code, fn.__name__) + if not candidates: + _packaging_reject() + target = _code_fingerprint(fn.__code__) + if not any(_code_fingerprint(candidate) == target for candidate in candidates): + _packaging_reject() + + +def _validate_signature(fn: FunctionType, config: _UdfConfig) -> None: + try: + signature: inspect.Signature | None = inspect.signature(fn) + except (TypeError, ValueError): + signature = None + if signature is None: + _packaging_reject() + parameters = list(signature.parameters.values()) + expected = [name for name, _ in config.inputs] + actual = [parameter.name for parameter in parameters] + if actual != expected: + _packaging_reject() + for parameter in parameters: + if parameter.kind not in _ALLOWED_PARAM_KINDS: + _packaging_reject() + + +def _validate_ambient_globals(fn: FunctionType, table: symtable.SymbolTable) -> None: + try: + closure_vars: inspect.ClosureVars | None = inspect.getclosurevars(fn) + except (TypeError, ValueError): + closure_vars = None + if closure_vars is None: + _packaging_reject() + if closure_vars.nonlocals: + _packaging_reject() + bound_names = _source_bound_names(table) + for name in closure_vars.globals: + if name not in bound_names: + _packaging_reject() + + +def _package_udf(fn: object) -> _PackagedUdf: + """Validate and snapshot a packagable ``@udf``-decorated function.""" + config = _get_udf_config(fn) + if not isinstance(fn, FunctionType) or not _is_ordinary_function(fn): + _packaging_reject() + + module_name = fn.__module__ + if ( + not isinstance(module_name, str) + or module_name == "" + or module_name == "__main__" + ): + _packaging_reject() + module = sys.modules.get(module_name) + if module is None: + _packaging_reject() + callable_name = fn.__name__ + if vars(module).get(callable_name) is not fn: + _packaging_reject() + + source_path = _resolve_source_path(fn, module) + try: + source: str | None = source_path.read_text(encoding="utf-8") + except (OSError, UnicodeError): + source = None + if source is None: + _packaging_reject() + + module_code, table = _validate_source(source, callable_name) + _validate_signature(fn, config) + _validate_ambient_globals(fn, table) + _validate_loaded_code_matches_source(fn, module_code) + + return _PackagedUdf( + source=source, + module=module_name, + callable_name=callable_name, + config=config, + ) diff --git a/python/python/tests/test_first_class_udf_packaging.py b/python/python/tests/test_first_class_udf_packaging.py new file mode 100644 index 000000000..6e8107912 --- /dev/null +++ b/python/python/tests/test_first_class_udf_packaging.py @@ -0,0 +1,486 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright The LanceDB Authors + +"""RED contract tests for private UDF packaging validation.""" + +from __future__ import annotations + +import importlib +import inspect +import json +import sys +from collections.abc import Iterator +from contextlib import contextmanager +from pathlib import Path + +import pyarrow as pa +import pytest + +from lancedb import Function, Job, udf +from lancedb._udf import _get_udf_config, _package_udf + +_BODY_MARKER = "packaging body marker unique-xyz" +_AMBIENT_SECRET = "ambient-secret-value-xyz" +_BUILTIN_SHADOW_SECRET = "builtin-shadow-secret-xyz" +_SOURCE_MISMATCH_SECRET = "source-mismatch-secret-xyz" +_INVALID_UTF8_SECRET = "invalid-utf8-secret-xyz" + +_OVERDESIGN_ATTRS = ( + "user_version", + "idempotency_key", + "deterministic", + "null_handling", + "null_policy", + "on_error", + "error_policy", + "FunctionVersion", + "function_version", + "artifact", + "digest", + "geneva", + "id", + "function_id", + "job", + "job_id", + "registration", + "catalog", + "retry_key", + "source_path", + "path", + "function", +) + +_PACKAGING_CONSTANT = 41 + + +def _packaging_helper(value: int) -> int: + return value + _PACKAGING_CONSTANT + + +@udf( + inputs={"x": pa.int32()}, + output=pa.int64(), + python="3.12", + packages=["pkg-a==1"], + output_nullable=False, +) +def packable_add(x): + """packaging body marker unique-xyz.""" + return _packaging_helper(x) + len(json.dumps({"k": 1})) + + +@udf( + inputs={"x": pa.int32(), "y": pa.int32()}, + output=pa.int32(), + python="3.12", +) +def packable_kwonly(x, *, y=2): + return x + y + + +@udf( + inputs={"x": pa.int32()}, + output=pa.int32(), + python="3.12", +) +def packable_rebind_target(x): + return x + + +@udf( + inputs={"x": pa.int32()}, + output=pa.int32(), + python="3.12", +) +def uses_injected_ambient(x): + return x + len(INJECTED_AMBIENT_GLOBAL) # noqa: F821 + + +@udf( + inputs={"x": pa.int32()}, + output=pa.int32(), + python="3.12", +) +def uses_shadowed_builtin_len(x): + return x + len((1, 2, 3)) + + +@udf( + inputs={"x": pa.int32(), "y": pa.int32()}, + output=pa.int32(), + python="3.12", +) +def mismatch_names(left, right): + return left + right + + +@udf( + inputs={"y": pa.int32(), "x": pa.int32()}, + output=pa.int32(), + python="3.12", +) +def mismatch_order(x, y): + return x + y + + +@udf( + inputs={"x": pa.int32(), "y": pa.int32()}, + output=pa.int32(), + python="3.12", +) +def positional_only(x, /, y): + return x + y + + +@udf( + inputs={"x": pa.int32()}, + output=pa.int32(), + python="3.12", +) +def varargs_fn(x, *args): + return x + + +@udf( + inputs={"x": pa.int32()}, + output=pa.int32(), + python="3.12", +) +def kwargs_fn(x, **kwargs): + return x + + +@udf( + inputs={"x": pa.int32()}, + output=pa.int32(), + python="3.12", +) +async def async_fn(x): + return x + + +@udf( + inputs={"x": pa.int32()}, + output=pa.int32(), + python="3.12", +) +async def async_gen_fn(x): + yield x + + +@udf( + inputs={"x": pa.int32()}, + output=pa.int32(), + python="3.12", +) +def generator_fn(x): + yield x + + +def _assert_sanitized_text(*parts: object, secret: str = _AMBIENT_SECRET) -> None: + combined = "\n".join(str(part) for part in parts) + lowered = combined.lower() + assert _BODY_MARKER.lower() not in lowered + assert secret.lower() not in lowered + assert str(Path(__file__).resolve()).lower() not in lowered + assert Path(__file__).resolve().as_posix().lower() not in lowered + + +def _assert_packaging_rejection(exc_info, *, secret: str = _AMBIENT_SECRET) -> None: + _assert_sanitized_text(exc_info.value, repr(exc_info.value), secret=secret) + assert exc_info.value.__cause__ is None + assert exc_info.value.__context__ is None + + +@contextmanager +def _temporary_imported_module( + directory: Path, module_name: str, source: str +) -> Iterator[tuple[Path, object]]: + path = directory / f"{module_name}.py" + path.write_text(source, encoding="utf-8") + inserted = str(directory) + sys.path.insert(0, inserted) + try: + sys.modules.pop(module_name, None) + module = importlib.import_module(module_name) + yield path, module + finally: + sys.modules.pop(module_name, None) + try: + sys.path.remove(inserted) + except ValueError: + pass + + +def _temp_udf_module_source(*, body: str, secret: str | None = None) -> str: + secret_line = f"_SECRET = {secret!r}\n" if secret is not None else "" + return ( + "import pyarrow as pa\n" + "from lancedb import udf\n" + f"{secret_line}\n" + "@udf(\n" + ' inputs={"x": pa.int32()},\n' + " output=pa.int32(),\n" + ' python="3.12",\n' + ")\n" + "def temp_pack_target(x):\n" + f" {body}\n" + ) + + +def test_package_udf_success_snapshot_source_module_callable_config_and_repr(): + packaged = _package_udf(packable_add) + source = Path(__file__).read_text(encoding="utf-8") + + assert packaged.source == source + assert packaged.module == __name__ + assert packaged.module != "__main__" + assert packaged.callable_name == "packable_add" + assert packable_add.__qualname__ == "packable_add" + assert packaged.config is _get_udf_config(packable_add) + assert packaged.config.inputs == (("x", pa.int32()),) + assert packaged.config.output == pa.int64() + assert packaged.config.output_nullable is False + assert packaged.config.python == "3.12" + assert packaged.config.packages == ("pkg-a==1",) + + for attr in ("source", "module", "callable_name", "config"): + with pytest.raises(AttributeError): + setattr(packaged, attr, None) + + text = repr(packaged) + _assert_sanitized_text(text) + assert _BODY_MARKER not in text + + +def test_package_udf_allows_source_bound_import_constant_and_helper(): + packaged = _package_udf(packable_add) + assert packaged.callable_name == "packable_add" + assert "import json" in packaged.source + assert "_PACKAGING_CONSTANT" in packaged.source + assert "_packaging_helper" in packaged.source + assert packable_add(1) == _packaging_helper(1) + len(json.dumps({"k": 1})) + + +def test_package_udf_accepts_positional_or_keyword_and_keyword_only_defaults(): + packaged = _package_udf(packable_kwonly) + assert packaged.callable_name == "packable_kwonly" + assert packaged.config.inputs == (("x", pa.int32()), ("y", pa.int32())) + assert str(inspect.signature(packable_kwonly)) == "(x, *, y=2)" + assert packable_kwonly(3) == 5 + assert packable_kwonly(3, y=7) == 10 + + +def test_package_udf_rejects_lambda_and_closure(): + lam = udf( + inputs={"n": pa.int32()}, + output=pa.int32(), + python="3.12", + )(lambda n: n + 1) + with pytest.raises(ValueError) as exc_info: + _package_udf(lam) + _assert_packaging_rejection(exc_info) + + ambient = _AMBIENT_SECRET + + def factory(offset): + @udf( + inputs={"n": pa.int32()}, + output=pa.int32(), + python="3.12", + ) + def closed(n): + return n + offset + len(ambient) + + return closed + + closed = factory(10) + with pytest.raises(ValueError) as exc_info: + _package_udf(closed) + _assert_packaging_rejection(exc_info) + + def outer(): + total = 0 + + @udf( + inputs={"n": pa.int32()}, + output=pa.int32(), + python="3.12", + ) + def nested(n): + nonlocal total + total += n + return total + + return nested + + with pytest.raises(ValueError) as exc_info: + _package_udf(outer()) + _assert_packaging_rejection(exc_info) + + +def test_package_udf_rejects_signature_mismatches_and_unsupported_parameter_kinds(): + for target in ( + mismatch_names, + mismatch_order, + positional_only, + varargs_fn, + kwargs_fn, + ): + with pytest.raises(ValueError) as exc_info: + _package_udf(target) + _assert_packaging_rejection(exc_info) + + +def test_package_udf_rejects_async_and_generator_functions(): + for target in (async_fn, async_gen_fn, generator_fn): + with pytest.raises(ValueError) as exc_info: + _package_udf(target) + _assert_packaging_rejection(exc_info) + + +def test_package_udf_rejects_dynamic_exec_source(): + namespace: dict[str, object] = {} + exec( + "def dynamic_pack_target(x):\n return x + 1\n", + namespace, + ) + dynamic = udf( + inputs={"x": pa.int32()}, + output=pa.int32(), + python="3.12", + )(namespace["dynamic_pack_target"]) + with pytest.raises(ValueError) as exc_info: + _package_udf(dynamic) + _assert_packaging_rejection(exc_info) + + +def test_package_udf_rejects_undecorated_and_wrong_input_types(): + def plain(x): + return x + + with pytest.raises(TypeError) as exc_info: + _package_udf(plain) + _assert_packaging_rejection(exc_info) + + with pytest.raises(TypeError) as exc_info: + _package_udf(object()) + _assert_packaging_rejection(exc_info) + + with pytest.raises(TypeError) as exc_info: + _package_udf(42) + _assert_packaging_rejection(exc_info) + + +def test_package_udf_rejects_rebound_module_attribute(): + module = sys.modules[__name__] + original = module.packable_rebind_target + replacement = udf( + inputs={"x": pa.int32()}, + output=pa.int32(), + python="3.12", + )(lambda x: x) + module.packable_rebind_target = replacement + try: + with pytest.raises(ValueError) as exc_info: + _package_udf(original) + _assert_packaging_rejection(exc_info) + finally: + module.packable_rebind_target = original + + +def test_package_udf_rejects_injected_ambient_global(): + module = sys.modules[__name__] + secret = _AMBIENT_SECRET + module.INJECTED_AMBIENT_GLOBAL = secret + try: + assert uses_injected_ambient(3) == 3 + len(secret) + with pytest.raises(ValueError) as exc_info: + _package_udf(uses_injected_ambient) + _assert_packaging_rejection(exc_info, secret=secret) + finally: + delattr(module, "INJECTED_AMBIENT_GLOBAL") + + +def test_package_udf_rejects_builtin_shadow_injection(): + module = sys.modules[__name__] + secret = _BUILTIN_SHADOW_SECRET + assert not hasattr(module, "len") + module.len = secret + try: + with pytest.raises(ValueError) as exc_info: + _package_udf(uses_shadowed_builtin_len) + _assert_packaging_rejection(exc_info, secret=secret) + finally: + delattr(module, "len") + + +def test_package_udf_rejects_loaded_code_source_mismatch(tmp_path: Path): + secret = _SOURCE_MISMATCH_SECRET + module_name = "udf_pkg_source_mismatch_mod" + original = _temp_udf_module_source(body="return x + 1") + replacement = _temp_udf_module_source( + body=f"return x + 99 # {secret}", + secret=secret, + ) + with _temporary_imported_module(tmp_path, module_name, original) as ( + path, + module, + ): + target = module.temp_pack_target + assert target(1) == 2 + path.write_text(replacement, encoding="utf-8") + with pytest.raises(ValueError) as exc_info: + _package_udf(target) + _assert_packaging_rejection(exc_info, secret=secret) + err_text = f"{exc_info.value}\n{exc_info.value!r}" + assert str(path.resolve()) not in err_text + assert path.resolve().as_posix() not in err_text + + +def test_package_udf_rejects_invalid_utf8_after_import(tmp_path: Path): + secret = _INVALID_UTF8_SECRET + module_name = "udf_pkg_invalid_utf8_mod" + original = _temp_udf_module_source(body="return x + 1") + with _temporary_imported_module(tmp_path, module_name, original) as ( + path, + module, + ): + target = module.temp_pack_target + assert target(1) == 2 + path.write_bytes(secret.encode("utf-8") + b"\xff\xfe invalid-bytes") + with pytest.raises(ValueError) as exc_info: + _package_udf(target) + assert type(exc_info.value) is ValueError + assert exc_info.value.__cause__ is None + assert exc_info.value.__context__ is None + err_text = f"{exc_info.value}\n{exc_info.value!r}" + assert secret not in err_text + assert "b'" not in err_text + assert r"\xff" not in err_text + assert str(path.resolve()) not in err_text + assert path.resolve().as_posix() not in err_text + + +def test_package_udf_snapshot_has_no_durable_overdesign_and_is_not_function_or_job(): + packaged = _package_udf(packable_add) + assert not isinstance(packaged, Function) + assert not isinstance(packaged, Job) + for attr in _OVERDESIGN_ATTRS: + assert not hasattr(packaged, attr) + + text = repr(packaged).lower() + for token in ( + "user_version", + "idempotency_key", + "deterministic", + "null_policy", + "on_error", + "functionversion", + "artifact", + "digest", + "geneva", + "retry_key", + ): + assert token not in text + _assert_sanitized_text(text)