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.
This commit is contained in:
Chris Lu authored and GitHub committed 2026-10-29 13:35:05 -07:00
1 parent 88565df885
commit 4ef123d764
12 files changed
+431 -118

No files matched your search

+55 -13
View File
@@ -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
+5
View File
@@ -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/
+23 -9
View File
@@ -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)
+23
View File
@@ -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.
+157
View File
@@ -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}))
+2
View File
@@ -0,0 +1,2 @@
boto3>=1.40
pytest>=8
+1 -62
View File
@@ -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 "$@"
+154 -20
View File
@@ -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"]
+3 -3
View File
@@ -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) {
+3 -3
View File
@@ -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 {
+1 -1
View File
@@ -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)
}
}
+4 -7
View File
@@ -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{}) {