From 983943181a10ba29fea72929c249e56603be36fd Mon Sep 17 00:00:00 2001 From: kaka11chen Date: Mon, 17 Aug 2026 19:41:02 +0800 Subject: [PATCH 1/4] [feat](build) Build self-contained loadable extension artifacts --- CMakeLists.txt | 1 + DEVELOPMENT.md | 24 +++ cmake/duckdb_loader.cmake | 157 +++++++++++++++++- .../fast/test_loadable_extension_artifacts.py | 47 ++++++ 4 files changed, 227 insertions(+), 2 deletions(-) create mode 100644 tests/fast/test_loadable_extension_artifacts.py diff --git a/CMakeLists.txt b/CMakeLists.txt index 255bd10b66e..1b06687d44a 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -189,6 +189,7 @@ pybind11_add_module( # Add Vane's Python and engine dependencies. target_link_libraries(_native PRIVATE vane_python_dependencies) duckdb_link_extensions(_native) +duckdb_stage_loadable_extensions() # ──────────────────────────────────────────── # Controlling symbol export diff --git a/DEVELOPMENT.md b/DEVELOPMENT.md index 8887512d80c..0b3d34e097c 100644 --- a/DEVELOPMENT.md +++ b/DEVELOPMENT.md @@ -44,6 +44,30 @@ Python-only changes do not require a native rebuild, but reinstall the non-editable package so the test environment receives them. Changes below `src/vane_py/` or `external/duckdb/src/` require an incremental native build. +## Building a loadable extension artifact + +`VANE_LOADABLE_EXTENSIONS` builds selected DuckDB extensions as self-contained +`.duckdb_extension` artifacts without linking them into `vane._native`. The +default is empty, so base Vane builds and wheels do not contain staged optional +extensions. For example, build and exercise the in-tree `tpch` artifact: + +```bash +export SKBUILD_BUILD_DIR="$PWD/build/python-release" +export SKBUILD_CMAKE_BUILD_TYPE=Release +uv pip install . --no-build-isolation \ + -Ccmake.define.VANE_LOADABLE_EXTENSIONS=tpch +cmake --build "$SKBUILD_BUILD_DIR" --target vane_loadable_extensions +VANE_TEST_LOADABLE_EXTENSION_PATH=\ +"$SKBUILD_BUILD_DIR/vane_extensions/tpch.duckdb_extension" \ + scripts/run_installed_pytest.sh tests/fast/test_loadable_extension_artifacts.py +``` + +Loadable artifacts require `EXTENSION_STATIC_BUILD=ON`. This keeps each +artifact self-contained and preserves Vane's private `_native` symbol boundary. +The staging directory is configurable with +`VANE_LOADABLE_EXTENSION_OUTPUT_DIRECTORY`; packaging and trusted artifact +metadata are intentionally handled separately. + ## Native C++ tests The complete native gate builds DuckDB, distributed exchange, and the test diff --git a/cmake/duckdb_loader.cmake b/cmake/duckdb_loader.cmake index 10861ea8482..1d8c30d49e9 100644 --- a/cmake/duckdb_loader.cmake +++ b/cmake/duckdb_loader.cmake @@ -44,6 +44,13 @@ _duckdb_set_default(DUCKDB_SOURCE_PATH # Extension list - commonly used extensions for Python _duckdb_set_default(BUILD_EXTENSIONS "core_functions;parquet;icu;json;httpfs") +# Optional extensions that are built as self-contained DuckDB loadable +# artifacts. They are deliberately configured with DONT_LINK so _native keeps +# its existing symbol-isolated static extension set. +_duckdb_set_default(VANE_LOADABLE_EXTENSIONS "") +_duckdb_set_default(VANE_LOADABLE_EXTENSION_OUTPUT_DIRECTORY + "${CMAKE_BINARY_DIR}/vane_extensions") + # Core build options - disable unnecessary components for Python builds _duckdb_set_default(BUILD_SHELL OFF) _duckdb_set_default(BUILD_UNITTESTS OFF) @@ -74,6 +81,14 @@ set(DUCKDB_SOURCE_PATH set(BUILD_EXTENSIONS "${BUILD_EXTENSIONS}" CACHE STRING "Semicolon-separated list of extensions to enable") +set(VANE_LOADABLE_EXTENSIONS + "${VANE_LOADABLE_EXTENSIONS}" + CACHE + STRING + "Semicolon-separated list of self-contained loadable extensions to build") +set(VANE_LOADABLE_EXTENSION_OUTPUT_DIRECTORY + "${VANE_LOADABLE_EXTENSION_OUTPUT_DIRECTORY}" + CACHE PATH "Directory used to stage Vane loadable extension artifacts") set(BUILD_SHELL "${BUILD_SHELL}" CACHE BOOL "Build the DuckDB shell executable") @@ -591,6 +606,89 @@ function(_duckdb_create_interface_target target_name) PROPERTIES INTERFACE_POSITION_INDEPENDENT_CODE ON) endfunction() +function(_duckdb_configure_loadable_extensions) + set(_VANE_LOADABLE_EXTENSION_NAMES) + foreach(_VANE_REQUESTED_EXTENSION IN LISTS VANE_LOADABLE_EXTENSIONS) + string(TOLOWER "${_VANE_REQUESTED_EXTENSION}" _VANE_LOADABLE_EXTENSION_NAME) + if(_VANE_LOADABLE_EXTENSION_NAME STREQUAL "") + continue() + endif() + if(NOT _VANE_LOADABLE_EXTENSION_NAME MATCHES "^[a-z][a-z0-9_]*$") + message( + FATAL_ERROR + "Invalid VANE_LOADABLE_EXTENSIONS entry '${_VANE_REQUESTED_EXTENSION}'. " + "Extension names must contain lowercase letters, digits, and underscores." + ) + endif() + list(APPEND _VANE_LOADABLE_EXTENSION_NAMES + "${_VANE_LOADABLE_EXTENSION_NAME}") + endforeach() + list(REMOVE_DUPLICATES _VANE_LOADABLE_EXTENSION_NAMES) + + # DuckDB processes BUILD_EXTENSIONS before DUCKDB_EXTENSION_CONFIGS. Remove + # selected artifacts first so the generated DONT_LINK configuration below owns + # their registration even when they are part of Vane's base build list. + foreach(_VANE_LOADABLE_EXTENSION_NAME IN LISTS _VANE_LOADABLE_EXTENSION_NAMES) + list(REMOVE_ITEM BUILD_EXTENSIONS "${_VANE_LOADABLE_EXTENSION_NAME}") + endforeach() + + if(_VANE_LOADABLE_EXTENSION_NAMES AND NOT DEFINED EXTENSION_STATIC_BUILD) + # DuckDB defaults this option to ON, but it reads it before declaring the + # option. Set it only for a requested Vane artifact so normal Vane builds + # retain their existing configuration order. + set(EXTENSION_STATIC_BUILD + ON + CACHE BOOL + "Build loadable extensions with a statically linked DuckDB engine" + ) + endif() + if(_VANE_LOADABLE_EXTENSION_NAMES AND NOT EXTENSION_STATIC_BUILD) + message( + FATAL_ERROR + "VANE_LOADABLE_EXTENSIONS requires EXTENSION_STATIC_BUILD=ON. Thin " + "extensions cannot resolve DuckDB symbols from Vane's private _native module." + ) + endif() + + if(_VANE_LOADABLE_EXTENSION_NAMES) + set(_VANE_LOADABLE_EXTENSION_CONFIG_CONTENT + "# Generated by cmake/duckdb_loader.cmake.\n") + foreach(_VANE_LOADABLE_EXTENSION_NAME IN + LISTS _VANE_LOADABLE_EXTENSION_NAMES) + string( + APPEND _VANE_LOADABLE_EXTENSION_CONFIG_CONTENT + "duckdb_extension_load(${_VANE_LOADABLE_EXTENSION_NAME} DONT_LINK)\n") + endforeach() + + set(_VANE_LOADABLE_EXTENSION_CONFIG_DIRECTORY + "${CMAKE_BINARY_DIR}/generated") + set(_VANE_LOADABLE_EXTENSION_CONFIG + "${_VANE_LOADABLE_EXTENSION_CONFIG_DIRECTORY}/vane_loadable_extensions.cmake" + ) + file(MAKE_DIRECTORY "${_VANE_LOADABLE_EXTENSION_CONFIG_DIRECTORY}") + file( + CONFIGURE + OUTPUT + "${_VANE_LOADABLE_EXTENSION_CONFIG}" + CONTENT + "${_VANE_LOADABLE_EXTENSION_CONFIG_CONTENT}" + @ONLY + NEWLINE_STYLE + UNIX) + list(PREPEND DUCKDB_EXTENSION_CONFIGS "${_VANE_LOADABLE_EXTENSION_CONFIG}") + endif() + + set(VANE_LOADABLE_EXTENSION_NAMES + "${_VANE_LOADABLE_EXTENSION_NAMES}" + PARENT_SCOPE) + set(BUILD_EXTENSIONS + "${BUILD_EXTENSIONS}" + PARENT_SCOPE) + set(DUCKDB_EXTENSION_CONFIGS + "${DUCKDB_EXTENSION_CONFIGS}" + PARENT_SCOPE) +endfunction() + function(_duckdb_print_summary) message(STATUS "DuckDB Configuration:") message(STATUS " Source: ${DUCKDB_SOURCE_PATH}") @@ -602,6 +700,14 @@ function(_duckdb_print_summary) message(STATUS " Build Type: ${CMAKE_BUILD_TYPE}") message(STATUS " Native Arch: ${NATIVE_ARCH}") message(STATUS " Unity Build Disabled: ${DISABLE_UNITY}") + if(VANE_LOADABLE_EXTENSION_NAMES) + message( + STATUS " Vane Loadable Extensions: ${VANE_LOADABLE_EXTENSION_NAMES}") + message( + STATUS + " Vane Loadable Extension Output: ${VANE_LOADABLE_EXTENSION_OUTPUT_DIRECTORY}" + ) + endif() set(debug_opts) if(FORCE_ASSERT) @@ -621,6 +727,7 @@ endfunction() # ════════════════════════════════════════════════════════════════════════════════ function(duckdb_add_library target_name) + _duckdb_configure_loadable_extensions() _duckdb_validate_source_path() _duckdb_resolve_source_id() _duckdb_resolve_fork_version() @@ -650,6 +757,9 @@ function(duckdb_add_library target_name) set(BUILD_EXTENSIONS "${BUILD_EXTENSIONS}" PARENT_SCOPE) + set(VANE_LOADABLE_EXTENSION_NAMES + "${VANE_LOADABLE_EXTENSION_NAMES}" + PARENT_SCOPE) endfunction() function(duckdb_link_extensions target_name) @@ -660,9 +770,13 @@ function(duckdb_link_extensions target_name) target_link_libraries( ${target_name} PRIVATE "$") - if(BUILD_EXTENSIONS) + set(_VANE_STATIC_EXTENSIONS ${BUILD_EXTENSIONS}) + foreach(_VANE_LOADABLE_EXTENSION_NAME IN LISTS VANE_LOADABLE_EXTENSION_NAMES) + list(REMOVE_ITEM _VANE_STATIC_EXTENSIONS "${_VANE_LOADABLE_EXTENSION_NAME}") + endforeach() + if(_VANE_STATIC_EXTENSIONS) message(STATUS "Linking DuckDB extensions:") - foreach(ext IN LISTS BUILD_EXTENSIONS) + foreach(ext IN LISTS _VANE_STATIC_EXTENSIONS) message(STATUS "- ${ext}") target_link_libraries(${target_name} PRIVATE ${ext}_extension) endforeach() @@ -671,6 +785,45 @@ function(duckdb_link_extensions target_name) endif() endfunction() +function(duckdb_stage_loadable_extensions) + if(NOT VANE_LOADABLE_EXTENSION_NAMES) + return() + endif() + + add_custom_target(vane_loadable_extensions) + foreach(_VANE_LOADABLE_EXTENSION_NAME IN LISTS VANE_LOADABLE_EXTENSION_NAMES) + set(_VANE_LOADABLE_EXTENSION_TARGET + "${_VANE_LOADABLE_EXTENSION_NAME}_loadable_extension") + if(NOT TARGET "${_VANE_LOADABLE_EXTENSION_TARGET}") + message( + FATAL_ERROR + "VANE_LOADABLE_EXTENSIONS requested '${_VANE_LOADABLE_EXTENSION_NAME}', " + "but DuckDB did not create target '${_VANE_LOADABLE_EXTENSION_TARGET}'." + ) + endif() + + set(_VANE_STAGED_LOADABLE_EXTENSION + "${VANE_LOADABLE_EXTENSION_OUTPUT_DIRECTORY}/${_VANE_LOADABLE_EXTENSION_NAME}.duckdb_extension" + ) + add_custom_command( + OUTPUT "${_VANE_STAGED_LOADABLE_EXTENSION}" + COMMAND ${CMAKE_COMMAND} -E make_directory + "${VANE_LOADABLE_EXTENSION_OUTPUT_DIRECTORY}" + COMMAND + ${CMAKE_COMMAND} -E copy_if_different + "$" + "${_VANE_STAGED_LOADABLE_EXTENSION}" + DEPENDS "${_VANE_LOADABLE_EXTENSION_TARGET}" + COMMENT "Staging Vane loadable extension ${_VANE_LOADABLE_EXTENSION_NAME}" + VERBATIM) + + add_custom_target("vane_loadable_extension_${_VANE_LOADABLE_EXTENSION_NAME}" + DEPENDS "${_VANE_STAGED_LOADABLE_EXTENSION}") + add_dependencies(vane_loadable_extensions + "vane_loadable_extension_${_VANE_LOADABLE_EXTENSION_NAME}") + endforeach() +endfunction() + # ════════════════════════════════════════════════════════════════════════════════ # Convenience Functions # ════════════════════════════════════════════════════════════════════════════════ diff --git a/tests/fast/test_loadable_extension_artifacts.py b/tests/fast/test_loadable_extension_artifacts.py new file mode 100644 index 00000000000..1aa1ff43026 --- /dev/null +++ b/tests/fast/test_loadable_extension_artifacts.py @@ -0,0 +1,47 @@ +# SPDX-FileCopyrightText: 2026 Vane contributors +# SPDX-License-Identifier: Apache-2.0 + +import os +from pathlib import Path + +import pytest + +import vane + + +@pytest.fixture(scope="module") +def loadable_extension_path() -> Path: + configured_path = os.environ.get("VANE_TEST_LOADABLE_EXTENSION_PATH") + if configured_path is None: + pytest.skip("set VANE_TEST_LOADABLE_EXTENSION_PATH to test a staged artifact") + + path = Path(configured_path).resolve() + assert path.is_file(), f"loadable extension artifact does not exist: {path}" + return path + + +def test_staged_tpch_extension_loads_without_static_linkage(loadable_extension_path: Path): + connection = vane.connect(config={"allow_unsigned_extensions": "true"}) + try: + initial_state = connection.execute( + """ + SELECT loaded, installed, install_mode + FROM duckdb_extensions() + WHERE extension_name = 'tpch' + """ + ).fetchone() + assert initial_state == (False, False, "NOT_INSTALLED") + + connection.load_extension(str(loadable_extension_path)) + + loaded_state = connection.execute( + """ + SELECT loaded, installed, install_mode + FROM duckdb_extensions() + WHERE extension_name = 'tpch' + """ + ).fetchone() + assert loaded_state == (True, False, "NOT_INSTALLED") + assert connection.execute("SELECT count(*) FROM tpch_queries()").fetchone() == (22,) + finally: + connection.close() From d617fe79d29574c59acc3015671b304c3aa6df77 Mon Sep 17 00:00:00 2001 From: kaka11chen Date: Mon, 17 Aug 2026 20:15:27 +0800 Subject: [PATCH 2/4] [feat](extensions) Add trusted dynamic-extension descriptors --- tests/fast/test_dynamic_extension_resolver.py | 309 ++++++++ vane/__init__.py | 19 + vane/extensions.py | 674 ++++++++++++++++++ 3 files changed, 1002 insertions(+) create mode 100644 tests/fast/test_dynamic_extension_resolver.py create mode 100644 vane/extensions.py diff --git a/tests/fast/test_dynamic_extension_resolver.py b/tests/fast/test_dynamic_extension_resolver.py new file mode 100644 index 00000000000..f0043216bce --- /dev/null +++ b/tests/fast/test_dynamic_extension_resolver.py @@ -0,0 +1,309 @@ +# SPDX-FileCopyrightText: 2026 Vane contributors +# SPDX-License-Identifier: Apache-2.0 + +import os +from dataclasses import replace +from pathlib import Path + +import pytest + +import vane +from vane.extensions import ( + DynamicExtensionDependency, + DynamicExtensionDescriptor, + DynamicExtensionError, + DynamicExtensionResolver, + LocalExtensionArtifact, + LocalExtensionProvider, + create_dynamic_extension_descriptor, +) + + +class _Result: + def __init__(self, row): + self._row = row + + def fetchone(self): + return self._row + + +class RecordingConnection: + def __init__(self, platform): + self.platform = platform + self.loaded_paths = [] + + def execute(self, query): + assert query == "SELECT platform FROM pragma_platform()" + return _Result((self.platform,)) + + def load_extension(self, extension): + self.loaded_paths.append(Path(extension)) + + +def _runtime_platform(): + connection = vane.connect() + try: + return connection.execute("SELECT platform FROM pragma_platform()").fetchone()[0] + finally: + connection.close() + + +def _write_extension_artifact(path, *, platform, source_id, extension_version="test-version", abi_type="CPP"): + path.parent.mkdir(parents=True, exist_ok=True) + footer = bytearray(512) + fields = ["", "", "", abi_type, extension_version, source_id, platform, "4"] + for index, value in enumerate(fields): + start = index * 32 + footer[start : start + len(value)] = value.encode("ascii") + path.write_bytes(b"Vane extension test payload" + footer) + return path + + +def _descriptor(path, *, name, trust_identity="local-tests"): + return create_dynamic_extension_descriptor(path, name=name, trust_identity=trust_identity) + + +def _resolver(*artifacts, trust_identity="local-tests"): + provider = LocalExtensionProvider(trust_identity, artifacts) + return DynamicExtensionResolver(trusted_identities={trust_identity}, providers=[provider]) + + +def test_descriptor_round_trip_preserves_ordered_dependency_identity(tmp_path): + platform = _runtime_platform() + dependency_path = _write_extension_artifact( + tmp_path / "dependency.duckdb_extension", + platform=platform, + source_id=vane.__git_revision__, + ) + root_path = _write_extension_artifact( + tmp_path / "root.duckdb_extension", + platform=platform, + source_id=vane.__git_revision__, + extension_version="root-version", + ) + dependency = _descriptor(dependency_path, name="dependency") + root = replace( + _descriptor(root_path, name="root"), + dependencies=( + DynamicExtensionDependency( + name=dependency.name, + extension_version=dependency.extension_version, + sha256=dependency.sha256, + ), + ), + ) + + restored = DynamicExtensionDescriptor.from_json(root.to_json()) + + assert restored == root + assert restored.dependencies[0].identity == dependency.identity + assert restored.to_json() == root.to_json() + + +def test_resolver_loads_dependencies_before_root_and_caches_digest(tmp_path): + platform = _runtime_platform() + dependency_path = _write_extension_artifact( + tmp_path / "dependency.duckdb_extension", + platform=platform, + source_id=vane.__git_revision__, + ) + root_path = _write_extension_artifact( + tmp_path / "root.duckdb_extension", + platform=platform, + source_id=vane.__git_revision__, + extension_version="root-version", + ) + dependency = _descriptor(dependency_path, name="dependency") + root = replace( + _descriptor(root_path, name="root"), + dependencies=( + DynamicExtensionDependency( + name=dependency.name, + extension_version=dependency.extension_version, + sha256=dependency.sha256, + ), + ), + ) + resolver = _resolver( + LocalExtensionArtifact(dependency, dependency_path), + LocalExtensionArtifact(root, root_path), + ) + connection = RecordingConnection(platform) + + loaded = resolver.load(connection, root) + resolver.load(connection, root) + + assert loaded.identity == root.identity + assert connection.loaded_paths == [dependency_path, root_path] + assert resolver.loaded_identities(connection) == (dependency.identity, root.identity) + + +def test_resolver_cache_is_scoped_to_the_connection(tmp_path): + platform = _runtime_platform() + artifact_path = _write_extension_artifact( + tmp_path / "root.duckdb_extension", + platform=platform, + source_id=vane.__git_revision__, + ) + descriptor = _descriptor(artifact_path, name="root") + resolver = _resolver(LocalExtensionArtifact(descriptor, artifact_path)) + first_connection = RecordingConnection(platform) + second_connection = RecordingConnection(platform) + + resolver.load(first_connection, descriptor) + resolver.load(second_connection, descriptor) + + assert first_connection.loaded_paths == [artifact_path] + assert second_connection.loaded_paths == [artifact_path] + + +def test_resolver_rejects_a_different_digest_for_an_already_loaded_extension_name(tmp_path): + platform = _runtime_platform() + first_path = _write_extension_artifact( + tmp_path / "first" / "root.duckdb_extension", + platform=platform, + source_id=vane.__git_revision__, + extension_version="first-version", + ) + second_path = _write_extension_artifact( + tmp_path / "second" / "root.duckdb_extension", + platform=platform, + source_id=vane.__git_revision__, + extension_version="second-version", + ) + first = _descriptor(first_path, name="root") + second = _descriptor(second_path, name="root") + resolver = DynamicExtensionResolver(trusted_identities={"local-tests"}) + connection = RecordingConnection(platform) + + resolver.load(connection, first, artifact=first_path) + with pytest.raises(DynamicExtensionError, match="LOADED_NAME_CONFLICT"): + resolver.load(connection, second, artifact=second_path) + + assert connection.loaded_paths == [first_path] + + +def test_descriptor_creation_rejects_a_name_that_does_not_match_the_artifact(tmp_path): + platform = _runtime_platform() + artifact_path = _write_extension_artifact( + tmp_path / "root.duckdb_extension", + platform=platform, + source_id=vane.__git_revision__, + ) + + with pytest.raises(DynamicExtensionError, match="NAME_MISMATCH"): + _descriptor(artifact_path, name="different") + + +def test_resolver_rejects_untrusted_descriptor_before_loading(tmp_path): + platform = _runtime_platform() + artifact_path = _write_extension_artifact( + tmp_path / "root.duckdb_extension", + platform=platform, + source_id=vane.__git_revision__, + ) + descriptor = _descriptor(artifact_path, name="root", trust_identity="untrusted") + resolver = DynamicExtensionResolver(trusted_identities={"local-tests"}) + connection = RecordingConnection(platform) + + with pytest.raises(DynamicExtensionError, match="TRUST_IDENTITY_UNTRUSTED"): + resolver.load(connection, descriptor, artifact=artifact_path) + + assert connection.loaded_paths == [] + + +def test_resolver_rejects_altered_artifact_before_loading(tmp_path): + platform = _runtime_platform() + artifact_path = _write_extension_artifact( + tmp_path / "root.duckdb_extension", + platform=platform, + source_id=vane.__git_revision__, + ) + descriptor = _descriptor(artifact_path, name="root") + artifact_path.write_bytes(artifact_path.read_bytes() + b"altered") + resolver = DynamicExtensionResolver(trusted_identities={"local-tests"}) + connection = RecordingConnection(platform) + + with pytest.raises(DynamicExtensionError, match="DIGEST_MISMATCH"): + resolver.load(connection, descriptor, artifact=artifact_path) + + assert connection.loaded_paths == [] + + +def test_resolver_rejects_platform_source_id_and_vane_version_mismatches_before_loading(tmp_path): + platform = _runtime_platform() + artifact_path = _write_extension_artifact( + tmp_path / "root.duckdb_extension", + platform=platform, + source_id=vane.__git_revision__, + ) + descriptor = _descriptor(artifact_path, name="root") + resolver = DynamicExtensionResolver(trusted_identities={"local-tests"}) + connection = RecordingConnection(platform) + + with pytest.raises(DynamicExtensionError, match="PLATFORM_MISMATCH"): + resolver.load(connection, replace(descriptor, platform="linux_arm64"), artifact=artifact_path) + with pytest.raises(DynamicExtensionError, match="SOURCE_ID_MISMATCH"): + resolver.load(connection, replace(descriptor, duckdb_source_id="a" * 40), artifact=artifact_path) + with pytest.raises(DynamicExtensionError, match="VANE_VERSION_MISMATCH"): + resolver.load(connection, replace(descriptor, vane_version="different-vane-version"), artifact=artifact_path) + + assert connection.loaded_paths == [] + + +def test_resolver_rejects_missing_ordered_dependency_before_loading_root(tmp_path): + platform = _runtime_platform() + root_path = _write_extension_artifact( + tmp_path / "root.duckdb_extension", + platform=platform, + source_id=vane.__git_revision__, + ) + root = replace( + _descriptor(root_path, name="root"), + dependencies=( + DynamicExtensionDependency( + name="missing", + extension_version="missing-version", + sha256="1" * 64, + ), + ), + ) + resolver = _resolver(LocalExtensionArtifact(root, root_path)) + connection = RecordingConnection(platform) + + with pytest.raises(DynamicExtensionError, match="DEPENDENCY_NOT_FOUND"): + resolver.load(connection, root) + + assert connection.loaded_paths == [] + + +@pytest.fixture(scope="module") +def staged_tpch_artifact(): + configured_path = os.environ.get("VANE_TEST_LOADABLE_EXTENSION_PATH") + if configured_path is None: + pytest.skip("set VANE_TEST_LOADABLE_EXTENSION_PATH to test a staged artifact") + artifact_path = Path(configured_path).resolve() + assert artifact_path.name == "tpch.duckdb_extension" + assert artifact_path.is_file() + return artifact_path + + +def test_resolver_loads_staged_tpch_artifact(staged_tpch_artifact): + descriptor = create_dynamic_extension_descriptor( + staged_tpch_artifact, + name="tpch", + trust_identity="vane-test-artifacts", + ) + resolver = _resolver( + LocalExtensionArtifact(descriptor, staged_tpch_artifact), + trust_identity="vane-test-artifacts", + ) + connection = vane.connect(config={"allow_unsigned_extensions": "true"}) + try: + loaded = resolver.load(connection, descriptor) + + assert loaded.identity == descriptor.identity + assert resolver.loaded_identities(connection) == (descriptor.identity,) + assert connection.execute("SELECT count(*) FROM tpch_queries()").fetchone() == (22,) + finally: + connection.close() diff --git a/vane/__init__.py b/vane/__init__.py index f26ea6f7f97..1a3305f807e 100644 --- a/vane/__init__.py +++ b/vane/__init__.py @@ -193,6 +193,16 @@ version, ) from vane.config import VaneConfig, configure, current_config +from vane.extensions import ( + DynamicExtensionDependency, + DynamicExtensionDescriptor, + DynamicExtensionError, + DynamicExtensionResolver, + LocalExtensionArtifact, + LocalExtensionProvider, + ResolvedDynamicExtension, + create_dynamic_extension_descriptor, +) from vane.value.constant import ( BinaryValue, BitValue, @@ -271,6 +281,7 @@ def set_runner_ray( "datasource", "execution", "experimental", + "extensions", "expressions", "filesystem", "query_graph", @@ -321,6 +332,10 @@ def __dir__() -> list[str]: "DependencyException", "DBAPITypeObject", "DoubleValue", + "DynamicExtensionDependency", + "DynamicExtensionDescriptor", + "DynamicExtensionError", + "DynamicExtensionResolver", "DuckDBPyConnection", "DuckDBPyRelation", "EnvRegistry", @@ -345,6 +360,8 @@ def __dir__() -> list[str]: "LambdaExpression", "ListValue", "LongValue", + "LocalExtensionArtifact", + "LocalExtensionProvider", "MapValue", "NUMBER", "NotImplementedException", @@ -359,6 +376,7 @@ def __dir__() -> list[str]: "PythonExceptionHandling", "RenderMode", "Relation", + "ResolvedDynamicExtension", "ROWID", "SQLExpression", "SequenceException", @@ -409,6 +427,7 @@ def __dir__() -> list[str]: "commit", "connect", "configure", + "create_dynamic_extension_descriptor", "current_config", "cursor", "decimal_type", diff --git a/vane/extensions.py b/vane/extensions.py new file mode 100644 index 00000000000..2af8cbb46c2 --- /dev/null +++ b/vane/extensions.py @@ -0,0 +1,674 @@ +# SPDX-FileCopyrightText: 2026 Vane contributors +# SPDX-License-Identifier: Apache-2.0 +"""Trusted local resolution for loadable DuckDB extensions. + +This module deliberately has no repository, download, or implicit-directory +lookup. A caller supplies an explicit artifact or an installed local provider, +and a descriptor pins the exact bytes that may be loaded into a connection. +""" + +from __future__ import annotations + +import hashlib +import json +import re +from collections.abc import Iterable, Mapping +from dataclasses import dataclass +from pathlib import Path +from typing import Any, NoReturn, Protocol + +_DESCRIPTOR_FORMAT_VERSION = 1 +_EXTENSION_FOOTER_SIZE = 512 +_EXTENSION_FOOTER_FIELD_SIZE = 32 +_EXTENSION_FOOTER_FIELD_COUNT = 8 +_VALID_ABI_TYPES = frozenset({"CPP", "C_STRUCT", "C_STRUCT_UNSTABLE"}) +_EXTENSION_NAME_RE = re.compile(r"^[a-z][a-z0-9_]*$") +_HEX_RE = re.compile(r"^[0-9a-f]{7,64}$") +_SHA256_RE = re.compile(r"^[0-9a-f]{64}$") +_PLATFORM_RE = re.compile(r"^[a-z0-9_]+$") +_TRUST_IDENTITY_RE = re.compile(r"^[A-Za-z0-9._:/-]+$") + + +class DynamicExtensionError(RuntimeError): + """A deterministic failure while resolving or loading an extension.""" + + def __init__(self, code: str, message: str): + self.code = code + super().__init__(f"VANE_DYNAMIC_EXTENSION_{code}: {message}") + + +def _fail(code: str, message: str) -> NoReturn: + raise DynamicExtensionError(code, message) + + +def _require_string(value: object, field_name: str) -> str: + if not isinstance(value, str) or not value: + _fail("DESCRIPTOR_INVALID", f"{field_name} must be a non-empty string") + if any(character.isspace() or ord(character) < 32 for character in value): + _fail("DESCRIPTOR_INVALID", f"{field_name} must not contain whitespace or control characters") + return value + + +def _validate_extension_name(name: object, field_name: str = "name") -> str: + value = _require_string(name, field_name) + if not _EXTENSION_NAME_RE.fullmatch(value): + _fail("DESCRIPTOR_INVALID", f"{field_name} must use lowercase ASCII extension-name syntax") + return value + + +def _validate_sha256(value: object, field_name: str = "sha256") -> str: + digest = _require_string(value, field_name) + if not _SHA256_RE.fullmatch(digest): + _fail("DESCRIPTOR_INVALID", f"{field_name} must be a lowercase SHA-256 digest") + return digest + + +def _validate_source_id(value: object, field_name: str = "duckdb_source_id") -> str: + source_id = _require_string(value, field_name) + if not _HEX_RE.fullmatch(source_id): + _fail("DESCRIPTOR_INVALID", f"{field_name} must be a lowercase Git-compatible object id") + return source_id + + +def _validate_platform(value: object) -> str: + platform = _require_string(value, "platform") + if not _PLATFORM_RE.fullmatch(platform): + _fail("DESCRIPTOR_INVALID", "platform must contain lowercase ASCII letters, digits, and underscores") + return platform + + +def _validate_trust_identity(value: object) -> str: + trust_identity = _require_string(value, "trust_identity") + if not _TRUST_IDENTITY_RE.fullmatch(trust_identity): + _fail("DESCRIPTOR_INVALID", "trust_identity contains unsupported characters") + return trust_identity + + +def _validate_abi_type(value: object) -> str: + abi_type = _require_string(value, "abi_type") + if abi_type not in _VALID_ABI_TYPES: + _fail("DESCRIPTOR_INVALID", f"abi_type must be one of {sorted(_VALID_ABI_TYPES)}") + return abi_type + + +def _validate_extension_version(value: object) -> str: + return _require_string(value, "extension_version") + + +def _runtime_identity() -> tuple[str, str]: + # Import lazily so this module can be imported from vane.__init__ while that + # package is still being initialized. + import vane + + source_id = _validate_source_id(getattr(vane, "__git_revision__", ""), "runtime DuckDB SourceID") + vane_version = _require_string(getattr(vane, "__version__", ""), "runtime Vane version") + return source_id, vane_version + + +@dataclass(frozen=True) +class DynamicExtensionDependency: + """An exact dependency identity, preserved in descriptor order.""" + + name: str + extension_version: str + sha256: str + + def __post_init__(self) -> None: + _validate_extension_name(self.name) + _validate_extension_version(self.extension_version) + _validate_sha256(self.sha256) + + @property + def identity(self) -> str: + """Return the immutable dependency identity.""" + return f"{self.name}@{self.extension_version}#sha256:{self.sha256}" + + def to_dict(self) -> dict[str, str]: + """Serialize this dependency for a descriptor document.""" + return { + "name": self.name, + "extension_version": self.extension_version, + "sha256": self.sha256, + } + + @classmethod + def from_dict(cls, value: Mapping[str, object]) -> DynamicExtensionDependency: + """Deserialize and validate one dependency document.""" + expected_keys = {"name", "extension_version", "sha256"} + if set(value) != expected_keys: + _fail("DESCRIPTOR_INVALID", "dependency must contain only name, extension_version, and sha256") + return cls( + name=_validate_extension_name(value["name"]), + extension_version=_validate_extension_version(value["extension_version"]), + sha256=_validate_sha256(value["sha256"]), + ) + + +@dataclass(frozen=True) +class DynamicExtensionDescriptor: + """Versioned, immutable identity of one loadable extension artifact.""" + + name: str + extension_version: str + abi_type: str + duckdb_source_id: str + vane_version: str + platform: str + sha256: str + trust_identity: str + dependencies: tuple[DynamicExtensionDependency, ...] = () + duckdb_capi_version: str | None = None + format_version: int = _DESCRIPTOR_FORMAT_VERSION + + def __post_init__(self) -> None: + if type(self.format_version) is not int or self.format_version != _DESCRIPTOR_FORMAT_VERSION: + _fail("DESCRIPTOR_INVALID", f"format_version must be {_DESCRIPTOR_FORMAT_VERSION}") + _validate_extension_name(self.name) + _validate_extension_version(self.extension_version) + abi_type = _validate_abi_type(self.abi_type) + _validate_source_id(self.duckdb_source_id) + _require_string(self.vane_version, "vane_version") + _validate_platform(self.platform) + _validate_sha256(self.sha256) + _validate_trust_identity(self.trust_identity) + + try: + dependencies = tuple(self.dependencies) + except TypeError as exception: + raise DynamicExtensionError( + "DESCRIPTOR_INVALID", "dependencies must be an iterable of DynamicExtensionDependency values" + ) from exception + if any(not isinstance(dependency, DynamicExtensionDependency) for dependency in dependencies): + _fail("DESCRIPTOR_INVALID", "dependencies must contain DynamicExtensionDependency values") + dependency_identities = [dependency.identity for dependency in dependencies] + if len(set(dependency_identities)) != len(dependency_identities): + _fail("DESCRIPTOR_INVALID", "dependencies must be unique and preserve declaration order") + object.__setattr__(self, "dependencies", dependencies) + + if abi_type == "C_STRUCT": + _require_string(self.duckdb_capi_version, "duckdb_capi_version") + elif self.duckdb_capi_version is not None: + _fail("DESCRIPTOR_INVALID", "duckdb_capi_version is valid only for C_STRUCT extensions") + + @property + def identity(self) -> str: + """Return the immutable descriptor identity.""" + return f"{self.name}@{self.extension_version}#sha256:{self.sha256}" + + def to_dict(self) -> dict[str, object]: + """Return the canonical descriptor mapping.""" + result: dict[str, object] = { + "format_version": self.format_version, + "name": self.name, + "extension_version": self.extension_version, + "abi_type": self.abi_type, + "duckdb_source_id": self.duckdb_source_id, + "vane_version": self.vane_version, + "platform": self.platform, + "sha256": self.sha256, + "trust_identity": self.trust_identity, + "dependencies": [dependency.to_dict() for dependency in self.dependencies], + } + if self.duckdb_capi_version is not None: + result["duckdb_capi_version"] = self.duckdb_capi_version + return result + + def to_json(self) -> str: + """Serialize the descriptor deterministically.""" + return json.dumps(self.to_dict(), sort_keys=True, separators=(",", ":")) + + @classmethod + def from_dict(cls, value: Mapping[str, object]) -> DynamicExtensionDescriptor: + """Deserialize and validate a version-one descriptor mapping.""" + required_keys = { + "format_version", + "name", + "extension_version", + "abi_type", + "duckdb_source_id", + "vane_version", + "platform", + "sha256", + "trust_identity", + "dependencies", + } + optional_keys = {"duckdb_capi_version"} + unknown_keys = set(value) - required_keys - optional_keys + missing_keys = required_keys - set(value) + if missing_keys or unknown_keys: + details = [] + if missing_keys: + details.append(f"missing {sorted(missing_keys)}") + if unknown_keys: + details.append(f"unknown {sorted(unknown_keys)}") + _fail("DESCRIPTOR_INVALID", f"descriptor keys are invalid: {', '.join(details)}") + + dependencies_value = value["dependencies"] + if not isinstance(dependencies_value, list): + _fail("DESCRIPTOR_INVALID", "dependencies must be a list") + dependencies: list[DynamicExtensionDependency] = [] + for dependency in dependencies_value: + if not isinstance(dependency, Mapping): + _fail("DESCRIPTOR_INVALID", "each dependency must be an object") + dependencies.append(DynamicExtensionDependency.from_dict(dependency)) + + format_version = value["format_version"] + if type(format_version) is not int: + _fail("DESCRIPTOR_INVALID", "format_version must be an integer") + capi_version = value.get("duckdb_capi_version") + if capi_version is not None and not isinstance(capi_version, str): + _fail("DESCRIPTOR_INVALID", "duckdb_capi_version must be a string") + return cls( + format_version=format_version, + name=_validate_extension_name(value["name"]), + extension_version=_validate_extension_version(value["extension_version"]), + abi_type=_validate_abi_type(value["abi_type"]), + duckdb_source_id=_validate_source_id(value["duckdb_source_id"]), + vane_version=_require_string(value["vane_version"], "vane_version"), + platform=_validate_platform(value["platform"]), + sha256=_validate_sha256(value["sha256"]), + trust_identity=_validate_trust_identity(value["trust_identity"]), + dependencies=tuple(dependencies), + duckdb_capi_version=capi_version, + ) + + @classmethod + def from_json(cls, value: str | bytes | bytearray) -> DynamicExtensionDescriptor: + """Deserialize one JSON descriptor document.""" + try: + parsed = json.loads(value) + except (TypeError, ValueError) as exception: + raise DynamicExtensionError("DESCRIPTOR_INVALID", "descriptor is not valid JSON") from exception + if not isinstance(parsed, Mapping): + _fail("DESCRIPTOR_INVALID", "descriptor JSON must contain an object") + return cls.from_dict(parsed) + + +@dataclass(frozen=True) +class LocalExtensionArtifact: + """A descriptor paired with an explicit local binary path.""" + + descriptor: DynamicExtensionDescriptor + path: Path + + def __post_init__(self) -> None: + if not isinstance(self.descriptor, DynamicExtensionDescriptor): + _fail("DESCRIPTOR_INVALID", "artifact descriptor must be a DynamicExtensionDescriptor") + object.__setattr__(self, "path", Path(self.path).expanduser().resolve()) + + +@dataclass(frozen=True) +class ResolvedDynamicExtension: + """A verified local artifact ready to be loaded into a connection.""" + + descriptor: DynamicExtensionDescriptor + path: Path + + @property + def identity(self) -> str: + """Return the verified descriptor identity.""" + return self.descriptor.identity + + +@dataclass(frozen=True) +class _ExtensionFooter: + abi_type: str + platform: str + engine_identity: str + extension_version: str + + +class LocalExtensionProvider: + """An installed, explicit provider of local extension artifacts. + + Packaging code can construct this provider from files installed by a + platform wheel. It intentionally has no directory scan or network fallback. + """ + + def __init__(self, trust_identity: str, artifacts: Iterable[LocalExtensionArtifact]): + self._trust_identity = _validate_trust_identity(trust_identity) + artifact_by_identity: dict[str, LocalExtensionArtifact] = {} + for artifact in artifacts: + if not isinstance(artifact, LocalExtensionArtifact): + _fail("DESCRIPTOR_INVALID", "provider artifacts must be LocalExtensionArtifact values") + if artifact.descriptor.trust_identity != self._trust_identity: + _fail( + "DESCRIPTOR_INVALID", + "provider trust_identity must equal each artifact descriptor trust_identity", + ) + if artifact.descriptor.identity in artifact_by_identity: + _fail("DESCRIPTOR_INVALID", f"provider declares {artifact.descriptor.identity} more than once") + artifact_by_identity[artifact.descriptor.identity] = artifact + self._artifact_by_identity = artifact_by_identity + + @property + def trust_identity(self) -> str: + """Return the local provider trust identity.""" + return self._trust_identity + + def find(self, identity: str) -> LocalExtensionArtifact | None: + """Return one exact artifact identity, without any fallback lookup.""" + return self._artifact_by_identity.get(identity) + + +class _ExtensionConnection(Protocol): + def execute(self, query: str) -> Any: ... + + def load_extension(self, extension: str) -> None: ... + + +class DynamicExtensionResolver: + """Resolve and load only explicitly trusted local extension artifacts.""" + + def __init__( + self, + *, + trusted_identities: Iterable[str], + providers: Iterable[LocalExtensionProvider] = (), + ): + trusted = frozenset(_validate_trust_identity(identity) for identity in trusted_identities) + if not trusted: + _fail("DESCRIPTOR_INVALID", "trusted_identities must not be empty") + self._trusted_identities = trusted + self._providers = tuple(providers) + # Keep the connection beside its id. A bare id() cache can be reused by + # a later connection after the original object is collected. + self._loaded_by_connection: dict[int, tuple[_ExtensionConnection, dict[str, ResolvedDynamicExtension]]] = {} + + def resolve( + self, + connection: _ExtensionConnection, + descriptor: DynamicExtensionDescriptor, + *, + artifact: str | Path | None = None, + ) -> tuple[ResolvedDynamicExtension, ...]: + """Verify dependencies first and return a deterministic load order.""" + current_source_id, current_vane_version = _runtime_identity() + current_platform = _connection_platform(connection) + resolved: list[ResolvedDynamicExtension] = [] + resolved_identities: set[str] = set() + visiting: set[str] = set() + + def visit(candidate: DynamicExtensionDescriptor, candidate_path: Path | None) -> None: + if candidate.identity in resolved_identities: + return + if candidate.identity in visiting: + _fail("DEPENDENCY_CYCLE", f"dependency cycle contains {candidate.identity}") + visiting.add(candidate.identity) + try: + if candidate_path is None: + candidate_path = self._provider_artifact(candidate.identity).path + self._verify_artifact( + candidate, + candidate_path, + current_source_id=current_source_id, + current_vane_version=current_vane_version, + current_platform=current_platform, + ) + for dependency in candidate.dependencies: + dependency_artifact = self._provider_artifact(dependency.identity) + if dependency_artifact.descriptor.identity != dependency.identity: + _fail("DEPENDENCY_NOT_FOUND", f"provider identity does not match {dependency.identity}") + visit(dependency_artifact.descriptor, dependency_artifact.path) + resolved.append(ResolvedDynamicExtension(candidate, candidate_path)) + resolved_identities.add(candidate.identity) + finally: + visiting.discard(candidate.identity) + + explicit_artifact = Path(artifact).expanduser().resolve() if artifact is not None else None + visit(descriptor, explicit_artifact) + return tuple(resolved) + + def load( + self, + connection: _ExtensionConnection, + descriptor: DynamicExtensionDescriptor, + *, + artifact: str | Path | None = None, + ) -> ResolvedDynamicExtension: + """Resolve and load dependencies before the requested extension.""" + resolved = self.resolve(connection, descriptor, artifact=artifact) + loaded = self._loaded_for_connection(connection) + for candidate in resolved: + existing = loaded.get(candidate.descriptor.sha256) + if existing is not None: + if existing.identity != candidate.identity: + _fail( + "LOADED_IDENTITY_CONFLICT", + f"digest {candidate.descriptor.sha256} is already cached as {existing.identity}", + ) + continue + for loaded_candidate in loaded.values(): + if loaded_candidate.descriptor.name == candidate.descriptor.name: + _fail( + "LOADED_NAME_CONFLICT", + f"{candidate.descriptor.name} is already loaded as {loaded_candidate.identity}", + ) + try: + connection.load_extension(str(candidate.path)) + except Exception as exception: + raise DynamicExtensionError( + "LOAD_FAILED", + f"failed to load {candidate.identity} from {candidate.path.name}: {exception}", + ) from exception + loaded[candidate.descriptor.sha256] = candidate + return resolved[-1] + + def loaded_identities(self, connection: _ExtensionConnection) -> tuple[str, ...]: + """Return cached artifact identities for one connection in load order.""" + cached = self._loaded_by_connection.get(id(connection)) + if cached is None or cached[0] is not connection: + return () + return tuple(candidate.identity for candidate in cached[1].values()) + + def _loaded_for_connection(self, connection: _ExtensionConnection) -> dict[str, ResolvedDynamicExtension]: + connection_id = id(connection) + cached = self._loaded_by_connection.get(connection_id) + if cached is not None and cached[0] is connection: + return cached[1] + loaded: dict[str, ResolvedDynamicExtension] = {} + self._loaded_by_connection[connection_id] = (connection, loaded) + return loaded + + def _provider_artifact(self, identity: str) -> LocalExtensionArtifact: + candidates = [artifact for provider in self._providers if (artifact := provider.find(identity)) is not None] + if not candidates: + _fail("DEPENDENCY_NOT_FOUND", f"no trusted local provider contains {identity}") + if len(candidates) != 1: + _fail("ARTIFACT_AMBIGUOUS", f"multiple local providers contain {identity}") + return candidates[0] + + def _verify_artifact( + self, + descriptor: DynamicExtensionDescriptor, + artifact_path: Path, + *, + current_source_id: str, + current_vane_version: str, + current_platform: str, + ) -> None: + if descriptor.trust_identity not in self._trusted_identities: + _fail("TRUST_IDENTITY_UNTRUSTED", f"{descriptor.trust_identity} is not in trusted_identities") + if descriptor.duckdb_source_id != current_source_id: + _fail( + "SOURCE_ID_MISMATCH", + f"{descriptor.identity} requires SourceID {descriptor.duckdb_source_id}, runtime has {current_source_id}", + ) + if descriptor.vane_version != current_vane_version: + _fail( + "VANE_VERSION_MISMATCH", + f"{descriptor.identity} requires Vane {descriptor.vane_version}, runtime has {current_vane_version}", + ) + if descriptor.platform != current_platform: + _fail( + "PLATFORM_MISMATCH", + f"{descriptor.identity} targets {descriptor.platform}, runtime is {current_platform}", + ) + expected_filename = f"{descriptor.name}.duckdb_extension" + if artifact_path.name != expected_filename: + _fail("NAME_MISMATCH", f"{descriptor.identity} must use artifact filename {expected_filename}") + if not artifact_path.is_file(): + _fail("ARTIFACT_NOT_FOUND", f"artifact does not exist: {artifact_path}") + actual_digest = _sha256_file(artifact_path) + if actual_digest != descriptor.sha256: + _fail( + "DIGEST_MISMATCH", + f"{descriptor.identity} expected SHA-256 {descriptor.sha256}, got {actual_digest}", + ) + footer = _parse_extension_footer(artifact_path) + if footer.abi_type != descriptor.abi_type: + _fail( + "ABI_MISMATCH", + f"{descriptor.identity} declares ABI {descriptor.abi_type}, artifact footer has {footer.abi_type}", + ) + if footer.platform != descriptor.platform: + _fail( + "PLATFORM_MISMATCH", + f"{descriptor.identity} declares platform {descriptor.platform}, artifact footer has {footer.platform}", + ) + if footer.extension_version != descriptor.extension_version: + _fail( + "EXTENSION_VERSION_MISMATCH", + f"{descriptor.identity} footer version is {footer.extension_version}", + ) + if descriptor.abi_type == "C_STRUCT": + if footer.engine_identity != descriptor.duckdb_capi_version: + _fail( + "CAPI_VERSION_MISMATCH", + f"{descriptor.identity} requires C API {descriptor.duckdb_capi_version}, " + f"artifact footer has {footer.engine_identity}", + ) + elif footer.engine_identity != descriptor.duckdb_source_id: + _fail( + "SOURCE_ID_MISMATCH", + f"{descriptor.identity} footer SourceID is {footer.engine_identity}", + ) + + +def create_dynamic_extension_descriptor( + artifact: str | Path, + *, + name: str, + trust_identity: str, + dependencies: Iterable[DynamicExtensionDependency] = (), + vane_version: str | None = None, +) -> DynamicExtensionDescriptor: + """Create a version-one descriptor directly from a built local artifact.""" + artifact_path = Path(artifact).expanduser().resolve() + validated_name = _validate_extension_name(name) + expected_filename = f"{validated_name}.duckdb_extension" + if artifact_path.name != expected_filename: + _fail("NAME_MISMATCH", f"descriptor name {validated_name} requires artifact filename {expected_filename}") + footer = _parse_extension_footer(artifact_path) + source_id, runtime_vane_version = _runtime_identity() + descriptor_vane_version = ( + runtime_vane_version if vane_version is None else _require_string(vane_version, "vane_version") + ) + if footer.abi_type == "C_STRUCT": + return DynamicExtensionDescriptor( + name=validated_name, + extension_version=footer.extension_version, + abi_type=footer.abi_type, + duckdb_source_id=source_id, + vane_version=descriptor_vane_version, + platform=footer.platform, + sha256=_sha256_file(artifact_path), + trust_identity=trust_identity, + dependencies=tuple(dependencies), + duckdb_capi_version=footer.engine_identity, + ) + return DynamicExtensionDescriptor( + name=validated_name, + extension_version=footer.extension_version, + abi_type=footer.abi_type, + duckdb_source_id=_validate_source_id(footer.engine_identity, "artifact footer SourceID"), + vane_version=descriptor_vane_version, + platform=footer.platform, + sha256=_sha256_file(artifact_path), + trust_identity=trust_identity, + dependencies=tuple(dependencies), + ) + + +def _connection_platform(connection: _ExtensionConnection) -> str: + try: + row = connection.execute("SELECT platform FROM pragma_platform()").fetchone() + except Exception as exception: + raise DynamicExtensionError( + "RUNTIME_IDENTITY_UNAVAILABLE", "could not query the DuckDB platform" + ) from exception + if not isinstance(row, tuple) or len(row) != 1 or not isinstance(row[0], str): + _fail("RUNTIME_IDENTITY_UNAVAILABLE", "pragma_platform() did not return one platform string") + return _validate_platform(row[0]) + + +def _sha256_file(path: Path) -> str: + digest = hashlib.sha256() + try: + with path.open("rb") as artifact_file: + for chunk in iter(lambda: artifact_file.read(1024 * 1024), b""): + digest.update(chunk) + except OSError as exception: + raise DynamicExtensionError("ARTIFACT_NOT_FOUND", f"could not read artifact: {path}") from exception + return digest.hexdigest() + + +def _parse_extension_footer(path: Path) -> _ExtensionFooter: + try: + with path.open("rb") as artifact_file: + artifact_file.seek(0, 2) + if artifact_file.tell() < _EXTENSION_FOOTER_SIZE: + _fail("FOOTER_INVALID", f"artifact is smaller than {_EXTENSION_FOOTER_SIZE} bytes: {path.name}") + artifact_file.seek(-_EXTENSION_FOOTER_SIZE, 2) + footer = artifact_file.read(_EXTENSION_FOOTER_SIZE) + except OSError as exception: + raise DynamicExtensionError("ARTIFACT_NOT_FOUND", f"could not read artifact: {path}") from exception + if len(footer) != _EXTENSION_FOOTER_SIZE: + _fail("FOOTER_INVALID", f"could not read a complete extension footer from {path.name}") + + fields = [] + for field_index in range(_EXTENSION_FOOTER_FIELD_COUNT): + start = field_index * _EXTENSION_FOOTER_FIELD_SIZE + raw_field = footer[start : start + _EXTENSION_FOOTER_FIELD_SIZE] + try: + field = raw_field.decode("ascii").rstrip("\0") + except UnicodeDecodeError as exception: + raise DynamicExtensionError( + "FOOTER_INVALID", f"extension footer contains non-ASCII field {field_index}" + ) from exception + if "\0" in field: + _fail("FOOTER_INVALID", f"extension footer contains an embedded NUL in field {field_index}") + fields.append(field) + + if fields[7] != "4": + _fail("FOOTER_INVALID", f"artifact footer magic is invalid: {path.name}") + abi_type = fields[3] or "CPP" + if abi_type not in _VALID_ABI_TYPES: + _fail("FOOTER_INVALID", f"artifact footer has unsupported ABI {abi_type!r}") + platform = fields[6] + if _PLATFORM_RE.fullmatch(platform) is None: + _fail("FOOTER_INVALID", "artifact footer platform is invalid") + engine_identity = fields[5] + if not engine_identity or any(character.isspace() or ord(character) < 32 for character in engine_identity): + _fail("FOOTER_INVALID", "artifact footer engine identity is invalid") + extension_version = fields[4] + if not extension_version or any(character.isspace() or ord(character) < 32 for character in extension_version): + _fail("FOOTER_INVALID", "artifact footer extension version is invalid") + return _ExtensionFooter( + abi_type=abi_type, + platform=platform, + engine_identity=engine_identity, + extension_version=extension_version, + ) + + +__all__ = [ + "DynamicExtensionDependency", + "DynamicExtensionDescriptor", + "DynamicExtensionError", + "DynamicExtensionResolver", + "LocalExtensionArtifact", + "LocalExtensionProvider", + "ResolvedDynamicExtension", + "create_dynamic_extension_descriptor", +] From ef7db72c6d2ad8be4ba34b37f6138eb6ef9f8e3a Mon Sep 17 00:00:00 2001 From: kaka11chen Date: Wed, 26 Aug 2026 00:41:37 +0800 Subject: [PATCH 3/4] [fix](extensions) Harden trusted extension loading --- src/vane_py/vane_python.cpp | 2 + tests/fast/test_dynamic_extension_resolver.py | 193 ++++++++++++- vane/_native/__init__.pyi | 1 + vane/extensions.py | 268 +++++++++++++----- 4 files changed, 386 insertions(+), 78 deletions(-) diff --git a/src/vane_py/vane_python.cpp b/src/vane_py/vane_python.cpp index 95ee64be29f..f0eb4e924a7 100644 --- a/src/vane_py/vane_python.cpp +++ b/src/vane_py/vane_python.cpp @@ -9,6 +9,7 @@ #include "duckdb/common/atomic.hpp" #include "duckdb/common/vector.hpp" +#include "duckdb/main/capi/extension_api.hpp" #include "duckdb/parser/parser.hpp" #include "vane_python/python_objects.hpp" @@ -1141,6 +1142,7 @@ PYBIND11_MODULE(_native, m) { // NOLINT m.attr("__version__") = std::string(DuckDB::LibraryVersion()).substr(1); m.attr("__standard_vector_size__") = DuckDB::StandardVectorSize(); m.attr("__git_revision__") = DuckDB::SourceID(); + m.attr("__duckdb_extension_api_version__") = DUCKDB_EXTENSION_API_VERSION_STRING; m.attr("__interactive__") = DuckDBPyConnection::DetectAndGetEnvironment(); m.attr("__jupyter__") = DuckDBPyConnection::IsJupyter(); m.attr("__formatted_python_version__") = DuckDBPyConnection::FormattedPythonVersion(); diff --git a/tests/fast/test_dynamic_extension_resolver.py b/tests/fast/test_dynamic_extension_resolver.py index f0043216bce..02c6efa7b89 100644 --- a/tests/fast/test_dynamic_extension_resolver.py +++ b/tests/fast/test_dynamic_extension_resolver.py @@ -1,7 +1,10 @@ # SPDX-FileCopyrightText: 2026 Vane contributors # SPDX-License-Identifier: Apache-2.0 +import gc +import hashlib import os +import weakref from dataclasses import replace from pathlib import Path @@ -20,24 +23,38 @@ class _Result: - def __init__(self, row): - self._row = row + def __init__(self, rows): + self._rows = rows def fetchone(self): - return self._row + return self._rows[0] if self._rows else None + + def fetchall(self): + return list(self._rows) class RecordingConnection: - def __init__(self, platform): + def __init__(self, platform, *, loaded_names=(), before_load=None): self.platform = platform self.loaded_paths = [] + self.loaded_payloads = [] + self.loaded_names = set(loaded_names) + self.before_load = before_load def execute(self, query): - assert query == "SELECT platform FROM pragma_platform()" - return _Result((self.platform,)) + if query == "SELECT platform FROM pragma_platform()": + return _Result([(self.platform,)]) + if query == "SELECT extension_name FROM duckdb_extensions() WHERE loaded ORDER BY extension_name": + return _Result([(name,) for name in sorted(self.loaded_names)]) + raise AssertionError(f"unexpected query: {query}") def load_extension(self, extension): - self.loaded_paths.append(Path(extension)) + path = Path(extension) + if self.before_load is not None: + self.before_load(path) + self.loaded_paths.append(path) + self.loaded_payloads.append(path.read_bytes()) + self.loaded_names.add(path.name.removesuffix(".duckdb_extension")) def _runtime_platform(): @@ -134,7 +151,8 @@ def test_resolver_loads_dependencies_before_root_and_caches_digest(tmp_path): resolver.load(connection, root) assert loaded.identity == root.identity - assert connection.loaded_paths == [dependency_path, root_path] + assert [path.name for path in connection.loaded_paths] == [dependency_path.name, root_path.name] + assert all(path != source for path, source in zip(connection.loaded_paths, [dependency_path, root_path])) assert resolver.loaded_identities(connection) == (dependency.identity, root.identity) @@ -153,8 +171,8 @@ def test_resolver_cache_is_scoped_to_the_connection(tmp_path): resolver.load(first_connection, descriptor) resolver.load(second_connection, descriptor) - assert first_connection.loaded_paths == [artifact_path] - assert second_connection.loaded_paths == [artifact_path] + assert [path.name for path in first_connection.loaded_paths] == [artifact_path.name] + assert [path.name for path in second_connection.loaded_paths] == [artifact_path.name] def test_resolver_rejects_a_different_digest_for_an_already_loaded_extension_name(tmp_path): @@ -180,7 +198,160 @@ def test_resolver_rejects_a_different_digest_for_an_already_loaded_extension_nam with pytest.raises(DynamicExtensionError, match="LOADED_NAME_CONFLICT"): resolver.load(connection, second, artifact=second_path) - assert connection.loaded_paths == [first_path] + assert [path.name for path in connection.loaded_paths] == [first_path.name] + + +def test_resolver_rejects_an_extension_loaded_outside_its_cache(tmp_path): + platform = _runtime_platform() + artifact_path = _write_extension_artifact( + tmp_path / "root.duckdb_extension", + platform=platform, + source_id=vane.__git_revision__, + ) + descriptor = _descriptor(artifact_path, name="root") + resolver = DynamicExtensionResolver(trusted_identities={"local-tests"}) + connection = RecordingConnection(platform, loaded_names={"root"}) + + with pytest.raises(DynamicExtensionError, match="LOADED_OUTSIDE_RESOLVER"): + resolver.load(connection, descriptor, artifact=artifact_path) + + assert connection.loaded_paths == [] + + +def test_resolver_loads_the_verified_snapshot_if_the_provider_path_is_replaced(tmp_path): + platform = _runtime_platform() + artifact_path = _write_extension_artifact( + tmp_path / "root.duckdb_extension", + platform=platform, + source_id=vane.__git_revision__, + ) + verified_payload = artifact_path.read_bytes() + descriptor = _descriptor(artifact_path, name="root") + resolver = DynamicExtensionResolver(trusted_identities={"local-tests"}) + + def replace_provider_path(_snapshot_path): + artifact_path.write_bytes(b"replacement after verification") + + connection = RecordingConnection(platform, before_load=replace_provider_path) + + loaded = resolver.load(connection, descriptor, artifact=artifact_path) + + assert loaded.path != artifact_path + assert connection.loaded_payloads == [verified_payload] + assert hashlib.sha256(connection.loaded_payloads[0]).hexdigest() == descriptor.sha256 + assert artifact_path.read_bytes() != verified_payload + + +def test_resolver_cache_does_not_keep_connections_alive(tmp_path): + platform = _runtime_platform() + artifact_path = _write_extension_artifact( + tmp_path / "root.duckdb_extension", + platform=platform, + source_id=vane.__git_revision__, + ) + descriptor = _descriptor(artifact_path, name="root") + resolver = DynamicExtensionResolver(trusted_identities={"local-tests"}) + connection = RecordingConnection(platform) + connection_reference = weakref.ref(connection) + + resolver.load(connection, descriptor, artifact=artifact_path) + assert len(resolver._loaded_by_connection) == 1 + + del connection + # The pytest timing plugin can retain the most recently returned Python + # frames while output capture is active. Advance it past load()/resolve() + # before checking that the resolver cache itself does not retain the key. + resolver.loaded_identities(RecordingConnection(platform)) + gc.collect() + + assert connection_reference() is None + assert len(resolver._loaded_by_connection) == 0 + + +def test_resolver_rejects_a_requested_descriptor_that_differs_from_the_provider(tmp_path): + platform = _runtime_platform() + artifact_path = _write_extension_artifact( + tmp_path / "root.duckdb_extension", + platform=platform, + source_id=vane.__git_revision__, + ) + provider_descriptor = _descriptor(artifact_path, name="root") + requested_descriptor = replace( + provider_descriptor, + dependencies=( + DynamicExtensionDependency( + name="injected", + extension_version="injected-version", + sha256="1" * 64, + ), + ), + ) + resolver = _resolver(LocalExtensionArtifact(provider_descriptor, artifact_path)) + connection = RecordingConnection(platform) + + with pytest.raises(DynamicExtensionError, match="PROVIDER_DESCRIPTOR_MISMATCH"): + resolver.load(connection, requested_descriptor) + + assert connection.loaded_paths == [] + + +def test_resolver_rejects_conflicting_dependency_names_before_loading_anything(tmp_path): + platform = _runtime_platform() + first_path = _write_extension_artifact( + tmp_path / "first" / "shared.duckdb_extension", + platform=platform, + source_id=vane.__git_revision__, + extension_version="first-version", + ) + second_path = _write_extension_artifact( + tmp_path / "second" / "shared.duckdb_extension", + platform=platform, + source_id=vane.__git_revision__, + extension_version="second-version", + ) + root_path = _write_extension_artifact( + tmp_path / "root.duckdb_extension", + platform=platform, + source_id=vane.__git_revision__, + ) + first = _descriptor(first_path, name="shared") + second = _descriptor(second_path, name="shared") + root = replace( + _descriptor(root_path, name="root"), + dependencies=( + DynamicExtensionDependency(first.name, first.extension_version, first.sha256), + DynamicExtensionDependency(second.name, second.extension_version, second.sha256), + ), + ) + resolver = _resolver( + LocalExtensionArtifact(first, first_path), + LocalExtensionArtifact(second, second_path), + LocalExtensionArtifact(root, root_path), + ) + connection = RecordingConnection(platform) + + with pytest.raises(DynamicExtensionError, match="RESOLVED_NAME_CONFLICT"): + resolver.load(connection, root) + + assert connection.loaded_paths == [] + + +def test_resolver_rejects_an_unsupported_c_struct_api_before_loading(tmp_path): + platform = _runtime_platform() + artifact_path = _write_extension_artifact( + tmp_path / "root.duckdb_extension", + platform=platform, + source_id="v999.0.0", + abi_type="C_STRUCT", + ) + descriptor = _descriptor(artifact_path, name="root") + resolver = DynamicExtensionResolver(trusted_identities={"local-tests"}) + connection = RecordingConnection(platform) + + with pytest.raises(DynamicExtensionError, match="CAPI_VERSION_MISMATCH"): + resolver.load(connection, descriptor, artifact=artifact_path) + + assert connection.loaded_paths == [] def test_descriptor_creation_rejects_a_name_that_does_not_match_the_artifact(tmp_path): diff --git a/vane/_native/__init__.pyi b/vane/_native/__init__.pyi index 7fb8ea33538..b40be87af6e 100644 --- a/vane/_native/__init__.pyi +++ b/vane/_native/__init__.pyi @@ -1560,6 +1560,7 @@ def write_csv( ) -> None: ... __formatted_python_version__: str +__duckdb_extension_api_version__: str __git_revision__: str __interactive__: bool __jupyter__: bool diff --git a/vane/extensions.py b/vane/extensions.py index 2af8cbb46c2..dc2697ff37f 100644 --- a/vane/extensions.py +++ b/vane/extensions.py @@ -11,7 +11,11 @@ import hashlib import json +import os import re +import tempfile +import threading +import weakref from collections.abc import Iterable, Mapping from dataclasses import dataclass from pathlib import Path @@ -27,6 +31,7 @@ _SHA256_RE = re.compile(r"^[0-9a-f]{64}$") _PLATFORM_RE = re.compile(r"^[a-z0-9_]+$") _TRUST_IDENTITY_RE = re.compile(r"^[A-Za-z0-9._:/-]+$") +_CAPI_VERSION_RE = re.compile(r"^v(0|[1-9][0-9]*)\.(0|[1-9][0-9]*)\.(0|[1-9][0-9]*)$") class DynamicExtensionError(RuntimeError): @@ -95,6 +100,14 @@ def _validate_extension_version(value: object) -> str: return _require_string(value, "extension_version") +def _parse_capi_version(value: object, field_name: str) -> tuple[int, int, int]: + capi_version = _require_string(value, field_name) + match = _CAPI_VERSION_RE.fullmatch(capi_version) + if match is None: + _fail("DESCRIPTOR_INVALID", f"{field_name} must use v.. syntax") + return int(match.group(1)), int(match.group(2)), int(match.group(3)) + + def _runtime_identity() -> tuple[str, str]: # Import lazily so this module can be imported from vane.__init__ while that # package is still being initialized. @@ -105,6 +118,21 @@ def _runtime_identity() -> tuple[str, str]: return source_id, vane_version +def _runtime_capi_version() -> str: + # Keep this value tied to the native runtime rather than duplicating the + # vendored DuckDB header's version in Python. + from vane import _native + + capi_version = getattr(_native, "__duckdb_extension_api_version__", "") + try: + _parse_capi_version(capi_version, "runtime DuckDB extension C API version") + except DynamicExtensionError as exception: + raise DynamicExtensionError( + "RUNTIME_IDENTITY_UNAVAILABLE", "native runtime did not expose a valid DuckDB extension C API version" + ) from exception + return capi_version + + @dataclass(frozen=True) class DynamicExtensionDependency: """An exact dependency identity, preserved in descriptor order.""" @@ -186,7 +214,7 @@ def __post_init__(self) -> None: object.__setattr__(self, "dependencies", dependencies) if abi_type == "C_STRUCT": - _require_string(self.duckdb_capi_version, "duckdb_capi_version") + _parse_capi_version(self.duckdb_capi_version, "duckdb_capi_version") elif self.duckdb_capi_version is not None: _fail("DESCRIPTOR_INVALID", "duckdb_capi_version is valid only for C_STRUCT extensions") @@ -371,9 +399,16 @@ def __init__( _fail("DESCRIPTOR_INVALID", "trusted_identities must not be empty") self._trusted_identities = trusted self._providers = tuple(providers) - # Keep the connection beside its id. A bare id() cache can be reused by - # a later connection after the original object is collected. - self._loaded_by_connection: dict[int, tuple[_ExtensionConnection, dict[str, ResolvedDynamicExtension]]] = {} + self._loaded_by_connection: weakref.WeakKeyDictionary[ + _ExtensionConnection, dict[str, ResolvedDynamicExtension] + ] = weakref.WeakKeyDictionary() + self._lock = threading.RLock() + # DuckDB re-opens an extension by path. Load a private copy of the exact + # bytes we verified so replacing the provider path cannot change what is + # handed to the native loader. + self._snapshot_directory = tempfile.TemporaryDirectory( + prefix="vane-dynamic-extensions-", ignore_cleanup_errors=True + ) def resolve( self, @@ -397,8 +432,14 @@ def visit(candidate: DynamicExtensionDescriptor, candidate_path: Path | None) -> visiting.add(candidate.identity) try: if candidate_path is None: - candidate_path = self._provider_artifact(candidate.identity).path - self._verify_artifact( + provider_artifact = self._provider_artifact(candidate.identity) + if provider_artifact.descriptor != candidate: + _fail( + "PROVIDER_DESCRIPTOR_MISMATCH", + f"provider descriptor for {candidate.identity} does not match the requested descriptor", + ) + candidate_path = provider_artifact.path + verified_path = self._verify_artifact( candidate, candidate_path, current_source_id=current_source_id, @@ -410,13 +451,14 @@ def visit(candidate: DynamicExtensionDescriptor, candidate_path: Path | None) -> if dependency_artifact.descriptor.identity != dependency.identity: _fail("DEPENDENCY_NOT_FOUND", f"provider identity does not match {dependency.identity}") visit(dependency_artifact.descriptor, dependency_artifact.path) - resolved.append(ResolvedDynamicExtension(candidate, candidate_path)) + resolved.append(ResolvedDynamicExtension(candidate, verified_path)) resolved_identities.add(candidate.identity) finally: visiting.discard(candidate.identity) explicit_artifact = Path(artifact).expanduser().resolve() if artifact is not None else None visit(descriptor, explicit_artifact) + self._validate_resolved_names(resolved) return tuple(resolved) def load( @@ -427,48 +469,88 @@ def load( artifact: str | Path | None = None, ) -> ResolvedDynamicExtension: """Resolve and load dependencies before the requested extension.""" - resolved = self.resolve(connection, descriptor, artifact=artifact) - loaded = self._loaded_for_connection(connection) - for candidate in resolved: - existing = loaded.get(candidate.descriptor.sha256) - if existing is not None: - if existing.identity != candidate.identity: - _fail( - "LOADED_IDENTITY_CONFLICT", - f"digest {candidate.descriptor.sha256} is already cached as {existing.identity}", - ) - continue - for loaded_candidate in loaded.values(): - if loaded_candidate.descriptor.name == candidate.descriptor.name: - _fail( - "LOADED_NAME_CONFLICT", - f"{candidate.descriptor.name} is already loaded as {loaded_candidate.identity}", - ) - try: - connection.load_extension(str(candidate.path)) - except Exception as exception: - raise DynamicExtensionError( - "LOAD_FAILED", - f"failed to load {candidate.identity} from {candidate.path.name}: {exception}", - ) from exception - loaded[candidate.descriptor.sha256] = candidate - return resolved[-1] + with self._lock: + resolved = self.resolve(connection, descriptor, artifact=artifact) + loaded = self._loaded_for_connection(connection) + database_loaded_names = _database_loaded_extension_names(connection) + self._validate_load_plan(resolved, loaded, database_loaded_names) + for candidate in resolved: + if candidate.descriptor.sha256 in loaded: + continue + try: + connection.load_extension(str(candidate.path)) + except Exception as exception: + raise DynamicExtensionError( + "LOAD_FAILED", + f"failed to load {candidate.identity} from {candidate.path.name}: {exception}", + ) from exception + loaded[candidate.descriptor.sha256] = candidate + return resolved[-1] def loaded_identities(self, connection: _ExtensionConnection) -> tuple[str, ...]: """Return cached artifact identities for one connection in load order.""" - cached = self._loaded_by_connection.get(id(connection)) - if cached is None or cached[0] is not connection: - return () - return tuple(candidate.identity for candidate in cached[1].values()) + with self._lock: + try: + cached = self._loaded_by_connection.get(connection) + except TypeError: + return () + if cached is None: + return () + return tuple(candidate.identity for candidate in cached.values()) def _loaded_for_connection(self, connection: _ExtensionConnection) -> dict[str, ResolvedDynamicExtension]: - connection_id = id(connection) - cached = self._loaded_by_connection.get(connection_id) - if cached is not None and cached[0] is connection: - return cached[1] - loaded: dict[str, ResolvedDynamicExtension] = {} - self._loaded_by_connection[connection_id] = (connection, loaded) - return loaded + try: + cached = self._loaded_by_connection.get(connection) + if cached is not None: + return cached + loaded: dict[str, ResolvedDynamicExtension] = {} + self._loaded_by_connection[connection] = loaded + return loaded + except TypeError as exception: + raise DynamicExtensionError( + "CONNECTION_UNSUPPORTED", "connection objects must support weak references and identity hashing" + ) from exception + + @staticmethod + def _validate_resolved_names(resolved: Iterable[ResolvedDynamicExtension]) -> None: + by_name: dict[str, ResolvedDynamicExtension] = {} + for candidate in resolved: + existing = by_name.get(candidate.descriptor.name) + if existing is not None and existing.identity != candidate.identity: + _fail( + "RESOLVED_NAME_CONFLICT", + f"dependency graph resolves {candidate.descriptor.name} as both " + f"{existing.identity} and {candidate.identity}", + ) + by_name[candidate.descriptor.name] = candidate + + @staticmethod + def _validate_load_plan( + resolved: Iterable[ResolvedDynamicExtension], + loaded: Mapping[str, ResolvedDynamicExtension], + database_loaded_names: frozenset[str], + ) -> None: + loaded_by_name = {candidate.descriptor.name: candidate for candidate in loaded.values()} + for candidate in resolved: + existing_digest = loaded.get(candidate.descriptor.sha256) + if existing_digest is not None: + if existing_digest.identity != candidate.identity: + _fail( + "LOADED_IDENTITY_CONFLICT", + f"digest {candidate.descriptor.sha256} is already cached as {existing_digest.identity}", + ) + continue + existing_name = loaded_by_name.get(candidate.descriptor.name) + if existing_name is not None: + _fail( + "LOADED_NAME_CONFLICT", + f"{candidate.descriptor.name} is already loaded as {existing_name.identity}", + ) + if candidate.descriptor.name in database_loaded_names: + _fail( + "LOADED_OUTSIDE_RESOLVER", + f"{candidate.descriptor.name} is already loaded but its artifact identity is not resolver-cached", + ) def _provider_artifact(self, identity: str) -> LocalExtensionArtifact: candidates = [artifact for provider in self._providers if (artifact := provider.find(identity)) is not None] @@ -486,7 +568,7 @@ def _verify_artifact( current_source_id: str, current_vane_version: str, current_platform: str, - ) -> None: + ) -> Path: if descriptor.trust_identity not in self._trusted_identities: _fail("TRUST_IDENTITY_UNTRUSTED", f"{descriptor.trust_identity} is not in trusted_identities") if descriptor.duckdb_source_id != current_source_id: @@ -509,13 +591,13 @@ def _verify_artifact( _fail("NAME_MISMATCH", f"{descriptor.identity} must use artifact filename {expected_filename}") if not artifact_path.is_file(): _fail("ARTIFACT_NOT_FOUND", f"artifact does not exist: {artifact_path}") - actual_digest = _sha256_file(artifact_path) + snapshot_path, actual_digest, footer_bytes = self._snapshot_artifact(artifact_path) if actual_digest != descriptor.sha256: _fail( "DIGEST_MISMATCH", f"{descriptor.identity} expected SHA-256 {descriptor.sha256}, got {actual_digest}", ) - footer = _parse_extension_footer(artifact_path) + footer = _parse_extension_footer(footer_bytes, artifact_path.name) if footer.abi_type != descriptor.abi_type: _fail( "ABI_MISMATCH", @@ -538,11 +620,31 @@ def _verify_artifact( f"{descriptor.identity} requires C API {descriptor.duckdb_capi_version}, " f"artifact footer has {footer.engine_identity}", ) + runtime_capi_version = _runtime_capi_version() + if not _is_supported_capi_version(descriptor.duckdb_capi_version, runtime_capi_version): + _fail( + "CAPI_VERSION_MISMATCH", + f"{descriptor.identity} requires C API {descriptor.duckdb_capi_version}, " + f"runtime supports up through {runtime_capi_version}", + ) elif footer.engine_identity != descriptor.duckdb_source_id: _fail( "SOURCE_ID_MISMATCH", f"{descriptor.identity} footer SourceID is {footer.engine_identity}", ) + return snapshot_path + + def _snapshot_artifact(self, artifact_path: Path) -> tuple[Path, str, bytes]: + snapshot_parent = Path(tempfile.mkdtemp(prefix="artifact-", dir=self._snapshot_directory.name)) + snapshot_path = snapshot_parent / artifact_path.name + digest, footer_bytes = _inspect_extension_artifact(artifact_path, snapshot_path=snapshot_path) + try: + snapshot_path.chmod(0o400) + except OSError as exception: + raise DynamicExtensionError( + "ARTIFACT_SNAPSHOT_FAILED", f"could not make verified artifact snapshot read-only: {artifact_path}" + ) from exception + return snapshot_path, digest, footer_bytes def create_dynamic_extension_descriptor( @@ -559,7 +661,8 @@ def create_dynamic_extension_descriptor( expected_filename = f"{validated_name}.duckdb_extension" if artifact_path.name != expected_filename: _fail("NAME_MISMATCH", f"descriptor name {validated_name} requires artifact filename {expected_filename}") - footer = _parse_extension_footer(artifact_path) + artifact_digest, footer_bytes = _inspect_extension_artifact(artifact_path) + footer = _parse_extension_footer(footer_bytes, artifact_path.name) source_id, runtime_vane_version = _runtime_identity() descriptor_vane_version = ( runtime_vane_version if vane_version is None else _require_string(vane_version, "vane_version") @@ -572,7 +675,7 @@ def create_dynamic_extension_descriptor( duckdb_source_id=source_id, vane_version=descriptor_vane_version, platform=footer.platform, - sha256=_sha256_file(artifact_path), + sha256=artifact_digest, trust_identity=trust_identity, dependencies=tuple(dependencies), duckdb_capi_version=footer.engine_identity, @@ -584,7 +687,7 @@ def create_dynamic_extension_descriptor( duckdb_source_id=_validate_source_id(footer.engine_identity, "artifact footer SourceID"), vane_version=descriptor_vane_version, platform=footer.platform, - sha256=_sha256_file(artifact_path), + sha256=artifact_digest, trust_identity=trust_identity, dependencies=tuple(dependencies), ) @@ -602,29 +705,60 @@ def _connection_platform(connection: _ExtensionConnection) -> str: return _validate_platform(row[0]) -def _sha256_file(path: Path) -> str: - digest = hashlib.sha256() +def _database_loaded_extension_names(connection: _ExtensionConnection) -> frozenset[str]: try: - with path.open("rb") as artifact_file: - for chunk in iter(lambda: artifact_file.read(1024 * 1024), b""): - digest.update(chunk) - except OSError as exception: - raise DynamicExtensionError("ARTIFACT_NOT_FOUND", f"could not read artifact: {path}") from exception - return digest.hexdigest() + rows = connection.execute( + "SELECT extension_name FROM duckdb_extensions() WHERE loaded ORDER BY extension_name" + ).fetchall() + except Exception as exception: + raise DynamicExtensionError( + "LOADED_STATE_UNAVAILABLE", "could not query the database's loaded extension state" + ) from exception + if not isinstance(rows, list) or any( + not isinstance(row, tuple) or len(row) != 1 or not isinstance(row[0], str) for row in rows + ): + _fail("LOADED_STATE_UNAVAILABLE", "duckdb_extensions() returned an invalid loaded extension result") + return frozenset(row[0] for row in rows) + + +def _is_supported_capi_version(required: str | None, runtime: str) -> bool: + required_major, required_minor, required_patch = _parse_capi_version(required, "duckdb_capi_version") + runtime_major, runtime_minor, runtime_patch = _parse_capi_version(runtime, "runtime DuckDB extension C API version") + if required_major != runtime_major: + return False + if required_minor != runtime_minor: + return required_minor < runtime_minor + return required_patch <= runtime_patch -def _parse_extension_footer(path: Path) -> _ExtensionFooter: +def _inspect_extension_artifact(path: Path, *, snapshot_path: Path | None = None) -> tuple[str, bytes]: + digest = hashlib.sha256() + footer = bytearray() try: with path.open("rb") as artifact_file: - artifact_file.seek(0, 2) - if artifact_file.tell() < _EXTENSION_FOOTER_SIZE: - _fail("FOOTER_INVALID", f"artifact is smaller than {_EXTENSION_FOOTER_SIZE} bytes: {path.name}") - artifact_file.seek(-_EXTENSION_FOOTER_SIZE, 2) - footer = artifact_file.read(_EXTENSION_FOOTER_SIZE) + if snapshot_path is None: + for chunk in iter(lambda: artifact_file.read(1024 * 1024), b""): + digest.update(chunk) + footer.extend(chunk) + del footer[:-_EXTENSION_FOOTER_SIZE] + else: + with snapshot_path.open("xb") as snapshot_file: + for chunk in iter(lambda: artifact_file.read(1024 * 1024), b""): + digest.update(chunk) + footer.extend(chunk) + del footer[:-_EXTENSION_FOOTER_SIZE] + snapshot_file.write(chunk) + snapshot_file.flush() + os.fsync(snapshot_file.fileno()) except OSError as exception: - raise DynamicExtensionError("ARTIFACT_NOT_FOUND", f"could not read artifact: {path}") from exception + code = "ARTIFACT_NOT_FOUND" if snapshot_path is None else "ARTIFACT_SNAPSHOT_FAILED" + raise DynamicExtensionError(code, f"could not read and snapshot artifact: {path}") from exception + return digest.hexdigest(), bytes(footer) + + +def _parse_extension_footer(footer: bytes, artifact_name: str) -> _ExtensionFooter: if len(footer) != _EXTENSION_FOOTER_SIZE: - _fail("FOOTER_INVALID", f"could not read a complete extension footer from {path.name}") + _fail("FOOTER_INVALID", f"artifact is smaller than {_EXTENSION_FOOTER_SIZE} bytes: {artifact_name}") fields = [] for field_index in range(_EXTENSION_FOOTER_FIELD_COUNT): @@ -641,7 +775,7 @@ def _parse_extension_footer(path: Path) -> _ExtensionFooter: fields.append(field) if fields[7] != "4": - _fail("FOOTER_INVALID", f"artifact footer magic is invalid: {path.name}") + _fail("FOOTER_INVALID", f"artifact footer magic is invalid: {artifact_name}") abi_type = fields[3] or "CPP" if abi_type not in _VALID_ABI_TYPES: _fail("FOOTER_INVALID", f"artifact footer has unsupported ABI {abi_type!r}") From 445c9950c6e8de5f04ad273bd49f58aa2aa27078 Mon Sep 17 00:00:00 2001 From: kaka11chen Date: Wed, 26 Aug 2026 00:49:26 +0800 Subject: [PATCH 4/4] [test](extensions) Stabilize weak cache coverage --- tests/fast/test_dynamic_extension_resolver.py | 17 +++-------------- 1 file changed, 3 insertions(+), 14 deletions(-) diff --git a/tests/fast/test_dynamic_extension_resolver.py b/tests/fast/test_dynamic_extension_resolver.py index 02c6efa7b89..d8259a9e61b 100644 --- a/tests/fast/test_dynamic_extension_resolver.py +++ b/tests/fast/test_dynamic_extension_resolver.py @@ -242,26 +242,15 @@ def replace_provider_path(_snapshot_path): assert artifact_path.read_bytes() != verified_payload -def test_resolver_cache_does_not_keep_connections_alive(tmp_path): - platform = _runtime_platform() - artifact_path = _write_extension_artifact( - tmp_path / "root.duckdb_extension", - platform=platform, - source_id=vane.__git_revision__, - ) - descriptor = _descriptor(artifact_path, name="root") +def test_resolver_cache_does_not_keep_connections_alive(): resolver = DynamicExtensionResolver(trusted_identities={"local-tests"}) - connection = RecordingConnection(platform) + connection = RecordingConnection("linux_amd64") connection_reference = weakref.ref(connection) - resolver.load(connection, descriptor, artifact=artifact_path) + resolver._loaded_by_connection[connection] = {} assert len(resolver._loaded_by_connection) == 1 del connection - # The pytest timing plugin can retain the most recently returned Python - # frames while output capture is active. Advance it past load()/resolve() - # before checking that the resolver cache itself does not retain the key. - resolver.loaded_identities(RecordingConnection(platform)) gc.collect() assert connection_reference() is None