[ENH]: Return database id in get collections call from sysdb (#4686)

sanketkedia · web-flow · commit 72985bd9965c · 2025-06-20T13:28:52.000-07:00
## Description of changes

_Summarize the changes made by this PR._

- Improvements &amp; Bug fixes
  - Organizing data of collections in s3 by s3 prefixes requires knowledge about database id in the write path. This PR populates and returns database id from sysdb in get_collections rpc
  - We skip serializing and returning this from FE to client
  - Even though the collection model is the same for both local and distributed, we don't set it in local currently. However, it can be done if needed
- New functionality
  - ...

## Test plan
_How are these changes tested?_
- [x] Tests pass locally with `pytest` for python, `yarn test` for js, `cargo test` for rust

## Documentation Changes
None
diff --git a/go/pkg/sysdb/grpc/proto_model_convert.go b/go/pkg/sysdb/grpc/proto_model_convert.go
@@ -42,6 +42,7 @@ func convertCollectionToProto(collection *model.Collection) *coordinatorpb.Colle
 		return nil
 	}
 
+	dbId := collection.DatabaseId.String()
 	collectionpb := &coordinatorpb.Collection{
 		Id:                         collection.ID.String(),
 		Name:                       collection.Name,
@@ -60,6 +61,7 @@ func convertCollectionToProto(collection *model.Collection) *coordinatorpb.Colle
 			Seconds: collection.UpdatedAt,
 			Nanos:   0,
 		},
+		DatabaseId: &dbId,
 	}
 
 	if collection.RootCollectionID != nil {
diff --git a/idl/chromadb/proto/chroma.proto b/idl/chromadb/proto/chroma.proto
@@ -64,6 +64,8 @@ message Collection {
   optional string root_collection_id = 14;
   optional string lineage_file_path = 15;
   google.protobuf.Timestamp updated_at = 16;
+  // This is the database id of the collection.
+  optional string database_id = 17;
 }
 
 message Database {
diff --git a/rust/garbage_collector/src/helper.rs b/rust/garbage_collector/src/helper.rs
@@ -2,9 +2,8 @@ use chroma_types::chroma_proto::log_service_client::LogServiceClient;
 use chroma_types::chroma_proto::query_executor_client::QueryExecutorClient;
 use chroma_types::chroma_proto::sys_db_client::SysDbClient;
 use chroma_types::chroma_proto::{
-    Collection, CreateCollectionRequest, CreateDatabaseRequest, CreateTenantRequest,
-    FilterOperator, GetCollectionWithSegmentsRequest, GetPlan, KnnOperator, KnnPlan,
-    KnnProjectionOperator, LimitOperator, ListCollectionVersionsRequest,
+    CreateCollectionRequest, CreateDatabaseRequest, CreateTenantRequest, FilterOperator,
+    GetCollectionWithSegmentsRequest, GetPlan, LimitOperator, ListCollectionVersionsRequest,
     ListCollectionVersionsResponse, OperationRecord, ProjectionOperator, PushLogsRequest,
     ScanOperator, Segment, SegmentScope, Vector,
 };
@@ -163,74 +162,6 @@ impl ChromaGrpcClients {
         }
     }
 
-    #[allow(dead_code)]
-    pub async fn query_collection(
-        &mut self,
-        collection_id: &str,
-        query_embedding: Vec<f32>,
-    ) -> Result<Vec<(String, f32)>, Box<dyn std::error::Error>> {
-        // Convert f32 vector to bytes
-        let vector_bytes: Vec<u8> = query_embedding
-            .iter()
-            .flat_map(|&x| x.to_le_bytes().to_vec())
-            .collect();
-
-        let knn_plan = KnnPlan {
-            scan: Some(ScanOperator {
-                collection: Some(Collection {
-                    id: collection_id.to_string(),
-                    name: String::new(),
-                    database: String::new(),
-                    tenant: String::new(),
-                    dimension: Some(query_embedding.len() as i32),
-                    configuration_json_str: String::new(),
-                    metadata: None,
-                    log_position: 0,
-                    version: 0,
-                    total_records_post_compaction: 0,
-                    size_bytes_post_compaction: 0,
-                    last_compaction_time_secs: 0,
-                    version_file_path: None,
-                    root_collection_id: None,
-                    lineage_file_path: None,
-                    updated_at: None,
-                }),
-                knn: None,
-                metadata: None,
-                record: None,
-            }),
-            filter: None,
-            knn: Some(KnnOperator {
-                embeddings: vec![Vector {
-                    dimension: query_embedding.len() as i32,
-                    vector: vector_bytes,
-                    encoding: 0,
-                }],
-                fetch: 2,
-            }),
-            projection: Some(KnnProjectionOperator {
-                projection: Some(ProjectionOperator {
-                    document: false,
-                    embedding: false,
-                    metadata: false,
-                }),
-                distance: true,
-            }),
-        };
-
-        let response = self.query_executor.knn(knn_plan).await?;
-        let results = response.into_inner().results;
-
-        let mut id_distances = Vec::new();
-        for result in results {
-            for record in result.records {
-                id_distances.push((record.record.unwrap().id, record.distance.unwrap_or(0.0)));
-            }
-        }
-
-        Ok(id_distances)
-    }
-
     pub async fn get_records(
         &mut self,
         collection_id: String,
diff --git a/rust/sysdb/src/sqlite.rs b/rust/sysdb/src/sqlite.rs
@@ -9,11 +9,12 @@ use chroma_sqlite::table;
 use chroma_types::{
     Collection, CollectionAndSegments, CollectionMetadataUpdate, CollectionUuid,
     CreateCollectionError, CreateCollectionResponse, CreateDatabaseError, CreateDatabaseResponse,
-    CreateTenantError, CreateTenantResponse, Database, DeleteCollectionError, DeleteDatabaseError,
-    DeleteDatabaseResponse, GetCollectionWithSegmentsError, GetCollectionsError, GetDatabaseError,
-    GetSegmentsError, GetTenantError, GetTenantResponse, InternalCollectionConfiguration,
-    ListDatabasesError, Metadata, MetadataValue, ResetError, ResetResponse, Segment, SegmentScope,
-    SegmentType, SegmentUuid, UpdateCollectionConfiguration, UpdateCollectionError,
+    CreateTenantError, CreateTenantResponse, Database, DatabaseUuid, DeleteCollectionError,
+    DeleteDatabaseError, DeleteDatabaseResponse, GetCollectionWithSegmentsError,
+    GetCollectionsError, GetDatabaseError, GetSegmentsError, GetTenantError, GetTenantResponse,
+    InternalCollectionConfiguration, ListDatabasesError, Metadata, MetadataValue, ResetError,
+    ResetResponse, Segment, SegmentScope, SegmentType, SegmentUuid, UpdateCollectionConfiguration,
+    UpdateCollectionError,
 };
 use futures::TryStreamExt;
 use sea_query_binder::SqlxBinder;
@@ -289,6 +290,8 @@ impl SqliteSysDb {
                     _ => CreateCollectionError::Internal(e.into()),
                 })?;
         let database_id = database_result.get::<&str, _>(0);
+        let database_uuid = DatabaseUuid::from_str(database_id)
+            .map_err(|_| CreateCollectionError::DatabaseIdParseError)?;
 
         sqlx::query(
             r#"
@@ -343,6 +346,7 @@ impl SqliteSysDb {
             root_collection_id: None,
             lineage_file_path: None,
             updated_at: SystemTime::UNIX_EPOCH,
+            database_id: database_uuid,
         })
     }
 
@@ -671,6 +675,7 @@ impl SqliteSysDb {
             .column((table::Collections::Table, table::Collections::Dimension))
             .column((table::Databases::Table, table::Databases::TenantId))
             .column((table::Databases::Table, table::Databases::Name))
+            .column((table::Collections::Table, table::Collections::DatabaseId))
             .columns([
                 table::CollectionMetadata::Key,
                 table::CollectionMetadata::StrValue,
@@ -720,6 +725,10 @@ impl SqliteSysDb {
                     }
                     None => InternalCollectionConfiguration::default_hnsw(),
                 };
+                let database_id = match DatabaseUuid::from_str(first_row.get(6)) {
+                    Ok(db_id) => db_id,
+                    Err(_) => return Some(Err(GetCollectionsError::DatabaseId)),
+                };
 
                 Some(Ok(Collection {
                     collection_id,
@@ -738,6 +747,7 @@ impl SqliteSysDb {
                     root_collection_id: None,
                     lineage_file_path: None,
                     updated_at: SystemTime::UNIX_EPOCH,
+                    database_id,
                 }))
             })
             .collect::<Result<Vec<_>, GetCollectionsError>>()?;
diff --git a/rust/sysdb/src/sysdb.rs b/rust/sysdb/src/sysdb.rs
@@ -20,9 +20,9 @@ use chroma_types::{
 };
 use chroma_types::{
     BatchGetCollectionSoftDeleteStatusError, BatchGetCollectionVersionFilePathsError, Collection,
-    CollectionConversionError, CollectionUuid, CountForksError, FinishDatabaseDeletionError,
-    FlushCompactionResponse, FlushCompactionResponseConversionError, ForkCollectionError, Segment,
-    SegmentConversionError, SegmentScope, Tenant,
+    CollectionConversionError, CollectionUuid, CountForksError, DatabaseUuid,
+    FinishDatabaseDeletionError, FlushCompactionResponse, FlushCompactionResponseConversionError,
+    ForkCollectionError, Segment, SegmentConversionError, SegmentScope, Tenant,
 };
 use std::collections::HashMap;
 use std::fmt::Debug;
@@ -266,6 +266,7 @@ impl SysDb {
                     root_collection_id: None,
                     lineage_file_path: None,
                     updated_at: SystemTime::now(),
+                    database_id: DatabaseUuid::new(),
                 };
 
                 test_sysdb.add_collection(collection.clone());
diff --git a/rust/types/src/api_types.rs b/rust/types/src/api_types.rs
@@ -647,6 +647,8 @@ pub enum CreateCollectionError {
     SpannNotImplemented,
     #[error("HNSW is not supported on this platform")]
     HnswNotSupported,
+    #[error("Failed to parse db id")]
+    DatabaseIdParseError,
 }
 
 impl ChromaError for CreateCollectionError {
@@ -663,6 +665,7 @@ impl ChromaError for CreateCollectionError {
             CreateCollectionError::Aborted(_) => ErrorCodes::Aborted,
             CreateCollectionError::SpannNotImplemented => ErrorCodes::InvalidArgument,
             CreateCollectionError::HnswNotSupported => ErrorCodes::InvalidArgument,
+            CreateCollectionError::DatabaseIdParseError => ErrorCodes::Internal,
         }
     }
 }
@@ -689,6 +692,8 @@ pub enum GetCollectionsError {
     Configuration(#[from] serde_json::Error),
     #[error("Could not deserialize collection ID")]
     CollectionId(#[from] uuid::Error),
+    #[error("Could not deserialize database ID")]
+    DatabaseId,
 }
 
 impl ChromaError for GetCollectionsError {
@@ -697,6 +702,7 @@ impl ChromaError for GetCollectionsError {
             GetCollectionsError::Internal(err) => err.code(),
             GetCollectionsError::Configuration(_) => ErrorCodes::Internal,
             GetCollectionsError::CollectionId(_) => ErrorCodes::Internal,
+            GetCollectionsError::DatabaseId => ErrorCodes::Internal,
         }
     }
 }
diff --git a/rust/types/src/collection.rs b/rust/types/src/collection.rs
@@ -1,3 +1,5 @@
+use std::str::FromStr;
+
 use super::{Metadata, MetadataValueConversionError};
 use crate::{
     chroma_proto, test_segment, CollectionConfiguration, InternalCollectionConfiguration, Segment,
@@ -30,6 +32,29 @@ use pyo3::types::PyAnyMethods;
 )]
 pub struct CollectionUuid(pub Uuid);
 
+/// DatabaseUuid is a wrapper around Uuid to provide a type for the database id.
+#[derive(
+    Copy,
+    Clone,
+    Debug,
+    Default,
+    Deserialize,
+    Eq,
+    PartialEq,
+    Ord,
+    PartialOrd,
+    Hash,
+    Serialize,
+    ToSchema,
+)]
+pub struct DatabaseUuid(pub Uuid);
+
+impl DatabaseUuid {
+    pub fn new() -> Self {
+        DatabaseUuid(Uuid::new_v4())
+    }
+}
+
 impl CollectionUuid {
     pub fn new() -> Self {
         CollectionUuid(Uuid::new_v4())
@@ -51,6 +76,17 @@ impl std::str::FromStr for CollectionUuid {
     }
 }
 
+impl std::str::FromStr for DatabaseUuid {
+    type Err = uuid::Error;
+
+    fn from_str(s: &str) -> Result<Self, Self::Err> {
+        match Uuid::parse_str(s) {
+            Ok(uuid) => Ok(DatabaseUuid(uuid)),
+            Err(err) => Err(err),
+        }
+    }
+}
+
 impl std::fmt::Display for CollectionUuid {
     fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
         write!(f, "{}", self.0)
@@ -107,6 +143,8 @@ pub struct Collection {
     pub lineage_file_path: Option<String>,
     #[serde(skip, default = "SystemTime::now")]
     pub updated_at: SystemTime,
+    #[serde(skip)]
+    pub database_id: DatabaseUuid,
 }
 
 impl Default for Collection {
@@ -128,6 +166,7 @@ impl Default for Collection {
             root_collection_id: None,
             lineage_file_path: None,
             updated_at: SystemTime::now(),
+            database_id: DatabaseUuid::new(),
         }
     }
 }
@@ -218,11 +257,8 @@ impl TryFrom<chroma_proto::Collection> for Collection {
     type Error = CollectionConversionError;
 
     fn try_from(proto_collection: chroma_proto::Collection) -> Result<Self, Self::Error> {
-        let collection_uuid = match Uuid::try_parse(&proto_collection.id) {
-            Ok(uuid) => uuid,
-            Err(_) => return Err(CollectionConversionError::InvalidUuid),
-        };
-        let collection_id = CollectionUuid(collection_uuid);
+        let collection_id = CollectionUuid::from_str(&proto_collection.id)
+            .map_err(|_| CollectionConversionError::InvalidUuid)?;
         let collection_metadata: Option<Metadata> = match proto_collection.metadata {
             Some(proto_metadata) => match proto_metadata.try_into() {
                 Ok(metadata) => Some(metadata),
@@ -238,6 +274,12 @@ impl TryFrom<chroma_proto::Collection> for Collection {
             }
             None => SystemTime::now(),
         };
+        // TOOD(Sanket): this should be updated to error with "missing field" once all SysDb deployments are up-to-date
+        let database_id = match proto_collection.database_id {
+            Some(db_id) => DatabaseUuid::from_str(&db_id)
+                .map_err(|_| CollectionConversionError::InvalidUuid)?,
+            None => DatabaseUuid::new(),
+        };
         Ok(Collection {
             collection_id,
             name: proto_collection.name,
@@ -257,6 +299,7 @@ impl TryFrom<chroma_proto::Collection> for Collection {
                 .map(|uuid| CollectionUuid(Uuid::try_parse(&uuid).unwrap())),
             lineage_file_path: proto_collection.lineage_file_path,
             updated_at,
+            database_id,
         })
     }
 }
@@ -296,6 +339,7 @@ impl TryFrom<Collection> for chroma_proto::Collection {
             root_collection_id: value.root_collection_id.map(|uuid| uuid.0.to_string()),
             lineage_file_path: value.lineage_file_path,
             updated_at: Some(value.updated_at.into()),
+            database_id: Some(value.database_id.0.to_string()),
         })
     }
 }
@@ -348,6 +392,7 @@ mod test {
                 seconds: 1,
                 nanos: 1,
             }),
+            database_id: Some("00000000-0000-0000-0000-000000000000".to_string()),
         };
         let converted_collection: Collection = proto_collection.try_into().unwrap();
         assert_eq!(
@@ -378,6 +423,7 @@ mod test {
             converted_collection.updated_at,
             SystemTime::UNIX_EPOCH + Duration::new(1, 1)
         );
+        assert_eq!(converted_collection.database_id, DatabaseUuid(Uuid::nil()));
     }
 
     #[test]
diff --git a/rust/worker/src/server.rs b/rust/worker/src/server.rs
@@ -513,6 +513,7 @@ mod tests {
 
     fn scan() -> chroma_proto::ScanOperator {
         let collection_id = Uuid::new_v4().to_string();
+        let database_id = Uuid::new_v4().to_string();
         chroma_proto::ScanOperator {
             collection: Some(chroma_proto::Collection {
                 id: collection_id.clone(),
@@ -522,6 +523,7 @@ mod tests {
                 dimension: None,
                 tenant: "test-tenant".to_string(),
                 database: "test-database".to_string(),
+                database_id: Some(database_id.clone()),
                 ..Default::default()
             }),
             knn: Some(chroma_proto::Segment {

Original file line number	Diff line number	Diff line change
`@@ -42,6 +42,7 @@ func convertCollectionToProto(collection model.Collection) coordinatorpb.Colle`
`42`	`42`	`return nil`
`43`	`43`	`}`
`44`	`44`
	`45`	`+ dbId := collection.DatabaseId.String()`
`45`	`46`	`collectionpb := &coordinatorpb.Collection{`
`46`	`47`	`Id: collection.ID.String(),`
`47`	`48`	`Name: collection.Name,`
`@@ -60,6 +61,7 @@ func convertCollectionToProto(collection model.Collection) coordinatorpb.Colle`
`60`	`61`	`Seconds: collection.UpdatedAt,`
`61`	`62`	`Nanos: 0,`
`62`	`63`	`},`
	`64`	`+ DatabaseId: &dbId,`
`63`	`65`	`}`
`64`	`66`
`65`	`67`	`if collection.RootCollectionID != nil {`
Original file line number	Diff line number	Diff line change
`@@ -64,6 +64,8 @@ message Collection {`
`64`	`64`	`optional string root_collection_id = 14;`
`65`	`65`	`optional string lineage_file_path = 15;`
`66`	`66`	`google.protobuf.Timestamp updated_at = 16;`
	`67`	`+ // This is the database id of the collection.`
	`68`	`+ optional string database_id = 17;`
`67`	`69`	`}`
`68`	`70`
`69`	`71`	`message Database {`
Original file line number	Diff line number	Diff line change
`@@ -647,6 +647,8 @@ pub enum CreateCollectionError {`
`647`	`647`	`SpannNotImplemented,`
`648`	`648`	`#[error("HNSW is not supported on this platform")]`
`649`	`649`	`HnswNotSupported,`
	`650`	`+ #[error("Failed to parse db id")]`
	`651`	`+ DatabaseIdParseError,`
`650`	`652`	`}`
`651`	`653`
`652`	`654`	`impl ChromaError for CreateCollectionError {`
`@@ -663,6 +665,7 @@ impl ChromaError for CreateCollectionError {`
`663`	`665`	`CreateCollectionError::Aborted(_) => ErrorCodes::Aborted,`
`664`	`666`	`CreateCollectionError::SpannNotImplemented => ErrorCodes::InvalidArgument,`
`665`	`667`	`CreateCollectionError::HnswNotSupported => ErrorCodes::InvalidArgument,`
	`668`	`+ CreateCollectionError::DatabaseIdParseError => ErrorCodes::Internal,`
`666`	`669`	`}`
`667`	`670`	`}`
`668`	`671`	`}`
`@@ -689,6 +692,8 @@ pub enum GetCollectionsError {`
`689`	`692`	`Configuration(#[from] serde_json::Error),`
`690`	`693`	`#[error("Could not deserialize collection ID")]`
`691`	`694`	`CollectionId(#[from] uuid::Error),`
	`695`	`+ #[error("Could not deserialize database ID")]`
	`696`	`+ DatabaseId,`
`692`	`697`	`}`
`693`	`698`
`694`	`699`	`impl ChromaError for GetCollectionsError {`
`@@ -697,6 +702,7 @@ impl ChromaError for GetCollectionsError {`
`697`	`702`	`GetCollectionsError::Internal(err) => err.code(),`
`698`	`703`	`GetCollectionsError::Configuration(_) => ErrorCodes::Internal,`
`699`	`704`	`GetCollectionsError::CollectionId(_) => ErrorCodes::Internal,`
	`705`	`+ GetCollectionsError::DatabaseId => ErrorCodes::Internal,`
`700`	`706`	`}`
`701`	`707`	`}`
`702`	`708`	`}`