mirror of
https://github.com/seaweedfs/seaweedfs.git
synced 2026-11-03 18:47:41 +01:00
* s3vectors: tag vector buckets and indexes
TagResource, UntagResource and ListTagsForResource answer on
/tags/{resourceArn} (POST, DELETE, GET) and by name. Each authorizes its
own s3vectors action on the bucket or index ARN, then reads or writes the
catalog record's tags. A resource holds at most 50 tags.
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
* s3vectors: apply the tags passed to CreateVectorBucket and CreateIndex
The tags field was accepted and dropped. It is now stored on the new
bucket or index, and creating with tags needs s3vectors:TagResource in
addition to the create action.
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
* test/s3vectors: cover resource tagging end to end
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
* s3tables: keep tag operations to a bound table bucket
TagResource, UntagResource and ListTagsForResource run through Execute
with WithTableBucketCreatedAt now answer not found when the bucket name
holds another bucket, as the bound creates and deletes already do.
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
* s3vectors: bind tagging to the authorized bucket and to vector indexes
The catalog tag calls carry the creation time of the bucket authorizeVector
approved, so a bucket created again under the name is not touched. An index
ARN must name a Lance table: another table in the vector namespace of a
format-less bucket is refused, as DeleteIndex does.
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
* s3vectors: match a TagResource grant scoped to the bucket being created
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
* s3tables: merge tag writes into the tags the record holds at the write
TagResource and UntagResource merged into the tags they read before the
compare-and-swap and wrote that map on every attempt, so a tag write that
landed in between was undone. The merge now runs inside the mutation.
WithMaxTags lets an in-process caller cap the merged set there.
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
* s3vectors: check the 50-tag limit inside the catalog's tag write
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
* s3tables: check a tag operation's bound bucket after its reads
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
* s3vectors: validate each tag key UntagResource removes
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
* s3vectors: validate untag keys one at a time
UntagResource sized a map by the number of submitted tag keys before
validating them, so a large body of repeated keys allocated far more than
the request itself, before authorization ran.
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
* s3vectors: authorize UntagResource before validating its keys
Validating a near-32 MiB body of tag keys costs about 13x its JSON
decode, so an unauthorized caller could make the gateway spend that
work on a request it then refuses.
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
---------
Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com>
496 lines
25 KiB
Python
496 lines
25 KiB
Python
"""Run the S3 Vectors API against Go/Rust volumes and single/distributed queries."""
|
|
|
|
import contextlib
|
|
import json
|
|
import time
|
|
import uuid
|
|
|
|
import boto3
|
|
from botocore.config import Config
|
|
from botocore.exceptions import ClientError
|
|
import pytest
|
|
|
|
BUCKET = "vectors-bucket"
|
|
INDEX = "idx"
|
|
|
|
|
|
def test_vector_lifecycle(vectors):
|
|
vectors.create_vector_bucket(vectorBucketName=BUCKET)
|
|
buckets = [b["vectorBucketName"] for b in vectors.list_vector_buckets()["vectorBuckets"]]
|
|
assert BUCKET in buckets
|
|
|
|
vectors.create_index(
|
|
vectorBucketName=BUCKET,
|
|
indexName=INDEX,
|
|
dataType="float32",
|
|
dimension=4,
|
|
distanceMetric="euclidean",
|
|
)
|
|
idx = vectors.get_index(vectorBucketName=BUCKET, indexName=INDEX)["index"]
|
|
assert idx["dimension"] == 4
|
|
|
|
vectors.put_vectors(
|
|
vectorBucketName=BUCKET,
|
|
indexName=INDEX,
|
|
vectors=[
|
|
{"key": "a", "data": {"float32": [1, 0, 0, 0]}, "metadata": {"color": "red"}},
|
|
{"key": "b", "data": {"float32": [0.95, 0.05, 0, 0]}, "metadata": {"color": "blue"}},
|
|
{"key": "c", "data": {"float32": [0.9, 0.1, 0, 0]}, "metadata": {"color": "red"}},
|
|
],
|
|
)
|
|
|
|
got = vectors.get_vectors(
|
|
vectorBucketName=BUCKET, indexName=INDEX, keys=["a"], returnData=True, returnMetadata=True
|
|
)["vectors"]
|
|
assert got[0]["key"] == "a" and got[0]["metadata"]["color"] == "red"
|
|
|
|
listed = vectors.list_vectors(vectorBucketName=BUCKET, indexName=INDEX)["vectors"]
|
|
assert sorted(v["key"] for v in listed) == ["a", "b", "c"]
|
|
|
|
hits = vectors.query_vectors(
|
|
vectorBucketName=BUCKET,
|
|
indexName=INDEX,
|
|
queryVector={"float32": [1, 0, 0, 0]},
|
|
topK=2,
|
|
returnDistance=True,
|
|
returnMetadata=True,
|
|
filter={"color": {"$eq": "red"}},
|
|
)["vectors"]
|
|
# 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)
|
|
assert distances["c"] == pytest.approx(0.02 ** 0.5, abs=1e-4)
|
|
assert all(h["metadata"]["color"] == "red" for h in hits)
|
|
|
|
vectors.delete_vectors(vectorBucketName=BUCKET, indexName=INDEX, keys=["b"])
|
|
listed = vectors.list_vectors(vectorBucketName=BUCKET, indexName=INDEX)["vectors"]
|
|
assert sorted(v["key"] for v in listed) == ["a", "c"]
|
|
|
|
vectors.delete_index(vectorBucketName=BUCKET, indexName=INDEX)
|
|
vectors.delete_vector_bucket(vectorBucketName=BUCKET)
|
|
|
|
|
|
def test_vector_bucket_policy(vectors):
|
|
vectors.create_vector_bucket(vectorBucketName=BUCKET)
|
|
try:
|
|
policy = json.dumps(
|
|
{
|
|
"Version": "2012-10-17",
|
|
"Statement": [
|
|
{
|
|
"Effect": "Allow",
|
|
"Principal": {"AWS": "*"},
|
|
"Action": "s3vectors:*",
|
|
"Resource": "*",
|
|
}
|
|
],
|
|
}
|
|
)
|
|
invalid = json.loads(policy)
|
|
invalid["Statement"][0]["Action"] = "s3tables:*"
|
|
expect_error("ValidationException", vectors.put_vector_bucket_policy,
|
|
vectorBucketName=BUCKET, policy=json.dumps(invalid))
|
|
vectors.put_vector_bucket_policy(vectorBucketName=BUCKET, policy=policy)
|
|
got = vectors.get_vector_bucket_policy(vectorBucketName=BUCKET)["policy"]
|
|
assert json.loads(got) == json.loads(policy)
|
|
|
|
vectors.delete_vector_bucket_policy(vectorBucketName=BUCKET)
|
|
with pytest.raises(ClientError) as exc:
|
|
vectors.get_vector_bucket_policy(vectorBucketName=BUCKET)
|
|
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")
|
|
for bad in ({"maxResults": 1001}, {"segmentCount": 17, "segmentIndex": 0}, {"segmentIndex": 0}):
|
|
expect_error("ValidationException", vectors.list_vectors, **index, **bad)
|
|
|
|
|
|
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)]
|
|
for op, members in (("$in", [1, "true"]), ("$nin", ["1", True])):
|
|
hits = vectors.query_vectors(**index, queryVector={"float32": [1, 0, 0, 0]}, topK=4,
|
|
filter={"value": {op: members}})["vectors"]
|
|
assert sorted(hit["key"] for hit in hits) == ["0", "3"]
|
|
|
|
|
|
def test_filter_arrays(vectors, index):
|
|
rows = {"space": ["space", "classic"], "nums": [1, 2.5], "strs": ["1", "2.5"], "flag": [True], "scalar": "space"}
|
|
vectors.put_vectors(**index, vectors=[{"key": key, "data": {"float32": [1, 0, 0, 0]}, "metadata": {"tags": tags}}
|
|
for key, tags in rows.items()])
|
|
|
|
def matches(condition):
|
|
hits = vectors.query_vectors(**index, queryVector={"float32": [1, 0, 0, 0]}, topK=10,
|
|
filter={"tags": condition})["vectors"]
|
|
return sorted(hit["key"] for hit in hits)
|
|
|
|
assert matches("space") == ["scalar", "space"]
|
|
assert matches({"$eq": 1}) == ["nums"]
|
|
assert matches({"$eq": "1"}) == ["strs"]
|
|
assert matches({"$eq": True}) == ["flag"]
|
|
assert matches({"$in": ["classic", 2.5]}) == ["nums", "space"]
|
|
assert matches({"$ne": "space"}) == ["flag", "nums", "strs"]
|
|
assert matches({"$nin": [1, "2.5"]}) == ["flag", "scalar", "space"]
|
|
|
|
|
|
def test_euclidean_distance(vectors, index):
|
|
vectors.put_vectors(**index, vectors=[{"key": "p", "data": {"float32": [3, 4, 0, 0]}}])
|
|
result = vectors.query_vectors(**index, queryVector={"float32": [0, 0, 0, 0]}, topK=1, returnDistance=True)
|
|
assert result["distanceMetric"] == "euclidean"
|
|
assert result["vectors"][0]["distance"] == pytest.approx(5.0)
|
|
|
|
|
|
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_vector_policy_permissions(vectors, vector_clients, index):
|
|
reader, writer = vector_clients["reader"], vector_clients["writer"]
|
|
arn = vectors.get_index(**index)["index"]["indexArn"]
|
|
row = {"key": "policy-row", "data": {"float32": [1, 0, 0, 0]}, "metadata": {"owner": "writer"}}
|
|
expect_error("AccessDeniedException", writer.put_vectors, **index, vectors=[row])
|
|
expect_error("AccessDeniedException", reader.get_vectors, **index, keys=[row["key"]])
|
|
statements = [
|
|
{"Effect": "Allow", "Principal": {"AWS": "writer"},
|
|
"Action": "s3vectors:PutVectors", "Resource": arn},
|
|
{"Effect": "Allow", "Principal": {"AWS": "reader"},
|
|
"Action": ["s3vectors:GetVectors", "s3vectors:QueryVectors"], "Resource": arn},
|
|
]
|
|
|
|
def put_policy():
|
|
vectors.put_vector_bucket_policy(vectorBucketName=index["vectorBucketName"],
|
|
policy=json.dumps({"Version": "2012-10-17", "Statement": statements}))
|
|
|
|
put_policy()
|
|
writer.put_vectors(**index, vectors=[row])
|
|
assert reader.get_vectors(**index, keys=[row["key"]], returnData=True, returnMetadata=True)["vectors"] == [row]
|
|
assert reader.query_vectors(**index, queryVector=row["data"], topK=1)["vectors"][0]["key"] == row["key"]
|
|
expect_error("AccessDeniedException", reader.put_vectors, **index, vectors=[row])
|
|
expect_error("AccessDeniedException", writer.get_vectors, **index, keys=[row["key"]])
|
|
expect_error("AccessDeniedException", reader.delete_index, **index)
|
|
expect_error("AccessDeniedException", reader.list_vectors, **index)
|
|
other = dict(index, indexName="other-index")
|
|
vectors.create_index(**other, dataType="float32", dimension=4, distanceMetric="euclidean")
|
|
try:
|
|
expect_error("AccessDeniedException", reader.get_vectors, **other, keys=[row["key"]])
|
|
expect_error("AccessDeniedException", writer.put_vectors, **other, vectors=[row])
|
|
statements.append({"Effect": "Deny", "Principal": {"AWS": "writer"},
|
|
"Action": "s3vectors:PutVectors", "Resource": arn})
|
|
put_policy()
|
|
expect_error("AccessDeniedException", writer.put_vectors, **index, vectors=[row])
|
|
assert reader.get_vectors(**index, keys=[row["key"]])["vectors"][0]["key"] == row["key"]
|
|
vectors.delete_vector_bucket_policy(vectorBucketName=index["vectorBucketName"])
|
|
expect_error("AccessDeniedException", reader.get_vectors, **index, keys=[row["key"]])
|
|
expect_error("AccessDeniedException", writer.put_vectors, **index, vectors=[row])
|
|
vectors.put_vectors(**index, vectors=[row])
|
|
finally:
|
|
vectors.delete_index(**other)
|
|
|
|
|
|
def index_object_encryption(vectors, bucket):
|
|
s3 = boto3.client("s3", endpoint_url=vectors.meta.endpoint_url, aws_access_key_id="test",
|
|
aws_secret_access_key="test", region_name="us-east-1")
|
|
keys = [o["Key"] for page in s3.get_paginator("list_objects_v2").paginate(Bucket=bucket, Prefix=f"vector/{INDEX}/")
|
|
for o in page.get("Contents", []) if o["Size"] > 0]
|
|
assert any("/data/" in key for key in keys) and any("/_versions/" in key for key in keys)
|
|
return {key: s3.head_object(Bucket=bucket, Key=key).get("ServerSideEncryption") for key in keys}
|
|
|
|
|
|
def test_vector_bucket_encryption(vectors):
|
|
bucket = "encryption-" + uuid.uuid4().hex[:16]
|
|
plain = bucket + "-plain"
|
|
config = {"sseType": "AES256"}
|
|
vectors.create_vector_bucket(vectorBucketName=bucket, encryptionConfiguration=config)
|
|
vectors.create_vector_bucket(vectorBucketName=plain)
|
|
try:
|
|
assert vectors.get_vector_bucket(vectorBucketName=bucket)["vectorBucket"]["encryptionConfiguration"] == config
|
|
expect_error("ValidationException", vectors.create_vector_bucket, vectorBucketName=bucket + "-kms",
|
|
encryptionConfiguration={"sseType": "aws:kms", "kmsKeyArn": "arn:aws:kms:us-east-1:123456789012:key/test"})
|
|
assert vectors.list_vector_buckets(prefix=bucket + "-kms")["vectorBuckets"] == []
|
|
|
|
for name in (bucket, plain):
|
|
vectors.create_index(vectorBucketName=name, indexName=INDEX, dataType="float32", dimension=4,
|
|
distanceMetric="euclidean")
|
|
vectors.put_vectors(vectorBucketName=name, indexName=INDEX, vectors=[
|
|
{"key": "a", "data": {"float32": [1, 0, 0, 0]}, "metadata": {"color": "red"}},
|
|
{"key": "b", "data": {"float32": [0, 1, 0, 0]}, "metadata": {"color": "blue"}},
|
|
])
|
|
hits = vectors.query_vectors(vectorBucketName=name, indexName=INDEX, queryVector={"float32": [1, 0, 0, 0]},
|
|
topK=1, returnMetadata=True)["vectors"]
|
|
assert [(h["key"], h["metadata"]) for h in hits] == [("a", {"color": "red"})]
|
|
assert set(index_object_encryption(vectors, bucket).values()) == {"AES256"}
|
|
assert set(index_object_encryption(vectors, plain).values()) == {None}
|
|
finally:
|
|
for name in (bucket, plain):
|
|
with contextlib.suppress(ClientError):
|
|
vectors.delete_index(vectorBucketName=name, indexName=INDEX)
|
|
vectors.delete_vector_bucket(vectorBucketName=name)
|
|
|
|
|
|
def test_catalog_pagination(vectors, index):
|
|
prefix = "pages-" + uuid.uuid4().hex[:16]
|
|
buckets, indexes = [], []
|
|
try:
|
|
for i in range(3):
|
|
bucket = f"{prefix}-{i}"
|
|
vectors.create_vector_bucket(vectorBucketName=bucket)
|
|
buckets.append(bucket)
|
|
for name in ("idx-a", "idx-b"):
|
|
args = dict(index, indexName=name)
|
|
vectors.create_index(**args, dataType="float32", dimension=4, distanceMetric="euclidean")
|
|
indexes.append(args)
|
|
pages = list(vectors.get_paginator("list_vector_buckets").paginate(
|
|
prefix=prefix, PaginationConfig={"PageSize": 1}))
|
|
assert all(len(page["vectorBuckets"]) <= 1 for page in pages)
|
|
assert sorted(b["vectorBucketName"] for page in pages for b in page["vectorBuckets"]) == buckets
|
|
pages = list(vectors.get_paginator("list_indexes").paginate(
|
|
vectorBucketName=index["vectorBucketName"], prefix="idx", PaginationConfig={"PageSize": 1}))
|
|
assert all(len(page["indexes"]) <= 1 for page in pages)
|
|
assert sorted(i["indexName"] for page in pages for i in page["indexes"]) == [INDEX, "idx-a", "idx-b"]
|
|
expect_error("ValidationException", vectors.list_vector_buckets, maxResults=501)
|
|
expect_error("ValidationException", vectors.list_indexes, vectorBucketName=index["vectorBucketName"], maxResults=501)
|
|
finally:
|
|
for args in indexes:
|
|
vectors.delete_index(**args)
|
|
for bucket in buckets:
|
|
vectors.delete_vector_bucket(vectorBucketName=bucket)
|
|
|
|
|
|
def test_delete_and_recreate_index(vectors, index):
|
|
vectors.put_vectors(**index, vectors=[{"key": "old", "data": {"float32": [1, 0, 0, 0]}}])
|
|
vectors.delete_index(**index)
|
|
expect_error("NotFoundException", vectors.get_index, **index)
|
|
vectors.create_index(**index, dataType="float32", dimension=2, distanceMetric="cosine")
|
|
assert vectors.list_vectors(**index)["vectors"] == []
|
|
row = {"key": "new", "data": {"float32": [1, 0]}}
|
|
vectors.put_vectors(**index, vectors=[row])
|
|
assert vectors.get_vectors(**index, keys=["old", "new"], returnData=True)["vectors"] == [row]
|
|
result = vectors.query_vectors(**index, queryVector=row["data"], topK=1, returnDistance=True)
|
|
assert result["distanceMetric"] == "cosine"
|
|
assert result["vectors"][0]["key"] == "new"
|
|
assert result["vectors"][0]["distance"] == pytest.approx(0, abs=1e-6)
|
|
|
|
|
|
def test_index_arn_validation(vectors, index):
|
|
arn = vectors.get_index(**index)["index"]["indexArn"]
|
|
foreign = []
|
|
for field, value in ((1, "aws-cn"), (3, "eu-west-1"), (4, "999999999999")):
|
|
parts = arn.split(":", 5)
|
|
assert parts[field] != value
|
|
parts[field] = value
|
|
foreign.append(":".join(parts))
|
|
for invalid in (
|
|
*foreign,
|
|
arn.replace(":s3vectors:", ":s3tables:"),
|
|
arn.replace(":s3vectors:", ":s3:"),
|
|
arn + "/index/extra",
|
|
arn + "/",
|
|
):
|
|
expect_error("ValidationException", vectors.get_index, indexArn=invalid)
|
|
expect_error("ValidationException", vectors.delete_index, indexArn=invalid)
|
|
expect_error("ValidationException", vectors.get_vector_bucket, vectorBucketArn=arn)
|
|
expect_error("ValidationException", vectors.list_indexes, vectorBucketArn=arn)
|
|
bucket_arn = arn.rsplit("/index/", 1)[0]
|
|
for request in (
|
|
{"vectorBucketName": "other-" + index["vectorBucketName"], "indexArn": arn},
|
|
{"indexName": "other-" + index["indexName"], "indexArn": arn},
|
|
):
|
|
expect_error("ValidationException", vectors.get_index, **request)
|
|
expect_error("ValidationException", vectors.delete_index, **request)
|
|
for operation in (vectors.get_vector_bucket, vectors.list_indexes, vectors.delete_vector_bucket):
|
|
expect_error("ValidationException", operation, vectorBucketName="other-" + index["vectorBucketName"],
|
|
vectorBucketArn=bucket_arn)
|
|
assert vectors.get_index(**index, indexArn=arn)["index"]["indexArn"] == arn
|
|
assert vectors.get_vector_bucket(vectorBucketName=index["vectorBucketName"],
|
|
vectorBucketArn=bucket_arn)["vectorBucket"]["vectorBucketArn"] == bucket_arn
|
|
assert vectors.get_index(**index)["index"]["indexArn"] == arn
|
|
|
|
|
|
def test_resource_tags(vectors, vector_clients, index):
|
|
reader = vector_clients["reader"]
|
|
bucket_arn = vectors.get_vector_bucket(vectorBucketName=index["vectorBucketName"])["vectorBucket"]["vectorBucketArn"]
|
|
index_arn = vectors.get_index(**index)["index"]["indexArn"]
|
|
assert vectors.list_tags_for_resource(resourceArn=bucket_arn)["tags"] == {}
|
|
vectors.tag_resource(resourceArn=bucket_arn, tags={"team": "search", "env": "dev"})
|
|
vectors.tag_resource(resourceArn=index_arn, tags={"team": "ranking"})
|
|
assert vectors.list_tags_for_resource(resourceArn=bucket_arn)["tags"] == {"team": "search", "env": "dev"}
|
|
assert vectors.list_tags_for_resource(resourceArn=index_arn)["tags"] == {"team": "ranking"}
|
|
vectors.untag_resource(resourceArn=bucket_arn, tagKeys=["env", "absent"])
|
|
assert vectors.list_tags_for_resource(resourceArn=bucket_arn)["tags"] == {"team": "search"}
|
|
expect_error("ValidationException", vectors.tag_resource, resourceArn=bucket_arn,
|
|
tags={f"k{i}": "v" for i in range(50)})
|
|
expect_error("NotFoundException", vectors.list_tags_for_resource, resourceArn=bucket_arn + "/index/missing")
|
|
expect_error("AccessDeniedException", reader.list_tags_for_resource, resourceArn=bucket_arn)
|
|
expect_error("AccessDeniedException", reader.tag_resource, resourceArn=index_arn, tags={"team": "other"})
|
|
assert vectors.list_tags_for_resource(resourceArn=index_arn)["tags"] == {"team": "ranking"}
|
|
|
|
|
|
def test_create_with_tags(vectors):
|
|
bucket = "tagged-" + uuid.uuid4().hex[:16]
|
|
bucket_arn = vectors.create_vector_bucket(vectorBucketName=bucket, tags={"team": "search"})["vectorBucketArn"]
|
|
try:
|
|
assert vectors.list_tags_for_resource(resourceArn=bucket_arn)["tags"] == {"team": "search"}
|
|
index_arn = vectors.create_index(vectorBucketName=bucket, indexName=INDEX, dataType="float32", dimension=4,
|
|
distanceMetric="cosine", tags={"tier": "hot"})["indexArn"]
|
|
try:
|
|
assert vectors.list_tags_for_resource(resourceArn=index_arn)["tags"] == {"tier": "hot"}
|
|
finally:
|
|
vectors.delete_index(vectorBucketName=bucket, indexName=INDEX)
|
|
finally:
|
|
vectors.delete_vector_bucket(vectorBucketName=bucket)
|
|
|
|
|
|
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"]
|
|
|
|
|
|
def test_engine_without_credentials(open_cluster):
|
|
"""With no S3 identities, an engine given no credentials still serves an index."""
|
|
client = boto3.client("s3vectors", endpoint_url=open_cluster["endpoint"], aws_access_key_id="any",
|
|
aws_secret_access_key="any", region_name="us-east-1",
|
|
config=Config(connect_timeout=3, read_timeout=20, retries={"max_attempts": 0}))
|
|
index = {"vectorBucketName": BUCKET, "indexName": INDEX}
|
|
client.create_vector_bucket(vectorBucketName=BUCKET)
|
|
client.create_index(**index, dataType="float32", dimension=2, distanceMetric="cosine")
|
|
client.put_vectors(**index, vectors=[{"key": "a", "data": {"float32": [1, 0]}},
|
|
{"key": "b", "data": {"float32": [0, 1]}}])
|
|
hits = client.query_vectors(**index, queryVector={"float32": [1, 0]}, topK=1)["vectors"]
|
|
assert [h["key"] for h in hits] == ["a"]
|
|
|