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{}) {