import os
import functools
import shutil
import tempfile
import pytest
import sqlite3
from typing import Generator, List, Callable, Dict, Union
from chromadb.db.impl.grpc.client import GrpcSysDB
from chromadb.db.impl.grpc.server import GrpcMockSysDB
from chromadb.errors import NotFoundError, UniqueConstraintError, InternalError
from chromadb.test.conftest import find_free_port
from chromadb.types import Collection, Segment, SegmentScope
from chromadb.db.impl.sqlite import SqliteDB
from chromadb.config import (
DEFAULT_DATABASE,
DEFAULT_TENANT,
System,
Settings,
)
from chromadb.db.system import SysDB
from pytest import FixtureRequest
import uuid
from chromadb.api.configuration import CollectionConfigurationInternal
import logging
logger = logging.getLogger(__name__)
TENANT = "default"
NAMESPACE = "default"
# These are the sample collections that are used in the tests below. Tests can override
# the fields as needed.
sample_collections: List[Collection] = [
Collection(
id=uuid.UUID(int=1),
name="test_collection_1",
configuration=CollectionConfigurationInternal(),
metadata={"test_str": "str1", "test_int": 1, "test_float": 1.3},
dimension=128,
database=DEFAULT_DATABASE,
tenant=DEFAULT_TENANT,
version=0,
),
Collection(
id=uuid.UUID(int=2),
name="test_collection_2",
configuration=CollectionConfigurationInternal(),
metadata={"test_str": "str2", "test_int": 2, "test_float": 2.3},
dimension=None,
database=DEFAULT_DATABASE,
tenant=DEFAULT_TENANT,
version=0,
),
Collection(
id=uuid.UUID(int=3),
name="test_collection_3",
configuration=CollectionConfigurationInternal(),
metadata={"test_str": "str3", "test_int": 3, "test_float": 3.3},
dimension=None,
database=DEFAULT_DATABASE,
tenant=DEFAULT_TENANT,
version=0,
),
Collection(
id=uuid.UUID(int=4),
name="test_collection_4",
configuration=CollectionConfigurationInternal(),
metadata={"test_str": "str4", "test_int": 4, "test_float": 4.4},
dimension=None,
database=DEFAULT_DATABASE,
tenant=DEFAULT_TENANT,
version=0,
),
]
def sqlite() -> Generator[SysDB, None, None]:
"""Fixture generator for sqlite DB"""
db = SqliteDB(
System(
Settings(
allow_reset=True,
)
)
)
db.start()
yield db
db.stop()
def sqlite_persistent() -> Generator[SysDB, None, None]:
"""Fixture generator for sqlite DB"""
save_path = tempfile.mkdtemp()
db = SqliteDB(
System(
Settings(
allow_reset=True,
is_persistent=True,
persist_directory=save_path,
)
)
)
db.start()
yield db
db.stop()
if os.path.exists(save_path):
shutil.rmtree(save_path)
def grpc_with_mock_server() -> Generator[SysDB, None, None]:
"""Fixture generator for sqlite DB that creates a mock grpc sysdb server
and a grpc client that connects to it."""
port = find_free_port()
system = System(
Settings(
allow_reset=True,
chroma_server_grpc_port=port,
)
)
system.instance(GrpcMockSysDB)
client = system.instance(GrpcSysDB)
system.start()
client.reset_and_wait_for_ready()
yield client
system.stop()
def grpc_with_real_server() -> Generator[SysDB, None, None]:
logger.debug("Setting up grpc_with_real_server")
system = System(
Settings(
allow_reset=True,
chroma_server_grpc_port=50051,
)
)
client = system.instance(GrpcSysDB)
logger.debug("Starting system")
system.start()
logger.debug("Resetting client and waiting for ready")
client.reset_and_wait_for_ready()
logger.debug("grpc_with_real_server setup complete")
yield client
logger.debug("Stopping system in grpc_with_real_server")
system.stop()
def db_fixtures() -> List[Callable[[], Generator[SysDB, None, None]]]:
if "CHROMA_CLUSTER_TEST_ONLY" in os.environ:
return [grpc_with_real_server]
else:
return [sqlite, sqlite_persistent, grpc_with_mock_server]
@pytest.fixture(scope="module", params=db_fixtures())
def sysdb(request: FixtureRequest) -> Generator[SysDB, None, None]:
logger.debug(f"Setting up sysdb fixture with {request.param.__name__}")
yield next(request.param())
logger.debug("Tearing down sysdb fixture")
def sample_segment(collection_id: uuid.UUID = uuid.uuid4(),
segment_type: str = "test_type_a",
scope: SegmentScope = SegmentScope.VECTOR,
metadata: Dict[str, Union[str, int, float]] = {
"test_str": "str1",
"test_int": 1,
"test_float": 1.3,
},
) -> Segment:
return Segment(
id=uuid.uuid4(),
type=segment_type,
scope=scope,
collection=collection_id,
metadata=metadata,
file_paths={},
)
# region Collection tests
def test_create_get_delete_collections(sysdb: SysDB) -> None:
logger.debug("Resetting state")
sysdb.reset_state()
segments_created_with_collection = []
for collection in sample_collections:
logger.debug(f"Creating collection: {collection.name}")
segments = [
sample_segment(collection_id=collection.id, scope=SegmentScope.METADATA),
sample_segment(collection_id=collection.id, scope=SegmentScope.RECORD),
sample_segment(collection_id=collection.id, scope=SegmentScope.VECTOR),
]
segments_created_with_collection.extend(segments)
sysdb.create_collection(
id=collection.id,
name=collection.name,
configuration=collection.get_configuration(),
segments=segments,
metadata=collection["metadata"],
dimension=collection["dimension"],
)
collection["database"] = DEFAULT_DATABASE
collection["tenant"] = DEFAULT_TENANT
logger.debug("Getting all collections")
results = sysdb.get_collections()
results = sorted(results, key=lambda c: c.name)
assert sorted(results, key=lambda c: c.name) == sample_collections
# Duplicate create fails
with pytest.raises(UniqueConstraintError):
sysdb.create_collection(
name=sample_collections[0].name,
id=sample_collections[0].id,
configuration=sample_collections[0].get_configuration(),
segments=[segments_created_with_collection[0]],
)
# Find by name
for collection in sample_collections:
result = sysdb.get_collections(name=collection["name"])
assert result == [collection]
# Find by id
for collection in sample_collections:
result = sysdb.get_collections(id=collection["id"])
assert result == [collection]
# Verify segment information
for collection in sample_collections:
collection_with_segments_result = sysdb.get_collection_with_segments(collection.id)
assert collection_with_segments_result["collection"] == collection
assert all([segment["collection"] == collection.id for segment in collection_with_segments_result["segments"]])
# Delete
c1 = sample_collections[0]
sysdb.delete_collection(id=c1.id)
results = sysdb.get_collections()
assert c1 not in results
assert len(results) == len(sample_collections) - 1
assert sorted(results, key=lambda c: c.name) == sample_collections[1:]
by_id_result = sysdb.get_collections(id=c1["id"])
assert by_id_result == []
# Check that the segment was deleted
by_collection_result = sysdb.get_segments(collection=c1.id)
assert by_collection_result == []
# Duplicate delete throws an exception
with pytest.raises(NotFoundError):
sysdb.delete_collection(c1.id)
# Create a new collection with two segments that have same id.
# Creation should fail.
# Check that collection was not created.
# Check that segments were not created.
with pytest.raises((InternalError, UniqueConstraintError, sqlite3.IntegrityError)):
sysdb.create_collection(
name=sample_collections[0].name,
id=sample_collections[0].id,
configuration=sample_collections[0].get_configuration(),
segments=[segments_created_with_collection[0], segments_created_with_collection[0]],
)
# Check that collection was not created.
# Check that segments were not created.
by_id_result = sysdb.get_collections(id=sample_collections[0].id)
assert by_id_result == []
by_collection_result = sysdb.get_segments(collection=sample_collections[0].id)
assert by_collection_result == []
def test_update_collections(sysdb: SysDB) -> None:
coll = Collection(
name=sample_collections[0].name,
id=sample_collections[0].id,
configuration=sample_collections[0].get_configuration(),
metadata=sample_collections[0]["metadata"],
dimension=sample_collections[0]["dimension"],
database=DEFAULT_DATABASE,
tenant=DEFAULT_TENANT,
version=0,
)
sysdb.reset_state()
sysdb.create_collection(
id=coll.id,
name=coll.name,
configuration=coll.get_configuration(),
segments=[
Segment(
id=uuid.uuid4(),
type="test_type_a",
scope=SegmentScope.VECTOR,
collection=coll.id,
metadata={"test_str": "str1", "test_int": 1, "test_float": 1.3},
file_paths={},
)
],
metadata=coll["metadata"],
dimension=coll["dimension"],
)
# Update name
coll["name"] = "new_name"
sysdb.update_collection(coll.id, name=coll.name)
result = sysdb.get_collections(name=coll.name)
assert result == [coll]
# Update dimension
coll["dimension"] = 128
sysdb.update_collection(coll.id, dimension=coll.dimension)
result = sysdb.get_collections(id=coll["id"])
assert result == [coll]
# Reset the metadata
coll["metadata"] = {"test_str2": "str2"}
sysdb.update_collection(coll.id, metadata=coll["metadata"])
result = sysdb.get_collections(id=coll["id"])
assert result == [coll]
# Delete all metadata keys
coll["metadata"] = None
sysdb.update_collection(coll.id, metadata=None)
result = sysdb.get_collections(id=coll["id"])
assert result == [coll]
def test_get_or_create_collection(sysdb: SysDB) -> None:
sysdb.reset_state()
# get_or_create = True returns existing collection
collection = sample_collections[0]
sysdb.create_collection(
id=collection.id,
name=collection.name,
configuration=collection.get_configuration(),
segments=[
Segment(
id=uuid.uuid4(),
type="test_type_a",
scope=SegmentScope.VECTOR,
collection=collection.id,
metadata={"test_str": "str1", "test_int": 1, "test_float": 1.3},
file_paths={},
)
],
metadata=collection["metadata"],
dimension=collection["dimension"],
)
# Create collection with same name, but different id.
# Since get_or_create is true, it should return the existing collection.
result, created = sysdb.create_collection(
name=collection.name,
id=uuid.uuid4(),
configuration=collection.get_configuration(),
get_or_create=True,
segments=[
Segment(
id=uuid.uuid4(),
type="test_type_a",
scope=SegmentScope.VECTOR,
collection=sample_collections[1].id,
metadata={"test_str": "str1", "test_int": 1, "test_float": 1.3},
file_paths={},
)
], # This could have been empty - [].
metadata=collection["metadata"],
)
assert result == collection
# Only one collection with the same name exists
get_result = sysdb.get_collections(name=collection["name"])
assert get_result == [collection]
# get_or_create = True creates new collection
result, created = sysdb.create_collection(
name=sample_collections[1].name,
id=sample_collections[1].id,
configuration=sample_collections[1].get_configuration(),
segments=[
Segment(
id=uuid.uuid4(),
type="test_type_a",
scope=SegmentScope.VECTOR,
collection=sample_collections[1].id,
metadata={"test_str": "str1", "test_int": 1, "test_float": 1.3},
file_paths={},
)
],
get_or_create=True,
metadata=sample_collections[1]["metadata"],
)
assert result == sample_collections[1]
# get_or_create = False creates new collection
result, created = sysdb.create_collection(
name=sample_collections[2].name,
id=sample_collections[2].id,
configuration=sample_collections[2].get_configuration(),
segments=[
Segment(
id=uuid.uuid4(),
type="test_type_a",
scope=SegmentScope.VECTOR,
collection=sample_collections[2].id,
metadata={"test_str": "str1", "test_int": 1, "test_float": 1.3},
file_paths={},
)
],
get_or_create=False,
metadata=sample_collections[2]["metadata"],
)
assert result == sample_collections[2]
# get_or_create = False fails if collection already exists
with pytest.raises(UniqueConstraintError):
sysdb.create_collection(
name=sample_collections[2].name,
id=sample_collections[2].id,
configuration=sample_collections[2].get_configuration(),
get_or_create=False,
segments=[
Segment(
id=uuid.uuid4(),
type="test_type_a",
scope=SegmentScope.VECTOR,
collection=sample_collections[2].id,
metadata={"test_str": "str1", "test_int": 1, "test_float": 1.3},
file_paths={},
)
],
metadata=collection["metadata"],
)
# get_or_create = True does not overwrite metadata
overlayed_metadata: Dict[str, Union[str, int, float]] = {
"test_new_str": "new_str",
"test_int": 1,
}
result, created = sysdb.create_collection(
name=sample_collections[2].name,
id=sample_collections[2].id,
configuration=sample_collections[2].get_configuration(),
segments=[
Segment(
id=uuid.uuid4(),
type="test_type_a",
scope=SegmentScope.VECTOR,
collection=sample_collections[2].id,
metadata={"test_str": "str1", "test_int": 1, "test_float": 1.3},
file_paths={},
)
],
get_or_create=True,
metadata=overlayed_metadata,
)
assert result["metadata"] != overlayed_metadata
assert result["metadata"] == sample_collections[2]["metadata"]
# get_or_create = True with None metadata does not overwrite metadata
result, created = sysdb.create_collection(
name=sample_collections[2].name,
id=sample_collections[2].id,
configuration=sample_collections[2].get_configuration(),
segments=[sample_segment(sample_collections[2].id)],
get_or_create=True,
metadata=None,
)
assert result["metadata"] == sample_collections[2]["metadata"]
def test_create_get_delete_database_and_collection(sysdb: SysDB) -> None:
sysdb.reset_state()
# Check that collection doesn't exist in new_database
result = sysdb.get_collections(id=sample_collections[1].id, database="new_database")
assert len(result) == 0
# Create a new database
sysdb.create_database(id=uuid.uuid4(), name="new_database")
# Create a new collection in the new database
sysdb.create_collection(
id=sample_collections[0].id,
name=sample_collections[0].name,
configuration=sample_collections[0].get_configuration(),
segments=[sample_segment(sample_collections[0].id)],
metadata=sample_collections[0]["metadata"],
dimension=sample_collections[0]["dimension"],
database="new_database",
)
# Create a new collection with the same id but different name in the new database
# and expect an error
with pytest.raises(UniqueConstraintError):
sysdb.create_collection(
id=sample_collections[0].id,
name="new_name",
configuration=sample_collections[0].get_configuration(),
metadata=sample_collections[0]["metadata"],
dimension=sample_collections[0]["dimension"],
segments=[sample_segment(sample_collections[0].id)],
database="new_database",
get_or_create=False,
)
# Create a new collection in the default database
sysdb.create_collection(
id=sample_collections[1].id,
name=sample_collections[1].name,
configuration=sample_collections[1].get_configuration(),
metadata=sample_collections[1]["metadata"],
dimension=sample_collections[1]["dimension"],
segments=[sample_segment(sample_collections[1].id)],
)
# Check that the new database and collections exist
result = sysdb.get_collections(
name=sample_collections[0]["name"], database="new_database"
)
assert len(result) == 1
sample_collections[0]["database"] = "new_database"
assert result[0] == sample_collections[0]
# Check that the collection in the default database exists
result = sysdb.get_collections(name=sample_collections[1]["name"])
assert len(result) == 1
assert result[0] == sample_collections[1]
# Get for a database that doesn't exist with a name that exists in the new database and expect no results
assert (
len(
sysdb.get_collections(
name=sample_collections[0]["name"], database="fake_db"
)
)
== 0
)
# Delete the collection in the new database
sysdb.delete_collection(id=sample_collections[0].id, database="new_database")
# Check that the collection in the new database was deleted
result = sysdb.get_collections(database="new_database")
assert len(result) == 0
# Check that the collection in the default database still exists
result = sysdb.get_collections(name=sample_collections[1].name)
assert len(result) == 1
assert result[0] == sample_collections[1]
# Delete the deleted collection in the default database and expect an error
with pytest.raises(NotFoundError):
sysdb.delete_collection(id=sample_collections[0].id)
# Delete the existing collection in the new database and expect an error
with pytest.raises(NotFoundError):
sysdb.delete_collection(id=sample_collections[1].id, database="new_database")
def test_create_update_with_database(sysdb: SysDB) -> None:
sysdb.reset_state()
# Create a new database
sysdb.create_database(id=uuid.uuid4(), name="new_database")
# Create a new collection in the new database
sysdb.create_collection(
id=sample_collections[0].id,
name=sample_collections[0].name,
configuration=sample_collections[0].get_configuration(),
segments=[sample_segment(sample_collections[0].id)],
metadata=sample_collections[0]["metadata"],
dimension=sample_collections[0]["dimension"],
database="new_database",
)
# Create a new collection in the default database
sysdb.create_collection(
id=sample_collections[1].id,
name=sample_collections[1].name,
configuration=sample_collections[1].get_configuration(),
segments=[sample_segment(sample_collections[1].id)],
metadata=sample_collections[1]["metadata"],
dimension=sample_collections[1]["dimension"],
)
# Update the collection in the default database
sysdb.update_collection(
id=sample_collections[1].id,
name="new_name_1",
)
# Check that the collection in the default database was updated
result = sysdb.get_collections(id=sample_collections[1]["id"])
assert len(result) == 1
assert result[0]["name"] == "new_name_1"
# Update the collection in the new database
sysdb.update_collection(
id=sample_collections[0].id,
name="new_name_0",
)
# Check that the collection in the new database was updated
result = sysdb.get_collections(
id=sample_collections[0]["id"], database="new_database"
)
assert len(result) == 1
assert result[0]["name"] == "new_name_0"
# Try to create the collection in the default database in the new database and expect an error
with pytest.raises(UniqueConstraintError):
sysdb.create_collection(
id=sample_collections[1].id,
name=sample_collections[1].name,
configuration=sample_collections[1].get_configuration(),
segments=[sample_segment(sample_collections[1].id)],
metadata=sample_collections[1]["metadata"],
dimension=sample_collections[1]["dimension"],
database="new_database",
)
def test_get_multiple_with_database(sysdb: SysDB) -> None:
sysdb.reset_state()
# Create a new database
sysdb.create_database(id=uuid.uuid4(), name="new_database")
# Create sample collections in the new database
for collection in sample_collections:
sysdb.create_collection(
id=collection.id,
name=collection.name,
configuration=collection.get_configuration(),
segments=[sample_segment(collection.id)],
metadata=collection["metadata"],
dimension=collection["dimension"],
database="new_database",
)
collection["database"] = "new_database"
# Get all collections in the new database
result = sysdb.get_collections(database="new_database")
assert len(result) == len(sample_collections)
assert sorted(result, key=lambda c: c.name) == sample_collections
# Get all collections in the default database
result = sysdb.get_collections()
assert len(result) == 0
def test_create_database_with_tenants(sysdb: SysDB) -> None:
sysdb.reset_state()
# Create a new tenant
sysdb.create_tenant(name="tenant1")
# Create tenant that already exits and expect an error
with pytest.raises(UniqueConstraintError):
sysdb.create_tenant(name="tenant1")
with pytest.raises(UniqueConstraintError):
sysdb.create_tenant(name=DEFAULT_TENANT)
# Create a new database within this tenant and also in the default tenant
sysdb.create_database(id=uuid.uuid4(), name="new_database", tenant="tenant1")
sysdb.create_database(id=uuid.uuid4(), name="new_database")
# Create a new collection in the new tenant
sysdb.create_collection(
id=sample_collections[0].id,
name=sample_collections[0].name,
configuration=sample_collections[0].get_configuration(),
segments=[sample_segment(sample_collections[0].id)],
metadata=sample_collections[0]["metadata"],
dimension=sample_collections[0]["dimension"],
database="new_database",
tenant="tenant1",
)
sample_collections[0]["tenant"] = "tenant1"
sample_collections[0]["database"] = "new_database"
# Create a new collection in the default tenant
sysdb.create_collection(
id=sample_collections[1].id,
name=sample_collections[1].name,
configuration=sample_collections[1].get_configuration(),
segments=[sample_segment(sample_collections[1].id)],
metadata=sample_collections[1]["metadata"],
dimension=sample_collections[1]["dimension"],
database="new_database",
)
sample_collections[1]["database"] = "new_database"
# Check that both tenants have the correct collections
result = sysdb.get_collections(database="new_database", tenant="tenant1")
assert len(result) == 1
assert result[0] == sample_collections[0]
result = sysdb.get_collections(database="new_database")
assert len(result) == 1
assert result[0] == sample_collections[1]
# Creating a collection id that already exists in a tenant that does not have it
# should error
with pytest.raises(UniqueConstraintError):
sysdb.create_collection(
id=sample_collections[0].id,
name=sample_collections[0].name,
configuration=sample_collections[0].get_configuration(),
segments=[sample_segment(sample_collections[0].id)],
metadata=sample_collections[0]["metadata"],
dimension=sample_collections[0]["dimension"],
database="new_database",
)
with pytest.raises(UniqueConstraintError):
sysdb.create_collection(
id=sample_collections[1].id,
name=sample_collections[1].name,
configuration=sample_collections[1].get_configuration(),
segments=[sample_segment(sample_collections[1].id)],
metadata=sample_collections[1]["metadata"],
dimension=sample_collections[1]["dimension"],
database="new_database",
tenant="tenant1",
)
# A new tenant DOES NOT have a default database. This does not error, instead 0
# results are returned
result = sysdb.get_collections(database=DEFAULT_DATABASE, tenant="tenant1")
assert len(result) == 0
def test_get_database_with_tenants(sysdb: SysDB) -> None:
sysdb.reset_state()
# Create a new tenant
sysdb.create_tenant(name="tenant1")
# Get the tenant and check that it exists
result = sysdb.get_tenant(name="tenant1")
assert result["name"] == "tenant1"
# Get a tenant that does not exist and expect an error
with pytest.raises(NotFoundError):
sysdb.get_tenant(name="tenant2")
# Create a new database within this tenant
sysdb.create_database(id=uuid.uuid4(), name="new_database", tenant="tenant1")
# Get the database and check that it exists
result = sysdb.get_database(name="new_database", tenant="tenant1")
assert result["name"] == "new_database"
assert result["tenant"] == "tenant1"
# Get a database that does not exist in a tenant that does exist and expect an error
with pytest.raises(NotFoundError):
sysdb.get_database(name="new_database1", tenant="tenant1")
# Get a database that does not exist in a tenant that does not exist and expect an
# error
with pytest.raises(NotFoundError):
sysdb.get_database(name="new_database1", tenant="tenant2")
# endregion
# region Segment tests
sample_segments = [
Segment(
id=uuid.UUID("00000000-d7d7-413b-92e1-731098a6e492"),
type="test_type_a",
scope=SegmentScope.VECTOR,
collection=sample_collections[0].id,
metadata={"test_str": "str1", "test_int": 1, "test_float": 1.3},
file_paths={},
),
Segment(
id=uuid.UUID("11111111-d7d7-413b-92e1-731098a6e492"),
type="test_type_b",
scope=SegmentScope.VECTOR,
collection=sample_collections[1].id,
metadata={"test_str": "str2", "test_int": 2, "test_float": 2.3},
file_paths={},
),
]
def test_create_get_delete_segments(sysdb: SysDB) -> None:
sysdb.reset_state()
# Keep track of segments created with a collection.
segments_created_with_collection = []
# Used to toggle between test_type_a and test_type_b
toggle_type = False
# Create collections along with segments.
for collection in sample_collections:
toggle_type = not toggle_type
segment = sample_segment(
collection_id=collection.id,
segment_type="test_type_a" if toggle_type else "test_type_b",
)
segments_created_with_collection.append(segment)
collection_result, created = sysdb.create_collection(
id=collection.id,
name=collection.name,
configuration=collection.get_configuration(),
segments=[segment],
metadata=collection["metadata"],
dimension=collection["dimension"],
)
assert created is True
results: List[Segment] = []
for collection in sample_collections:
results.extend(sysdb.get_segments(collection=collection.id))
results = sorted(results, key=lambda c: c["id"])
sorted_segments = sorted(segments_created_with_collection, key=lambda c: c["id"])
assert results == sorted_segments
# Duplicate create fails
with pytest.raises(UniqueConstraintError):
sysdb.create_segment(segments_created_with_collection[0])
# Find by id
for segment in segments_created_with_collection:
result = sysdb.get_segments(id=segment["id"], collection=segment["collection"])
assert result == [segment]
# Find by type
result = sysdb.get_segments(type="test_type_a", collection=sample_collections[0].id)
assert len(result) == 1
assert result[0]["collection"] == sample_collections[0].id
assert result[0] == segments_created_with_collection[0]
result = sysdb.get_segments(type="test_type_b", collection=sample_collections[1].id)
assert len(result) == 1
assert result[0] == segments_created_with_collection[1]
# Find by collection ID
result = sysdb.get_segments(collection=sample_collections[0].id)
assert len(result) == 1
assert result[0] == segments_created_with_collection[0]
# Find by type and collection ID (positive case)
result = sysdb.get_segments(type="test_type_a", collection=sample_collections[0].id)
assert len(result) == 1
assert result[0] == segments_created_with_collection[0]
# Find by type and collection ID (negative case)
result = sysdb.get_segments(type="test_type_b", collection=sample_collections[0].id)
assert len(result) == 0
# Delete Segments will be a NoOp due to atomic delete of collection and segments.
# See comments in coordinator.go for more details.
# Execute the call and expect no error or exception.
s1 = segments_created_with_collection[0]
sysdb.delete_segment(s1["collection"], s1["id"])
def test_update_segment(sysdb: SysDB) -> None:
metadata: Dict[str, Union[str, int, float]] = {
"test_str": "str1",
"test_int": 1,
"test_float": 1.3,
}
segment = Segment(
id=uuid.uuid4(),
type="test_type_a",
scope=SegmentScope.VECTOR,
collection=sample_collections[0].id,
metadata=metadata,
file_paths={},
)
sysdb.reset_state()
for c in sample_collections:
sysdb.create_collection(
id=c.id,
name=c.name,
configuration=c.get_configuration(),
segments=[sample_segment(c.id)],
metadata=c["metadata"],
dimension=c["dimension"],
)
sysdb.create_segment(segment)
# TODO: revisit update segment - push collection id
result = sysdb.get_segments(id=segment["id"], collection=segment["collection"])
result[0]["collection"] = segment["collection"]
assert result == [segment]
result = sysdb.get_segments(id=segment["id"], collection=segment["collection"])
result[0]["collection"] = segment["collection"]
assert result == [segment]
# Add a new metadata key
metadata["test_str2"] = "str2"
sysdb.update_segment(
segment["collection"], segment["id"], metadata={"test_str2": "str2"}
)
result = sysdb.get_segments(id=segment["id"], collection=segment["collection"])
result[0]["collection"] = segment["collection"]
assert result == [segment]
# Update a metadata key
metadata["test_str"] = "str3"
sysdb.update_segment(
segment["collection"], segment["id"], metadata={"test_str": "str3"}
)
result = sysdb.get_segments(id=segment["id"], collection=segment["collection"])
result[0]["collection"] = segment["collection"]
assert result == [segment]
# Delete a metadata key
del metadata["test_str"]
sysdb.update_segment(
segment["collection"], segment["id"], metadata={"test_str": None}
)
result = sysdb.get_segments(id=segment["id"], collection=segment["collection"])
result[0]["collection"] = segment["collection"]
assert result == [segment]
# Delete all metadata keys
segment["metadata"] = None
sysdb.update_segment(segment["collection"], segment["id"], metadata=None)
result = sysdb.get_segments(id=segment["id"], collection=segment["collection"])
result[0]["collection"] = segment["collection"]
assert result == [segment]
# endregion