from concurrent import futures from typing import Any, Dict, List, cast from uuid import UUID from overrides import overrides from chromadb.api.configuration import CollectionConfigurationInternal from chromadb.config import DEFAULT_DATABASE, DEFAULT_TENANT, Component, System from chromadb.proto.convert import ( from_proto_metadata, from_proto_update_metadata, from_proto_segment, from_proto_segment_scope, to_proto_collection, to_proto_segment, ) import chromadb.proto.chroma_pb2 as proto from chromadb.proto.coordinator_pb2 import ( CreateCollectionRequest, CreateCollectionResponse, CreateDatabaseRequest, CreateDatabaseResponse, CreateSegmentRequest, CreateSegmentResponse, CreateTenantRequest, CreateTenantResponse, DeleteCollectionRequest, DeleteCollectionResponse, DeleteSegmentRequest, DeleteSegmentResponse, GetCollectionsRequest, GetCollectionsResponse, GetCollectionWithSegmentsRequest, GetCollectionWithSegmentsResponse, GetDatabaseRequest, GetDatabaseResponse, GetSegmentsRequest, GetSegmentsResponse, GetTenantRequest, GetTenantResponse, ResetStateResponse, UpdateCollectionRequest, UpdateCollectionResponse, UpdateSegmentRequest, UpdateSegmentResponse, ) from chromadb.proto.coordinator_pb2_grpc import ( SysDBServicer, add_SysDBServicer_to_server, ) import grpc from google.protobuf.empty_pb2 import Empty from chromadb.types import Collection, Metadata, Segment, SegmentScope class GrpcMockSysDB(SysDBServicer, Component): """A mock sysdb implementation that can be used for testing the grpc client. It stores state in simple python data structures instead of a database.""" _server: grpc.Server _server_port: int _segments: Dict[str, Segment] = {} _collection_to_segments: Dict[str, List[str]] = {} _tenants_to_databases_to_collections: Dict[ str, Dict[str, Dict[str, Collection]] ] = {} _tenants_to_database_to_id: Dict[str, Dict[str, UUID]] = {} def __init__(self, system: System): self._server_port = system.settings.require("chroma_server_grpc_port") return super().__init__(system) @overrides def start(self) -> None: self._server = grpc.server(futures.ThreadPoolExecutor(max_workers=10)) add_SysDBServicer_to_server(self, self._server) # type: ignore self._server.add_insecure_port(f"[::]:{self._server_port}") self._server.start() return super().start() @overrides def stop(self) -> None: self._server.stop(None) return super().stop() @overrides def reset_state(self) -> None: self._segments = {} self._tenants_to_databases_to_collections = {} # Create defaults self._tenants_to_databases_to_collections[DEFAULT_TENANT] = {} self._tenants_to_databases_to_collections[DEFAULT_TENANT][DEFAULT_DATABASE] = {} self._tenants_to_database_to_id[DEFAULT_TENANT] = {} self._tenants_to_database_to_id[DEFAULT_TENANT][DEFAULT_DATABASE] = UUID(int=0) return super().reset_state() @overrides(check_signature=False) def CreateDatabase( self, request: CreateDatabaseRequest, context: grpc.ServicerContext ) -> CreateDatabaseResponse: tenant = request.tenant database = request.name if tenant not in self._tenants_to_databases_to_collections: context.abort(grpc.StatusCode.NOT_FOUND, f"Tenant {tenant} not found") if database in self._tenants_to_databases_to_collections[tenant]: context.abort( grpc.StatusCode.ALREADY_EXISTS, f"Database {database} already exists" ) self._tenants_to_databases_to_collections[tenant][database] = {} self._tenants_to_database_to_id[tenant][database] = UUID(hex=request.id) return CreateDatabaseResponse() @overrides(check_signature=False) def GetDatabase( self, request: GetDatabaseRequest, context: grpc.ServicerContext ) -> GetDatabaseResponse: tenant = request.tenant database = request.name if tenant not in self._tenants_to_databases_to_collections: context.abort(grpc.StatusCode.NOT_FOUND, f"Tenant {tenant} not found") if database not in self._tenants_to_databases_to_collections[tenant]: context.abort(grpc.StatusCode.NOT_FOUND, f"Database {database} not found") id = self._tenants_to_database_to_id[tenant][database] return GetDatabaseResponse( database=proto.Database(id=id.hex, name=database, tenant=tenant), ) @overrides(check_signature=False) def CreateTenant( self, request: CreateTenantRequest, context: grpc.ServicerContext ) -> CreateTenantResponse: tenant = request.name if tenant in self._tenants_to_databases_to_collections: context.abort( grpc.StatusCode.ALREADY_EXISTS, f"Tenant {tenant} already exists" ) self._tenants_to_databases_to_collections[tenant] = {} self._tenants_to_database_to_id[tenant] = {} return CreateTenantResponse() @overrides(check_signature=False) def GetTenant( self, request: GetTenantRequest, context: grpc.ServicerContext ) -> GetTenantResponse: tenant = request.name if tenant not in self._tenants_to_databases_to_collections: context.abort(grpc.StatusCode.NOT_FOUND, f"Tenant {tenant} not found") return GetTenantResponse( tenant=proto.Tenant(name=tenant), ) # We are forced to use check_signature=False because the generated proto code # does not have type annotations for the request and response objects. # TODO: investigate generating types for the request and response objects @overrides(check_signature=False) def CreateSegment( self, request: CreateSegmentRequest, context: grpc.ServicerContext ) -> CreateSegmentResponse: segment = from_proto_segment(request.segment) return self.CreateSegmentHelper(segment, context) def CreateSegmentHelper(self, segment: Segment, context: grpc.ServicerContext) -> CreateSegmentResponse: if segment["id"].hex in self._segments: context.abort( grpc.StatusCode.ALREADY_EXISTS, f"Segment {segment['id']} already exists", ) self._segments[segment["id"].hex] = segment return CreateSegmentResponse() @overrides(check_signature=False) def DeleteSegment( self, request: DeleteSegmentRequest, context: grpc.ServicerContext ) -> DeleteSegmentResponse: id_to_delete = request.id if id_to_delete in self._segments: del self._segments[id_to_delete] return DeleteSegmentResponse() else: context.abort( grpc.StatusCode.NOT_FOUND, f"Segment {id_to_delete} not found" ) @overrides(check_signature=False) def GetSegments( self, request: GetSegmentsRequest, context: grpc.ServicerContext ) -> GetSegmentsResponse: target_id = UUID(hex=request.id) if request.HasField("id") else None target_type = request.type if request.HasField("type") else None target_scope = ( from_proto_segment_scope(request.scope) if request.HasField("scope") else None ) target_collection = UUID(hex=request.collection) found_segments = [] for segment in self._segments.values(): if target_id and segment["id"] != target_id: continue if target_type and segment["type"] != target_type: continue if target_scope and segment["scope"] != target_scope: continue if target_collection and segment["collection"] != target_collection: continue found_segments.append(segment) return GetSegmentsResponse( segments=[to_proto_segment(segment) for segment in found_segments] ) @overrides(check_signature=False) def UpdateSegment( self, request: UpdateSegmentRequest, context: grpc.ServicerContext ) -> UpdateSegmentResponse: id_to_update = UUID(request.id) if id_to_update.hex not in self._segments: context.abort( grpc.StatusCode.NOT_FOUND, f"Segment {id_to_update} not found" ) else: segment = self._segments[id_to_update.hex] if request.HasField("metadata"): target = cast(Dict[str, Any], segment["metadata"]) if segment["metadata"] is None: segment["metadata"] = {} self._merge_metadata(target, request.metadata) if request.HasField("reset_metadata") and request.reset_metadata: segment["metadata"] = {} return UpdateSegmentResponse() @overrides(check_signature=False) def CreateCollection( self, request: CreateCollectionRequest, context: grpc.ServicerContext ) -> CreateCollectionResponse: collection_name = request.name tenant = request.tenant database = request.database if tenant not in self._tenants_to_databases_to_collections: context.abort(grpc.StatusCode.NOT_FOUND, f"Tenant {tenant} not found") if database not in self._tenants_to_databases_to_collections[tenant]: context.abort(grpc.StatusCode.NOT_FOUND, f"Database {database} not found") # Check if the collection already exists globally by id for ( search_tenant, databases, ) in self._tenants_to_databases_to_collections.items(): for search_database, search_collections in databases.items(): if request.id in search_collections: if ( search_tenant != request.tenant or search_database != request.database ): context.abort( grpc.StatusCode.ALREADY_EXISTS, f"Collection {request.id} already exists in tenant {search_tenant} database {search_database}", ) elif not request.get_or_create: # If the id exists for this tenant and database, and we are not doing a get_or_create, then # we should return an already exists error context.abort( grpc.StatusCode.ALREADY_EXISTS, f"Collection {request.id} already exists in tenant {search_tenant} database {search_database}", ) # Check if the collection already exists in this database by name collections = self._tenants_to_databases_to_collections[tenant][database] matches = [c for c in collections.values() if c["name"] == collection_name] assert len(matches) <= 1 if len(matches) > 0: if request.get_or_create: existing_collection = matches[0] return CreateCollectionResponse( collection=to_proto_collection(existing_collection), created=False, ) context.abort( grpc.StatusCode.ALREADY_EXISTS, f"Collection {collection_name} already exists", ) configuration = CollectionConfigurationInternal.from_json_str( request.configuration_json_str ) id = UUID(hex=request.id) new_collection = Collection( id=id, name=request.name, configuration=configuration, metadata=from_proto_metadata(request.metadata), dimension=request.dimension, database=database, tenant=tenant, version=0, ) # Check that segments are unique and do not already exist # Keep a track of the segments that are being added segments_added = [] # Create segments for the collection for segment_proto in request.segments: segment = from_proto_segment(segment_proto) if segment["id"].hex in self._segments: # Remove the already added segment since we need to roll back for s in segments_added: self.DeleteSegment(DeleteSegmentRequest(id=s), context) context.abort(grpc.StatusCode.ALREADY_EXISTS, f"Segment {segment['id']} already exists") self.CreateSegmentHelper(segment, context) segments_added.append(segment["id"].hex) collections[request.id] = new_collection collection_unique_key = f"{tenant}:{database}:{request.id}" self._collection_to_segments[collection_unique_key] = segments_added return CreateCollectionResponse( collection=to_proto_collection(new_collection), created=True, ) @overrides(check_signature=False) def DeleteCollection( self, request: DeleteCollectionRequest, context: grpc.ServicerContext ) -> DeleteCollectionResponse: collection_id = request.id tenant = request.tenant database = request.database if tenant not in self._tenants_to_databases_to_collections: context.abort(grpc.StatusCode.NOT_FOUND, f"Tenant {tenant} not found") if database not in self._tenants_to_databases_to_collections[tenant]: context.abort(grpc.StatusCode.NOT_FOUND, f"Database {database} not found") collections = self._tenants_to_databases_to_collections[tenant][database] if collection_id in collections: del collections[collection_id] collection_unique_key = f"{tenant}:{database}:{collection_id}" segment_ids = self._collection_to_segments[collection_unique_key] if segment_ids: # Delete segments if provided. for segment_id in segment_ids: del self._segments[segment_id] return DeleteCollectionResponse() else: context.abort( grpc.StatusCode.NOT_FOUND, f"Collection {collection_id} not found" ) @overrides(check_signature=False) def GetCollections( self, request: GetCollectionsRequest, context: grpc.ServicerContext ) -> GetCollectionsResponse: target_id = UUID(hex=request.id) if request.HasField("id") else None target_name = request.name if request.HasField("name") else None allCollections = {} for tenant, databases in self._tenants_to_databases_to_collections.items(): for database, collections in databases.items(): if request.tenant != "" and tenant != request.tenant: continue if request.database != "" and database != request.database: continue allCollections.update(collections) print( f"Tenant: {tenant}, Database: {database}, Collections: {collections}" ) found_collections = [] for collection in allCollections.values(): if target_id and collection["id"] != target_id: continue if target_name and collection["name"] != target_name: continue found_collections.append(collection) return GetCollectionsResponse( collections=[ to_proto_collection(collection) for collection in found_collections ] ) @overrides(check_signature=False) def GetCollectionWithSegments( self, request: GetCollectionWithSegmentsRequest, context: grpc.ServicerContext ) -> GetCollectionWithSegmentsResponse: allCollections = {} for tenant, databases in self._tenants_to_databases_to_collections.items(): for database, collections in databases.items(): allCollections.update(collections) print( f"Tenant: {tenant}, Database: {database}, Collections: {collections}" ) collection = allCollections.get(request.id, None) if collection is None: context.abort(grpc.StatusCode.NOT_FOUND, f"Collection with id {request.id} not found") collection_unique_key = f"{collection.tenant}:{collection.database}:{request.id}" segments = [self._segments[id] for id in self._collection_to_segments[collection_unique_key]] if {segment["scope"] for segment in segments} != {SegmentScope.METADATA, SegmentScope.RECORD, SegmentScope.VECTOR}: context.abort(grpc.StatusCode.INTERNAL, f"Incomplete segments for collection {collection}: {segments}") return GetCollectionWithSegmentsResponse( collection=to_proto_collection(collection), segments=[to_proto_segment(segment) for segment in segments] ) @overrides(check_signature=False) def UpdateCollection( self, request: UpdateCollectionRequest, context: grpc.ServicerContext ) -> UpdateCollectionResponse: id_to_update = UUID(request.id) # Find the collection with this id collections = {} for tenant, databases in self._tenants_to_databases_to_collections.items(): for database, maybe_collections in databases.items(): if id_to_update.hex in maybe_collections: collections = maybe_collections if id_to_update.hex not in collections: context.abort( grpc.StatusCode.NOT_FOUND, f"Collection {id_to_update} not found" ) else: collection = collections[id_to_update.hex] if request.HasField("name"): collection["name"] = request.name if request.HasField("dimension"): collection["dimension"] = request.dimension if request.HasField("metadata"): # TODO: IN SysDB SQlite we have technical debt where we # replace the entire metadata dict with the new one. We should # fix that by merging it. For now we just do the same thing here update_metadata = from_proto_update_metadata(request.metadata) cleaned_metadata = None if update_metadata is not None: cleaned_metadata = {} for key, value in update_metadata.items(): if value is not None: cleaned_metadata[key] = value collection["metadata"] = cleaned_metadata elif request.HasField("reset_metadata"): if request.reset_metadata: collection["metadata"] = {} return UpdateCollectionResponse() @overrides(check_signature=False) def ResetState( self, request: Empty, context: grpc.ServicerContext ) -> ResetStateResponse: self.reset_state() return ResetStateResponse() def _merge_metadata(self, target: Metadata, source: proto.UpdateMetadata) -> None: target_metadata = cast(Dict[str, Any], target) source_metadata = cast(Dict[str, Any], from_proto_update_metadata(source)) target_metadata.update(source_metadata) # If a key has a None value, remove it from the metadata for key, value in source_metadata.items(): if value is None and key in target: del target_metadata[key]