seaweedfs/test/s3vectors/s3vectors_test.py

"""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_tag_conditions(vectors, vector_clients, index):
    reader, writer = vector_clients["reader"], vector_clients["writer"]
    bucket = index["vectorBucketName"]
    bucket_arn = vectors.get_vector_bucket(vectorBucketName=bucket)["vectorBucket"]["vectorBucketArn"]
    index_arn = vectors.get_index(**index)["index"]["indexArn"]
    other = dict(index, indexName="other-index")
    vectors.create_index(**other, dataType="float32", dimension=4, distanceMetric="euclidean")
    vectors.tag_resource(resourceArn=bucket_arn, tags={"env": "prod"})
    vectors.tag_resource(resourceArn=index_arn, tags={"stage": "live"})
    indexes = bucket_arn + "/index/*"
    statements = [
        {"Effect": "Allow", "Principal": {"AWS": "reader"}, "Action": "s3vectors:GetIndex", "Resource": indexes,
         "Condition": {"StringEquals": {"aws:ResourceTag/stage": "live"}}},
        {"Effect": "Allow", "Principal": {"AWS": "reader"}, "Action": "s3vectors:QueryVectors", "Resource": indexes,
         "Condition": {"StringEquals": {"s3vectors:VectorBucketTag/env": "prod"}}},
        {"Effect": "Allow", "Principal": {"AWS": "writer"}, "Action": "s3vectors:TagResource", "Resource": indexes,
         "Condition": {"StringEquals": {"aws:RequestTag/team": "search"},
                       "ForAllValues:StringEquals": {"aws:TagKeys": ["team"]}}},
    ]
    try:
        vectors.put_vector_bucket_policy(vectorBucketName=bucket,
                                         policy=json.dumps({"Version": "2012-10-17", "Statement": statements}))
        assert reader.get_index(**index)["index"]["indexArn"] == index_arn
        expect_error("AccessDeniedException", reader.get_index, **other)
        assert reader.query_vectors(**other, queryVector={"float32": [1, 0, 0, 0]}, topK=1)["vectors"] == []
        writer.tag_resource(resourceArn=index_arn, tags={"team": "search"})
        expect_error("AccessDeniedException", writer.tag_resource, resourceArn=index_arn, tags={"team": "other"})
        expect_error("AccessDeniedException", writer.tag_resource, resourceArn=index_arn,
                     tags={"team": "search", "env": "dev"})
        assert vectors.list_tags_for_resource(resourceArn=index_arn)["tags"] == {"stage": "live", "team": "search"}
        vectors.tag_resource(resourceArn=bucket_arn, tags={"env": "dev"})
        expect_error("AccessDeniedException", reader.query_vectors, **other,
                     queryVector={"float32": [1, 0, 0, 0]}, topK=1)
        unsupported = dict(statements[0], Condition={"IpAddress": {"aws:SourceIp": "10.0.0.0/8"}})
        expect_error("ValidationException", vectors.put_vector_bucket_policy, vectorBucketName=bucket,
                     policy=json.dumps({"Version": "2012-10-17", "Statement": [unsupported]}))
    finally:
        vectors.delete_vector_bucket_policy(vectorBucketName=bucket)
        vectors.delete_index(**other)


def exact_neighbours(rows, query, k):
    import math
    qn = math.sqrt(sum(x * x for x in query))
    scored = sorted(((sum(a * b for a, b in zip(v, query)) / (math.sqrt(sum(a * a for a in v)) * qn), key)
                     for key, v in rows.items()), reverse=True)
    return {key for _, key in scored[:k]}


@pytest.mark.worker
def test_vector_index_job(vectors, cluster):
    """The maintenance worker indexes a table past the row threshold through the
    admin server, queries then run through the index, and rows written after
    the build are still found."""
    import os
    import random
    import subprocess
    import urllib.request
    from pathlib import Path

    bucket, index, dim = "vectors-indexed", "idx", 16
    vectors.create_vector_bucket(vectorBucketName=bucket)
    vectors.create_index(vectorBucketName=bucket, indexName=index, dataType="float32",
                         dimension=dim, distanceMetric="cosine")
    # Embedding-like data: points around shared centres. Independent Gaussian
    # vectors have no locality for an IVF index to use, so recall through one
    # would only track the fraction of partitions probed.
    rng = random.Random(7)
    centres = [[rng.gauss(0, 1) for _ in range(dim)] for _ in range(30)]

    def point():
        centre = rng.choice(centres)
        return [c + 0.3 * rng.gauss(0, 1) for c in centre]

    rows = {f"v{i}": point() for i in range(1200)}
    keys = list(rows)
    for offset in range(0, len(keys), 400):
        vectors.put_vectors(vectorBucketName=bucket, indexName=index,
                            vectors=[{"key": key, "data": {"float32": rows[key]}} for key in keys[offset:offset + 400]])

    admin = cluster["admin"]
    job = "lance_vector_index"

    def call(path, method="GET", body=None):
        data = json.dumps(body).encode() if body is not None else None
        request = urllib.request.Request(f"{admin}/api/plugin/job-types/{job}{path}", data=data, method=method,
                                         headers={"Content-Type": "application/json"} if data else {})
        with urllib.request.urlopen(request, timeout=120) as response:
            return json.load(response)

    call("/config", "PUT", {"admin_runtime": {"enabled": True},
                            "worker_config_values": {"min_rows": {"int64_value": "1000"}}})
    run = call("/run", "POST", {"timeout_seconds": 110})
    assert run["success_count"] == 1 and run["error_count"] == 0, run

    cli = Path(os.environ.get("VECTOR_CLI_BINARY", Path(__file__).resolve().parent / "weed-vector-cli"))
    coverage = json.loads(subprocess.run(
        [str(cli), "--namespace", cluster["catalog"], "--access-key", "test", "--secret-key", "test",
         "coverage", f"{bucket}$vector${index}"], check=True, capture_output=True, text=True).stdout)
    assert len(coverage) == 1 and coverage[0]["segments"] == 1 and coverage[0]["unindexed_rows"] == 0, coverage

    found = 0
    for _ in range(20):
        query = point()
        hits = vectors.query_vectors(vectorBucketName=bucket, indexName=index,
                                     queryVector={"float32": query}, topK=10)["vectors"]
        assert len(hits) == 10
        found += len({h["key"] for h in hits} & exact_neighbours(rows, query, 10))
    assert found >= 180, f"recall {found}/200"

    late = point()
    vectors.put_vectors(vectorBucketName=bucket, indexName=index, vectors=[{"key": "late", "data": {"float32": late}}])
    hits = vectors.query_vectors(vectorBucketName=bucket, indexName=index, queryVector={"float32": late}, topK=1)
    assert [h["key"] for h in hits["vectors"]] == ["late"]

    again = call("/run", "POST", {"timeout_seconds": 60})
    assert again["detected_count"] == 0, again

    def coverage_now():
        return json.loads(subprocess.run(
            [str(cli), "--namespace", cluster["catalog"], "--access-key", "test", "--secret-key", "test",
             "coverage", f"{bucket}$vector${index}"], check=True, capture_output=True, text=True).stdout)[0]

    def optimize(expect_detected):
        request = urllib.request.Request(
            f"{admin}/api/plugin/job-types/lance_optimize_indices/run", method="POST",
            data=json.dumps({"timeout_seconds": 110}).encode(), headers={"Content-Type": "application/json"})
        with urllib.request.urlopen(request, timeout=120) as response:
            result = json.load(response)
        assert result["detected_count"] == expect_detected and result["error_count"] == 0, result

    # A tenth of the index written after the build brings the optimize job
    # in, even though it is far below its absolute budget.
    optimize(0)
    more = {f"m{i}": point() for i in range(200)}
    vectors.put_vectors(vectorBucketName=bucket, indexName=index,
                        vectors=[{"key": key, "data": {"float32": v}} for key, v in more.items()])
    before = coverage_now()
    assert before["unindexed_rows"] == 201, before
    optimize(1)
    after = coverage_now()
    assert after["unindexed_rows"] == 0 and after["segments"] == 2, after

    # Five times the rows the index was sized for: the job trains it again
    # instead of extending it, and the partitions follow the new row count.
    rows.update(more)
    for offset in range(0, 4600, 400):
        batch = {f"g{offset + i}": point() for i in range(400)}
        rows.update(batch)
        vectors.put_vectors(vectorBucketName=bucket, indexName=index,
                            vectors=[{"key": key, "data": {"float32": v}} for key, v in batch.items()])
    optimize(1)
    retrained = coverage_now()
    assert retrained["unindexed_rows"] == 0 and retrained["segments"] == 1, retrained
    assert retrained["partitions"] > after["partitions"], (before, after, retrained)
    found = 0
    for _ in range(20):
        query = point()
        hits = vectors.query_vectors(vectorBucketName=bucket, indexName=index,
                                     queryVector={"float32": query}, topK=10)["vectors"]
        found += len({h["key"] for h in hits} & exact_neighbours(rows, query, 10))
    assert found >= 180, f"recall after retrain {found}/200"


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"]