From 4ef123d764cbb0bad28b13f08aef5759e8098aa5 Mon Sep 17 00:00:00 2001 From: Chris Lu Date: Thu, 29 Oct 2026 13:35:05 -0700 Subject: [PATCH] s3vectors: fix shared API behavior and test Go/Rust volume parity (#12216) * s3vectors: reuse shared authenticated identity context * s3vectors: encode empty results as arrays * test: cover S3 Vectors with Go and Rust volume servers * ci: require both volume implementations in vector integration tests * s3vectors: keep identityOf as a shim over the shared identity context Concurrent branches add identityOf call sites; deleting the helper would break the build when they merge in either order. Route it through GetIdentityNameFromContext so every caller resolves the same identity. * test: wait for the vector engine to register before running The gateway resolves weed-vector lazily through the master cluster list, so a listening port does not mean the engine is reachable yet; the first vector write or query could race KeepConnected and fail. Restore the cluster.ps registration wait the previous harness had, and ignore the pytest cache directory the suite creates. * test: wait for master registration across OSS and enterprise Use the exact engine registration event after AddClusterNode completes. The current enterprise cluster.ps does not list vector nodes, so waiting on its output prevents the shared fixture from starting there. * s3vectors: keep contextWithIdentity as a shim alongside identityOf In-flight callers still write the caller name through it; route the write through SetIdentityNameInContext so both helpers agree on the shared key whichever lands first. --- .github/workflows/rust-vector-tests.yml | 68 +++++++-- .gitignore | 5 + test/vectors/Makefile | 32 +++-- test/vectors/README.md | 23 ++++ test/vectors/conftest.py | 157 +++++++++++++++++++++ test/vectors/requirements.txt | 2 + test/vectors/run_test.sh | 63 +-------- test/vectors/s3vectors_test.py | 174 +++++++++++++++++++++--- weed/s3api/s3vectors/auth_test.go | 6 +- weed/s3api/s3vectors/handlers.go | 6 +- weed/s3api/s3vectors/server.go | 2 +- weed/s3api/s3vectors/types.go | 11 +- 12 files changed, 431 insertions(+), 118 deletions(-) create mode 100644 test/vectors/README.md create mode 100644 test/vectors/conftest.py create mode 100644 test/vectors/requirements.txt diff --git a/.github/workflows/rust-vector-tests.yml b/.github/workflows/rust-vector-tests.yml index 9bad9cefa3..fbe3e01a9a 100644 --- a/.github/workflows/rust-vector-tests.yml +++ b/.github/workflows/rust-vector-tests.yml @@ -5,6 +5,7 @@ on: branches: [ master ] paths: - 'seaweed-vector/**' + - 'seaweed-volume/**' - 'seaweed-worker/**' - 'seaweed-common/**' - 'weed/pb/plugin.proto' @@ -18,6 +19,7 @@ on: branches: [ master, main ] paths: - 'seaweed-vector/**' + - 'seaweed-volume/**' - 'seaweed-worker/**' - 'seaweed-common/**' - 'weed/pb/plugin.proto' @@ -87,8 +89,15 @@ jobs: vectors-e2e: - name: S3 Vectors End-to-End - runs-on: ubuntu-latest + name: S3 Vectors Go/Rust Volume Parity + runs-on: ubuntu-24.04 + timeout-minutes: 45 + env: + CARGO_PROFILE_DEV_DEBUG: "0" + CARGO_BUILD_JOBS: "2" + CARGO_TARGET_DIR: ${{ github.workspace }}/seaweed-vector/target + VOLUME_SERVER_IMPL: go,rust + VECTOR_QUERY_MODE: single steps: - uses: actions/checkout@v7 with: @@ -96,11 +105,22 @@ jobs: - uses: actions/setup-go@v5 with: go-version-file: go.mod - - name: Build weed - run: go build -o test/vectors/weed ./weed + cache: ${{ runner.environment == 'github-hosted' }} - uses: dtolnay/rust-toolchain@stable - # Same as the unit-test job: lance's build scripts need a protoc; use - # the vendored one so no system package is required. + - name: Fingerprint build toolchain + id: toolchain + run: echo "fingerprint=$(getconf GNU_LIBC_VERSION | tr ' ' '-')-rustc-$(rustc -V | awk '{print $2}')" >> "$GITHUB_OUTPUT" + - name: Cache cargo registry and target + if: runner.environment == 'github-hosted' + uses: actions/cache@v6 + with: + path: | + ~/.cargo/registry + ~/.cargo/git + seaweed-vector/target/debug + key: vectors-e2e-${{ runner.os }}-${{ runner.arch }}-${{ steps.toolchain.outputs.fingerprint }}-${{ hashFiles('seaweed-vector/Cargo.lock', 'seaweed-volume/Cargo.lock') }} + restore-keys: | + vectors-e2e-${{ runner.os }}-${{ runner.arch }}-${{ steps.toolchain.outputs.fingerprint }}- - name: Use the vendored protoc run: | cd seaweed-vector @@ -110,15 +130,37 @@ jobs: protoc=$(find ~/.cargo/registry/src -path "*protoc-bin-vendored-linux-x86_64-$version/bin/protoc" | head -1) test -x "$protoc" || { echo "no vendored protoc $version in the registry" >&2; exit 1; } echo "PROTOC=$protoc" >> "$GITHUB_ENV" - - name: Build weed-vector-server - run: cargo build --release --manifest-path seaweed-vector/Cargo.toml -p weed-vector-server + - name: Test Go vector gateway + run: go test ./weed/s3api/s3vectors ./weed/command -run 'Vector|Vectors' -count=1 + - name: Build integration binaries + working-directory: test/vectors + run: make build CARGO_PROFILE=dev - uses: actions/setup-python@v5 with: python-version: "3.12" - - name: Install test deps - run: pip install boto3 pytest - - name: Run the s3vectors conformance subset + - name: Install test dependencies + run: python -m pip install -r test/vectors/requirements.txt + - name: Run the S3 Vectors integration matrix working-directory: test/vectors run: | - cp ../../seaweed-vector/target/release/weed-vector weed-vector - ./run_test.sh + ./run_test.sh --basetemp="$RUNNER_TEMP/vector-tests" --junitxml="$RUNNER_TEMP/vector-results.xml" + python - <<'PYTHON' + import os + from pathlib import Path + import xml.etree.ElementTree as ET + cases = ET.parse(Path(os.environ["RUNNER_TEMP"]) / "vector-results.xml").findall(".//testcase") + assert len(cases) >= 24, f"expected at least 24 cases, got {len(cases)}" + assert not any(case.find("skipped") is not None for case in cases), "integration cases were skipped" + for volume in ("go", "rust"): + assert sum(case.get("name", "").endswith(f"[{volume}-single]") for case in cases) >= 12, volume + PYTHON + - name: Upload integration diagnostics + if: always() + uses: actions/upload-artifact@v7 + with: + name: vectors-e2e-diagnostics + path: | + ${{ runner.temp }}/vector-tests + ${{ runner.temp }}/vector-results.xml + if-no-files-found: ignore + retention-days: 7 diff --git a/.gitignore b/.gitignore index c383623df9..fac1150200 100644 --- a/.gitignore +++ b/.gitignore @@ -145,3 +145,8 @@ telemetry/server/telemetry-server .aider* /seaweed-volume/docs __pycache__/ +/test/vectors/weed +/test/vectors/weed-vector +/test/vectors/weed-vector-worker +/test/vectors/weed-volume +.pytest_cache/ diff --git a/test/vectors/Makefile b/test/vectors/Makefile index fe222c9fc0..4d19b0eab2 100644 --- a/test/vectors/Makefile +++ b/test/vectors/Makefile @@ -1,28 +1,42 @@ -# End-to-end S3 Vectors conformance run: -# weed mini (cluster + lance catalog) + weed s3 (s3vectors port) + -# weed-vector (VectorService gRPC), then boto3 s3vectors calls. WEED := ./weed VEC := ./weed-vector +WORKER_MANIFEST := ../../rust/sw-vector-distributed/Cargo.toml +ifneq ($(wildcard $(WORKER_MANIFEST)),) +WORKER := ./weed-vector-worker +endif +VOLUME := ./weed-volume +CARGO_PROFILE ?= release +CARGO_TARGET_DIR ?= $(abspath ../../seaweed-vector/target) +CARGO_OUTPUT := $(CARGO_TARGET_DIR)/$(if $(filter dev,$(CARGO_PROFILE)),debug,$(CARGO_PROFILE)) +export CARGO_TARGET_DIR .PHONY: all test build clean FORCE all: test -# No prerequisites: go/cargo decide for themselves whether a rebuild is needed. $(WEED): FORCE go build -o $(WEED) ../../weed $(VEC): FORCE - cargo build --release --manifest-path ../../seaweed-vector/Cargo.toml -p weed-vector-server - cp ../../seaweed-vector/target/release/weed-vector $(VEC) + cargo build --locked --profile $(CARGO_PROFILE) --manifest-path ../../seaweed-vector/Cargo.toml -p weed-vector-server + cp $(CARGO_OUTPUT)/weed-vector $(VEC) + +ifneq ($(WORKER),) +$(WORKER): $(VEC) FORCE + cargo build --locked --profile $(CARGO_PROFILE) --manifest-path ../../rust/Cargo.toml -p sw-vector-distributed --bin weed-vector + cp $(CARGO_OUTPUT)/weed-vector $(WORKER) +endif + +$(VOLUME): $(VEC) $(WORKER) FORCE + cargo build --locked --profile $(CARGO_PROFILE) --manifest-path ../../seaweed-volume/Cargo.toml --bin weed-volume + cp $(CARGO_OUTPUT)/weed-volume $(VOLUME) FORCE: -build: $(WEED) $(VEC) +build: $(WEED) $(VEC) $(WORKER) $(VOLUME) test: build ./run_test.sh clean: - rm -f $(WEED) $(VEC) - rm -rf mini-data + rm -f $(WEED) $(VEC) ./weed-vector-worker $(VOLUME) diff --git a/test/vectors/README.md b/test/vectors/README.md new file mode 100644 index 0000000000..a0efbea15f --- /dev/null +++ b/test/vectors/README.md @@ -0,0 +1,23 @@ +# S3 Vectors integration tests + +Run `python3 -m pip install -r requirements.txt`, then `make test`. + +The suite starts a master, volume, filer, authenticated S3/Lance gateway, and +vector engine on temporary ports. The same 12 API tests run with both Go and +Rust volume servers. Coverage includes bucket/index lifecycle, bucket policies, +empty results, ARN addressing, metadata, filters, cosine distance, atomic upserts, +pagination, segmented listing, authentication, and writer loss. + +Server logs and data stay in pytest's temporary directory for diagnosis. Every +child process stops on exit. To select a volume server, set +`VOLUME_SERVER_IMPL=go` or `rust`. + +In an enterprise checkout containing `rust/sw-vector-distributed`, the build and +default matrix also include distributed queries with two search workers. Those +queries must still succeed after the writer stops. `VECTOR_QUERY_MODE=single` +or `distributed` selects one query mode. + +To use prebuilt binaries, set `WEED_BINARY`, `VECTOR_BINARY`, +`RUST_VOLUME_SERVER_BINARY`, and (for distributed queries) `VECTOR_WORKER_BINARY`, +then run `./run_test.sh`. `PYTHON` selects the Python interpreter. The Makefile +accepts `CARGO_PROFILE=dev` for faster local builds. diff --git a/test/vectors/conftest.py b/test/vectors/conftest.py new file mode 100644 index 0000000000..b395f13136 --- /dev/null +++ b/test/vectors/conftest.py @@ -0,0 +1,157 @@ +import contextlib +import os +from pathlib import Path +import socket +import subprocess +import time +import urllib.error +import urllib.request + +import boto3 +from botocore.config import Config +import pytest + +HERE = Path(__file__).resolve().parent + + +def combinations(): + volumes = os.environ.get("VOLUME_SERVER_IMPL", "go,rust").split(",") + default_modes = "single,distributed" if (HERE / "../../rust/sw-vector-distributed/Cargo.toml").is_file() else "single" + modes = os.environ.get("VECTOR_QUERY_MODE", default_modes).split(",") + if not set(volumes) <= {"go", "rust"} or not set(modes) <= {"single", "distributed"}: + raise ValueError("invalid VOLUME_SERVER_IMPL or VECTOR_QUERY_MODE") + return [(volume, mode) for volume in volumes for mode in modes] + + +def binary(variable, name): + path = Path(os.environ.get(variable, str(HERE / name))).resolve() + if not path.is_file() or not os.access(path, os.X_OK): + pytest.fail(f"Missing executable {path}; run make build or set {variable}") + return str(path) + + +def wait_for(check, processes, message): + deadline = time.monotonic() + 60 + while time.monotonic() < deadline: + for name, process in processes.items(): + if process.poll() is not None: + pytest.fail(f"{name} exited with {process.returncode}; {message}") + try: + if check(): + return + except (OSError, urllib.error.URLError): + pass + time.sleep(0.2) + pytest.fail(message) + + +def listening(port): + with socket.create_connection(("127.0.0.1", port), timeout=0.5): + return True + + +def http_ready(url): + try: + with urllib.request.urlopen(url, timeout=0.5): + return True + except urllib.error.HTTPError: + return True + + +def stop(process): + if process.poll() is None: + process.terminate() + try: + process.wait(timeout=10) + except subprocess.TimeoutExpired: + process.kill() + process.wait(timeout=10) + + +@pytest.fixture(scope="module", params=combinations(), ids=lambda p: "-".join(p)) +def cluster(request, tmp_path_factory): + volume, mode = request.param + directory = tmp_path_factory.mktemp(f"vectors-{volume}-{mode}") + print(f"\nVector integration logs: {directory}") + weed = binary("WEED_BINARY", "weed") + engine = binary("VECTOR_BINARY", "weed-vector") + worker = binary("VECTOR_WORKER_BINARY", "weed-vector-worker") if mode == "distributed" else None + volume_binary = binary("RUST_VOLUME_SERVER_BINARY", "weed-volume") if volume == "rust" else weed + processes = {} + with contextlib.ExitStack() as stack: + with contextlib.ExitStack() as reservations: + ports = [] + for _ in range(12): + sock = reservations.enter_context(socket.socket()) + sock.bind(("127.0.0.1", 0)) + ports.append(sock.getsockname()[1]) + master, master_grpc, vol, vol_grpc, filer, filer_grpc, s3, s3_grpc, lance, vector, worker1, worker2 = ports + master_addr = f"127.0.0.1:{master}.{master_grpc}" + filer_addr = f"127.0.0.1:{filer}.{filer_grpc}" + + def start(name, args): + log = stack.enter_context((directory / f"{name}.log").open("w")) + env = dict(os.environ, AWS_ACCESS_KEY_ID="test", AWS_SECRET_ACCESS_KEY="test", + AWS_DEFAULT_REGION="us-east-1", AWS_EC2_METADATA_DISABLED="true") + process = subprocess.Popen(args, cwd=directory, stdout=log, stderr=log, env=env) + processes[name] = process + stack.callback(stop, process) + return process + + for name in ("master", "volume", "filer"): + (directory / name).mkdir() + start("master", [weed, "-v=1", "master", "-ip=127.0.0.1", "-ip.bind=127.0.0.1", + f"-port={master}", f"-port.grpc={master_grpc}", + f"-mdir={directory / 'master'}", "-peers=none", "-volumeSizeLimitMB=32"]) + wait_for(lambda: http_ready(f"http://127.0.0.1:{master}/dir/status"), processes, "master not ready") + prefix = [volume_binary] if volume == "rust" else [weed, "volume"] + flag = "--" if volume == "rust" else "-" + start("volume", prefix + [f"{flag}{key}={value}" for key, value in { + "master": master_addr, "ip": "127.0.0.1", "ip.bind": "127.0.0.1", + "port": vol, "port.grpc": vol_grpc, "dir": directory / "volume", + "max": 16, "preStopSeconds": 0, + }.items()]) + wait_for(lambda: listening(vol_grpc), processes, "volume not ready") + start("filer", [weed, "filer", f"-master={master_addr}", "-ip=127.0.0.1", + "-ip.bind=127.0.0.1", f"-port={filer}", f"-port.grpc={filer_grpc}", + f"-defaultStoreDir={directory / 'filer'}"]) + wait_for(lambda: listening(filer_grpc), processes, "filer not ready") + endpoint = f"http://127.0.0.1:{s3}" + start("s3", [weed, "-v=1", "s3", f"-filer={filer_addr}", f"-config={HERE / 's3.json'}", + "-ip=127.0.0.1", "-ip.bind=127.0.0.1", f"-port={s3}", + f"-port.grpc={s3_grpc}", "-port.iceberg=0", f"-port.lance={lance}"]) + wait_for(lambda: http_ready(f"http://127.0.0.1:{lance}/v1/namespace/root/list"), processes, "catalog not ready") + start("vector", [engine, "--namespace", f"http://127.0.0.1:{lance}", + "--masters", master_addr, "--listen", f"127.0.0.1:{vector}", + "--storage-option", "aws_access_key_id=test", + "--storage-option", "aws_secret_access_key=test"]) + wait_for(lambda: listening(vector), processes, "vector engine not ready") + + wait_for( + lambda: f"Cluster: vector node 127.0.0.1:{vector} added to group ''" + in (directory / "master.log").read_text(), + processes, + "vector engine never registered with the master", + ) + + if worker: + for i, port in enumerate((worker1, worker2)): + start(f"worker-{i}", [worker, "--masters", master_addr, "--listen", f"127.0.0.1:{port}"]) + wait_for(lambda: listening(port), processes, "vector worker not ready") + wait_for( + lambda: all( + f"S3 Vectors: discovered search worker 127.0.0.1:{port}" + in (directory / "s3.log").read_text() + for port in (worker1, worker2) + ), + processes, + "gateway did not discover both search workers", + ) + yield {"endpoint": endpoint, "mode": mode, "processes": processes, "directory": directory} + + +@pytest.fixture(scope="module") +def vectors(cluster): + return boto3.client("s3vectors", endpoint_url=cluster["endpoint"], + aws_access_key_id="test", aws_secret_access_key="test", region_name="us-east-1", + config=Config(connect_timeout=3, read_timeout=20, retries={"max_attempts": 0})) diff --git a/test/vectors/requirements.txt b/test/vectors/requirements.txt new file mode 100644 index 0000000000..17d6fd2b49 --- /dev/null +++ b/test/vectors/requirements.txt @@ -0,0 +1,2 @@ +boto3>=1.40 +pytest>=8 diff --git a/test/vectors/run_test.sh b/test/vectors/run_test.sh index 3a44eb9009..cc8f14624a 100755 --- a/test/vectors/run_test.sh +++ b/test/vectors/run_test.sh @@ -1,65 +1,4 @@ #!/usr/bin/env bash -# Spin a mini cluster plus the vectors gateway and engine, run the boto3 -# conformance subset, tear everything down. set -euo pipefail cd "$(dirname "$0")" - -DATA=mini-data -rm -rf "$DATA" -PIDS=() -cleanup() { - kill "${PIDS[@]}" 2>/dev/null || true - rm -rf "$DATA" -} -trap cleanup EXIT - -./weed mini -dir "$DATA" \ - -s3=false -webdav=false -admin.ui=false & -PIDS+=($!) - -./weed s3 \ - -filer 127.0.0.1:8888 \ - -config s3.json \ - -ip 127.0.0.1 \ - -ip.bind 127.0.0.1 \ - -port 8334 \ - -port.lance 9102 & -PIDS+=($!) - -./weed-vector \ - --namespace http://127.0.0.1:9102 \ - --masters 127.0.0.1:9333 \ - --storage-option aws_access_key_id=test \ - --storage-option aws_secret_access_key=test \ - --listen 127.0.0.1:9301 & -PIDS+=($!) - -# Any HTTP response (even an auth rejection) means the lance catalog is up. -for _ in $(seq 1 60); do - code=$(curl -s -o /dev/null -w '%{http_code}' http://127.0.0.1:9102/v1/namespace/root/list || true) - [ "$code" != "000" ] && break - sleep 1 -done -[ "${code:-000}" != "000" ] || { echo "lance catalog on :9102 never came up" >&2; exit 1; } - -for _ in $(seq 1 60); do - (exec 3<>/dev/tcp/127.0.0.1/9301) 2>/dev/null && break - sleep 1 -done -(exec 3<>/dev/tcp/127.0.0.1/9301) 2>/dev/null || { echo "vector server on :9301 never came up" >&2; exit 1; } - -# A listening engine still has to land its KeepConnected registration before -# the gateway's lazy discovery can find it — poll the cluster list for the -# vector node rather than race the first request against it. -vector_registered() { - # grep -q would close the pipe on first match and SIGPIPE the shell mid-write - # under `set -o pipefail`; a full read lets cluster.ps finish. - printf 'cluster.ps\n' | ./weed shell -master 127.0.0.1:9333 2>/dev/null | grep "vector servers" >/dev/null -} -for _ in $(seq 1 60); do - vector_registered && break - sleep 1 -done -vector_registered || { echo "vector server never registered with the master" >&2; exit 1; } - -python3 -m pytest -xvs s3vectors_test.py +exec "${PYTHON:-python3}" -m pytest -xvs s3vectors_test.py "$@" diff --git a/test/vectors/s3vectors_test.py b/test/vectors/s3vectors_test.py index f9866d7a25..4204c93370 100644 --- a/test/vectors/s3vectors_test.py +++ b/test/vectors/s3vectors_test.py @@ -1,31 +1,17 @@ -"""Boto3 conformance subset for the S3 Vectors API gateway. - -Run via `make test` — a mini cluster, an s3 gateway serving the vectors API on its main port, and a -weed-vector are up by then. -""" +"""Run the S3 Vectors API against Go/Rust volumes and single/distributed queries.""" import json +import time +import uuid import boto3 -import pytest from botocore.exceptions import ClientError +import pytest -ENDPOINT = "http://127.0.0.1:8334" BUCKET = "vectors-bucket" INDEX = "idx" -@pytest.fixture(scope="module") -def vectors(): - return boto3.client( - "s3vectors", - endpoint_url=ENDPOINT, - aws_access_key_id="test", - aws_secret_access_key="test", - region_name="us-east-1", - ) - - def test_vector_lifecycle(vectors): vectors.create_vector_bucket(vectorBucketName=BUCKET) buckets = [b["vectorBucketName"] for b in vectors.list_vector_buckets()["vectorBuckets"]] @@ -68,8 +54,7 @@ def test_vector_lifecycle(vectors): returnMetadata=True, filter={"color": {"$eq": "red"}}, )["vectors"] - # b is the second-nearest vector but blue, so an ignored filter would - # return ["a", "b"] here. + # b is closer than c but excluded by the filter. assert [h["key"] for h in hits] == ["a", "c"] distances = {h["key"]: h["distance"] for h in hits} assert distances["a"] == pytest.approx(0.0, abs=1e-6) @@ -110,3 +95,152 @@ def test_vector_bucket_policy(vectors): assert exc.value.response["Error"]["Code"] == "NotFoundException" finally: vectors.delete_vector_bucket(vectorBucketName=BUCKET) + + +@pytest.fixture +def index(vectors): + bucket = "vectors-" + uuid.uuid4().hex[:16] + vectors.create_vector_bucket(vectorBucketName=bucket) + vectors.create_index(vectorBucketName=bucket, indexName=INDEX, dataType="float32", + dimension=4, distanceMetric="euclidean", + metadataConfiguration={"nonFilterableMetadataKeys": ["private"]}) + args = {"vectorBucketName": bucket, "indexName": INDEX} + yield args + vectors.delete_index(**args) + vectors.delete_vector_bucket(vectorBucketName=bucket) + + +def expect_error(code, operation, **kwargs): + with pytest.raises(ClientError) as exc: + operation(**kwargs) + assert exc.value.response["Error"]["Code"] == code + + +def test_empty_catalog(vectors): + bucket = "empty-" + uuid.uuid4().hex[:16] + assert vectors.list_vector_buckets(prefix=bucket)["vectorBuckets"] == [] + vectors.create_vector_bucket(vectorBucketName=bucket) + try: + assert vectors.list_indexes(vectorBucketName=bucket)["indexes"] == [] + finally: + vectors.delete_vector_bucket(vectorBucketName=bucket) + + +def test_empty_index(vectors, index): + assert vectors.get_vectors(**index, keys=["absent"])["vectors"] == [] + assert vectors.list_vectors(**index)["vectors"] == [] + assert vectors.query_vectors(**index, queryVector={"float32": [1, 0, 0, 0]}, topK=3)["vectors"] == [] + expect_error("ValidationException", vectors.query_vectors, **index, + queryVector={"float32": [1, 0, 0]}, topK=3) + + +def test_upsert_and_invalid_batch(vectors, index): + vectors.put_vectors(**index, vectors=[{"key": "a", "data": {"float32": [1, 0, 0, 0]}, "metadata": {"revision": 1}}]) + vectors.put_vectors(**index, vectors=[{"key": "a", "data": {"float32": [0, 1, 0, 0]}, "metadata": {"revision": 2}}]) + expected = [{"key": "a", "data": {"float32": [0, 1, 0, 0]}, "metadata": {"revision": 2}}] + for invalid in ( + [{"key": "a", "data": {"float32": [1, 1, 1, 1]}}, {"key": "bad", "data": {"float32": [1, 2]}}], + [{"key": "a", "data": {"float32": [1, 1, 1, 1]}}, {"key": "a", "data": {"float32": [2, 2, 2, 2]}}], + ): + expect_error("ValidationException", vectors.put_vectors, **index, vectors=invalid) + assert vectors.list_vectors(**index, returnData=True, returnMetadata=True)["vectors"] == expected + + +def test_metadata_and_arn(vectors, index): + description = vectors.get_index(**index)["index"] + assert description["metadataConfiguration"] == {"nonFilterableMetadataKeys": ["private"]} + arn = {"indexArn": description["indexArn"]} + metadata = {"color": "red", "private": "secret", "nested": {"x": [1, True, "text"]}} + vectors.put_vectors(**arn, vectors=[{"key": "quote'key", "data": {"float32": [1, 0, 0, 0]}, "metadata": metadata}]) + assert vectors.get_vectors(**arn, keys=["quote'key"], returnMetadata=True)["vectors"][0]["metadata"] == metadata + hits = vectors.query_vectors(**arn, queryVector={"float32": [1, 0, 0, 0]}, topK=1, returnMetadata=True)["vectors"] + assert hits[0]["metadata"] == metadata + assert "distance" not in hits[0] + expect_error("ValidationException", vectors.query_vectors, **arn, queryVector={"float32": [1, 0, 0, 0]}, + topK=1, filter={"$and": [{"color": "red"}, {"private": "secret"}]}) + vectors.delete_vectors(**arn, keys=["quote'key"]) + assert vectors.list_vectors(**arn)["vectors"] == [] + + +def test_list_pagination_and_segments(vectors, index): + rows = [{"key": f"key-{i:02}", "data": {"float32": [i, 1, 0, 0]}, "metadata": {"row": i}} for i in range(17)] + vectors.put_vectors(**index, vectors=rows) + + def collect(**extra): + result, tokens = [], set() + while True: + page = vectors.list_vectors(**index, maxResults=2, returnData=True, returnMetadata=True, **extra) + result.extend(page["vectors"]) + token = page.get("nextToken") + if not token: + return result + assert token not in tokens + tokens.add(token) + extra["nextToken"] = token + + assert sorted(collect(), key=lambda r: r["key"]) == rows + segmented = [row for segment in range(4) for row in collect(segmentCount=4, segmentIndex=segment)] + assert sorted(segmented, key=lambda r: r["key"]) == rows + expect_error("ValidationException", vectors.list_vectors, **index, nextToken="invalid") + + +def test_filter_types(vectors, index): + values = [1, "1", True, "true"] + vectors.put_vectors(**index, vectors=[{"key": str(i), "data": {"float32": [1, 0, 0, 0]}, "metadata": {"value": value}} + for i, value in enumerate(values)]) + for i, value in enumerate(values): + hits = vectors.query_vectors(**index, queryVector={"float32": [1, 0, 0, 0]}, topK=4, + filter={"value": {"$eq": value}})["vectors"] + assert [hit["key"] for hit in hits] == [str(i)] + + +def test_cosine_metric(vectors): + bucket = "cosine-" + uuid.uuid4().hex[:16] + vectors.create_vector_bucket(vectorBucketName=bucket) + args = {"vectorBucketName": bucket, "indexName": INDEX} + vectors.create_index(**args, dataType="float32", dimension=2, distanceMetric="cosine") + try: + vectors.put_vectors(**args, vectors=[{"key": "far", "data": {"float32": [10, 0]}}, + {"key": "near", "data": {"float32": [1, 1]}}]) + result = vectors.query_vectors(**args, queryVector={"float32": [1, 0]}, topK=2, returnDistance=True) + assert result["distanceMetric"] == "cosine" + assert [v["key"] for v in result["vectors"]] == ["far", "near"] + assert result["vectors"][0]["distance"] == pytest.approx(0, abs=1e-6) + expect_error("ValidationException", vectors.query_vectors, **args, queryVector={"float32": [1, 0, 0]}, topK=1) + finally: + vectors.delete_index(**args) + vectors.delete_vector_bucket(vectorBucketName=bucket) + + +def test_delete_bucket_conflict(vectors, index): + expect_error("ConflictException", vectors.delete_vector_bucket, vectorBucketName=index["vectorBucketName"]) + assert vectors.get_index(**index)["index"]["dimension"] == 4 + + +def test_invalid_credentials(cluster): + client = boto3.client("s3vectors", endpoint_url=cluster["endpoint"], aws_access_key_id="unknown", + aws_secret_access_key="unknown", region_name="us-east-1") + with pytest.raises(ClientError) as exc: + client.list_vector_buckets() + assert exc.value.response["ResponseMetadata"]["HTTPStatusCode"] == 403 + + +def test_query_without_writer(vectors, index, cluster): + from conftest import stop + + vectors.put_vectors(**index, vectors=[{"key": "survivor", "data": {"float32": [1, 0, 0, 0]}}]) + stop(cluster["processes"].pop("vector")) + args = dict(index, queryVector={"float32": [1, 0, 0, 0]}, topK=1) + if cluster["mode"] == "single": + expect_error("ServiceUnavailableException", vectors.query_vectors, **args) + return + deadline = time.monotonic() + 15 + while True: + try: + result = vectors.query_vectors(**args) + break + except ClientError as exc: + if exc.response["Error"]["Code"] != "ServiceUnavailableException" or time.monotonic() >= deadline: + raise + time.sleep(0.2) + assert [v["key"] for v in result["vectors"]] == ["survivor"] diff --git a/weed/s3api/s3vectors/auth_test.go b/weed/s3api/s3vectors/auth_test.go index 43c5d8e26b..a333265d0d 100644 --- a/weed/s3api/s3vectors/auth_test.go +++ b/weed/s3api/s3vectors/auth_test.go @@ -20,15 +20,15 @@ func (f *fakeAuthenticator) AuthenticateRequest(r *http.Request) (string, interf func (f *fakeAuthenticator) DefaultAllow() bool { return true } -// The authenticated identity object must ride the request context — only the -// name survives otherwise, and downstream authorization loses actions, -// account, and claims. func TestAuthKeepsIdentityObject(t *testing.T) { want := &s3api.Identity{Name: "alice", Actions: []s3api.Action{"Admin"}} s := &Server{authenticator: &fakeAuthenticator{identity: want}} var got interface{} h := s.Auth(func(w http.ResponseWriter, r *http.Request) { got = s3_constants.GetIdentityFromContext(r) + if name := s3_constants.GetIdentityNameFromContext(r); name != want.Name { + t.Errorf("identity name = %q, want %q", name, want.Name) + } }) h(httptest.NewRecorder(), httptest.NewRequest(http.MethodPost, "/", nil)) if got != interface{}(want) { diff --git a/weed/s3api/s3vectors/handlers.go b/weed/s3api/s3vectors/handlers.go index a2630a2cce..0145952cca 100644 --- a/weed/s3api/s3vectors/handlers.go +++ b/weed/s3api/s3vectors/handlers.go @@ -362,7 +362,7 @@ func (s *Server) handleListVectorBuckets(w http.ResponseWriter, r *http.Request) } // The catalog token continues the filtered page: ICEBERG buckets drop out // of the listing, Lance and format-less legacy buckets stay. - out := &ListVectorBucketsResponse{NextToken: resp.ContinuationToken} + out := &ListVectorBucketsResponse{NextToken: resp.ContinuationToken, VectorBuckets: []*VectorBucket{}} for _, b := range resp.TableBuckets { if b.Format != "" && b.Format != s3tables.FormatLance { continue @@ -750,7 +750,7 @@ func (s *Server) handleListIndexes(w http.ResponseWriter, r *http.Request) { } // Summaries carry only what the catalog listing knows; a summary that // asserted a dimension or metric without a record lookup would lie. - out := &ListIndexesResponse{NextToken: resp.ContinuationToken} + out := &ListIndexesResponse{NextToken: resp.ContinuationToken, Indexes: []*IndexSummary{}} for _, t := range resp.TableSummaries { // A non-Lance table in the index namespace is not an index. if t.Format != "" && t.Format != s3tables.FormatLance { @@ -1412,7 +1412,7 @@ func (s *Server) handleQueryVectors(w http.ResponseWriter, r *http.Request) { vectorErr(w, err) return } - out := &QueryVectorsResponse{DistanceMetric: metric} + out := &QueryVectorsResponse{DistanceMetric: metric, Vectors: []*QueryOutputVector{}} for _, v := range resp.Vectors { o := &QueryOutputVector{Key: v.Key} if req.ReturnDistance { diff --git a/weed/s3api/s3vectors/server.go b/weed/s3api/s3vectors/server.go index 71130c1f50..8d0f22eeeb 100644 --- a/weed/s3api/s3vectors/server.go +++ b/weed/s3api/s3vectors/server.go @@ -97,7 +97,7 @@ func (s *Server) Auth(handler http.HandlerFunc) http.HandlerFunc { return } r = r.WithContext(s3_constants.SetIdentityInContext( - contextWithIdentity(r.Context(), identity), identityObj)) + s3_constants.SetIdentityNameInContext(r.Context(), identity), identityObj)) handler(w, r) } } diff --git a/weed/s3api/s3vectors/types.go b/weed/s3api/s3vectors/types.go index 087291f8c4..066208f47a 100644 --- a/weed/s3api/s3vectors/types.go +++ b/weed/s3api/s3vectors/types.go @@ -6,19 +6,16 @@ import ( "io" "net/http" "time" + + "github.com/seaweedfs/seaweedfs/weed/s3api/s3_constants" ) -type contextKey string - -const identityKey contextKey = "s3vectors-identity" - func contextWithIdentity(ctx context.Context, identity string) context.Context { - return context.WithValue(ctx, identityKey, identity) + return s3_constants.SetIdentityNameInContext(ctx, identity) } func identityOf(r *http.Request) string { - v, _ := r.Context().Value(identityKey).(string) - return v + return s3_constants.GetIdentityNameFromContext(r) } func writeJSON(w http.ResponseWriter, status int, payload interface{}) {