[ENH]: add batch get version file paths method to Sysdb

codetheweb · codetheweb · commit ea366a572965 · 2025-05-15T13:07:59.000-07:00
diff --git a/go/pkg/sysdb/coordinator/coordinator.go b/go/pkg/sysdb/coordinator/coordinator.go
@@ -263,6 +263,10 @@ func (s *Coordinator) DeleteCollectionVersion(ctx context.Context, req *coordina
 	return s.catalog.DeleteCollectionVersion(ctx, req)
 }
 
+func (s *Coordinator) BatchGetCollectionVersionFilePaths(ctx context.Context, req *coordinatorpb.BatchGetCollectionVersionFilePathsRequest) (*coordinatorpb.BatchGetCollectionVersionFilePathsResponse, error) {
+	return s.catalog.BatchGetCollectionVersionFilePaths(ctx, req)
+}
+
 // SetDeleteMode sets the delete mode for testing
 func (c *Coordinator) SetDeleteMode(mode DeleteMode) {
 	c.deleteMode = mode
diff --git a/go/pkg/sysdb/coordinator/coordinator_test.go b/go/pkg/sysdb/coordinator/coordinator_test.go
@@ -1564,6 +1564,37 @@ func (suite *APIsTestSuite) TestForkCollection() {
 	suite.True(exists, "Lineage file should exist in S3")
 }
 
+func (suite *APIsTestSuite) TestBatchGetCollectionVersionFilePaths() {
+	ctx := context.Background()
+
+	// Create a new collection
+	newCollection := &model.CreateCollection{
+		ID:           types.NewUniqueID(),
+		Name:         "test_batch_get_collection_version_file_paths",
+		TenantID:     suite.tenantName,
+		DatabaseName: suite.databaseName,
+	}
+
+	newSegments := []*model.CreateSegment{}
+
+	// Create the collection
+	suite.coordinator.catalog.versionFileEnabled = true
+	_, _, err := suite.coordinator.CreateCollectionAndSegments(ctx, newCollection, newSegments)
+	suite.NoError(err)
+
+	// Get the version file paths for the collection
+	versionFilePaths, err := suite.coordinator.BatchGetCollectionVersionFilePaths(ctx, &coordinatorpb.BatchGetCollectionVersionFilePathsRequest{
+		CollectionIds: []string{newCollection.ID.String()},
+	})
+	suite.NoError(err)
+	suite.Len(versionFilePaths.CollectionIdToVersionFilePath, 1)
+
+	// Verify version file exists in S3
+	exists, err := suite.s3MetaStore.HasObjectWithPrefix(ctx, versionFilePaths.CollectionIdToVersionFilePath[newCollection.ID.String()])
+	suite.NoError(err)
+	suite.True(exists, "Version file should exist in S3")
+}
+
 func (suite *APIsTestSuite) TestCountForks() {
 	ctx := context.Background()
 
diff --git a/go/pkg/sysdb/coordinator/table_catalog.go b/go/pkg/sysdb/coordinator/table_catalog.go
@@ -2000,6 +2000,20 @@ func (tc *Catalog) DeleteCollectionVersion(ctx context.Context, req *coordinator
 	return &result, nil
 }
 
+func (tc *Catalog) BatchGetCollectionVersionFilePaths(ctx context.Context, req *coordinatorpb.BatchGetCollectionVersionFilePathsRequest) (*coordinatorpb.BatchGetCollectionVersionFilePathsResponse, error) {
+	result := coordinatorpb.BatchGetCollectionVersionFilePathsResponse{
+		CollectionIdToVersionFilePath: make(map[string]string),
+	}
+
+	paths, err := tc.metaDomain.CollectionDb(ctx).BatchGetCollectionVersionFilePaths(req.CollectionIds)
+	if err != nil {
+		return nil, err
+	}
+	result.CollectionIdToVersionFilePath = paths
+
+	return &result, nil
+}
+
 func (tc *Catalog) GetVersionFileNamesForCollection(ctx context.Context, tenantID string, collectionID string) (string, error) {
 	collectionIDPtr := &collectionID
 	collectionEntry, err := tc.metaDomain.CollectionDb(ctx).GetCollectionEntry(collectionIDPtr, nil)
diff --git a/go/pkg/sysdb/grpc/collection_service.go b/go/pkg/sysdb/grpc/collection_service.go
@@ -545,3 +545,11 @@ func (s *Server) DeleteCollectionVersion(ctx context.Context, req *coordinatorpb
 	}
 	return res, nil
 }
+
+func (s *Server) BatchGetCollectionVersionFilePaths(ctx context.Context, req *coordinatorpb.BatchGetCollectionVersionFilePathsRequest) (*coordinatorpb.BatchGetCollectionVersionFilePathsResponse, error) {
+	res, err := s.coordinator.BatchGetCollectionVersionFilePaths(ctx, req)
+	if err != nil {
+		return nil, grpcutils.BuildInternalGrpcError(err.Error())
+	}
+	return res, nil
+}
diff --git a/go/pkg/sysdb/metastore/db/dao/collection.go b/go/pkg/sysdb/metastore/db/dao/collection.go
@@ -527,3 +527,20 @@ func (s *collectionDb) UpdateCollectionLineageFilePath(collectionID string, curr
 		}).Error
 
 }
+
+func (s *collectionDb) BatchGetCollectionVersionFilePaths(collectionIDs []string) (map[string]string, error) {
+	var collections []dbmodel.Collection
+	err := s.read_db.Model(&dbmodel.Collection{}).
+		Select("id, version_file_name").
+		Where("id IN ?", collectionIDs).
+		Find(&collections).Error
+	if err != nil {
+		return nil, err
+	}
+
+	result := make(map[string]string)
+	for _, collection := range collections {
+		result[collection.ID] = collection.VersionFileName
+	}
+	return result, nil
+}
diff --git a/go/pkg/sysdb/metastore/db/dbmodel/collection.go b/go/pkg/sysdb/metastore/db/dbmodel/collection.go
@@ -68,4 +68,5 @@ type ICollectionDb interface {
 	UpdateVersionRelatedFields(collectionID, existingVersionFileName, newVersionFileName string, oldestVersionTs *time.Time, numActiveVersions *int) (int64, error)
 	LockCollection(collectionID string) error
 	UpdateCollectionLineageFilePath(collectionID string, currentLineageFilePath string, newLineageFilePath string) error
+	BatchGetCollectionVersionFilePaths(collectionIDs []string) (map[string]string, error)
 }
diff --git a/go/pkg/sysdb/metastore/db/dbmodel/mocks/ICollectionDb.go b/go/pkg/sysdb/metastore/db/dbmodel/mocks/ICollectionDb.go
diff --git a/idl/chromadb/proto/coordinator.proto b/idl/chromadb/proto/coordinator.proto
@@ -467,6 +467,14 @@ message DeleteCollectionVersionResponse {
   map<string, bool> collection_id_to_success = 1;
 }
 
+message BatchGetCollectionVersionFilePathsRequest {
+  repeated string collection_ids = 1;
+}
+
+message BatchGetCollectionVersionFilePathsResponse {
+  map<string, string> collection_id_to_version_file_path = 1;
+}
+
 service SysDB {
   rpc CreateDatabase(CreateDatabaseRequest) returns (CreateDatabaseResponse) {}
   rpc GetDatabase(GetDatabaseRequest) returns (GetDatabaseResponse) {}
@@ -498,4 +506,5 @@ service SysDB {
   rpc ListCollectionsToGc(ListCollectionsToGcRequest) returns (ListCollectionsToGcResponse) {}
   rpc MarkVersionForDeletion(MarkVersionForDeletionRequest) returns (MarkVersionForDeletionResponse) {}
   rpc DeleteCollectionVersion(DeleteCollectionVersionRequest) returns (DeleteCollectionVersionResponse) {}
+  rpc BatchGetCollectionVersionFilePaths(BatchGetCollectionVersionFilePathsRequest) returns (BatchGetCollectionVersionFilePathsResponse) {}
 }
diff --git a/rust/sysdb/src/sysdb.rs b/rust/sysdb/src/sysdb.rs
@@ -19,9 +19,9 @@ use chroma_types::{
     UpdateCollectionConfiguration, UpdateCollectionError, VectorIndexConfiguration,
 };
 use chroma_types::{
-    Collection, CollectionConversionError, CollectionUuid, CountForksError,
-    FlushCompactionResponse, FlushCompactionResponseConversionError, ForkCollectionError, Segment,
-    SegmentConversionError, SegmentScope, Tenant,
+    BatchGetCollectionVersionFilePathsError, Collection, CollectionConversionError, CollectionUuid,
+    CountForksError, FlushCompactionResponse, FlushCompactionResponseConversionError,
+    ForkCollectionError, Segment, SegmentConversionError, SegmentScope, Tenant,
 };
 use std::collections::HashMap;
 use std::fmt::Debug;
@@ -400,6 +400,20 @@ impl SysDb {
         }
     }
 
+    pub async fn batch_get_collection_version_file_paths(
+        &mut self,
+        collection_ids: Vec<CollectionUuid>,
+    ) -> Result<HashMap<CollectionUuid, String>, BatchGetCollectionVersionFilePathsError> {
+        match self {
+            SysDb::Grpc(grpc) => {
+                grpc.batch_get_collection_version_file_paths(collection_ids)
+                    .await
+            }
+            SysDb::Sqlite(_) => todo!(),
+            SysDb::Test(_) => todo!(),
+        }
+    }
+
     pub async fn get_last_compaction_time(
         &mut self,
         tanant_ids: Vec<String>,
@@ -1113,6 +1127,33 @@ impl GrpcSysDb {
         })
     }
 
+    async fn batch_get_collection_version_file_paths(
+        &mut self,
+        collection_ids: Vec<CollectionUuid>,
+    ) -> Result<HashMap<CollectionUuid, String>, BatchGetCollectionVersionFilePathsError> {
+        let res = self
+            .client
+            .batch_get_collection_version_file_paths(
+                chroma_proto::BatchGetCollectionVersionFilePathsRequest {
+                    collection_ids: collection_ids
+                        .into_iter()
+                        .map(|id| id.0.to_string())
+                        .collect(),
+                },
+            )
+            .await?;
+        let collection_id_to_path = res.into_inner().collection_id_to_version_file_path;
+        let mut result = HashMap::new();
+        for (key, value) in collection_id_to_path {
+            let collection_id = CollectionUuid(
+                Uuid::try_parse(&key)
+                    .map_err(|err| BatchGetCollectionVersionFilePathsError::Uuid(err, key))?,
+            );
+            result.insert(collection_id, value);
+        }
+        Ok(result)
+    }
+
     async fn get_last_compaction_time(
         &mut self,
         tenant_ids: Vec<String>,
diff --git a/rust/types/src/api_types.rs b/rust/types/src/api_types.rs
@@ -97,6 +97,23 @@ impl ChromaError for GetCollectionWithSegmentsError {
     }
 }
 
+#[derive(Debug, Error)]
+pub enum BatchGetCollectionVersionFilePathsError {
+    #[error("Grpc error: {0}")]
+    Grpc(#[from] Status),
+    #[error("Could not parse UUID from string {1}: {0}")]
+    Uuid(uuid::Error, String),
+}
+
+impl ChromaError for BatchGetCollectionVersionFilePathsError {
+    fn code(&self) -> ErrorCodes {
+        match self {
+            BatchGetCollectionVersionFilePathsError::Grpc(status) => status.code().into(),
+            BatchGetCollectionVersionFilePathsError::Uuid(_, _) => ErrorCodes::InvalidArgument,
+        }
+    }
+}
+
 #[derive(Serialize, ToSchema)]
 pub struct ResetResponse {}
 

Original file line number	Diff line number	Diff line change
`@@ -545,3 +545,11 @@ func (s Server) DeleteCollectionVersion(ctx context.Context, req coordinatorpb`
`545`	`545`	`}`
`546`	`546`	`return res, nil`
`547`	`547`	`}`
	`548`	`+`
	`549`	`+func (s Server) BatchGetCollectionVersionFilePaths(ctx context.Context, req coordinatorpb.BatchGetCollectionVersionFilePathsRequest) (*coordinatorpb.BatchGetCollectionVersionFilePathsResponse, error) {`
	`550`	`+ res, err := s.coordinator.BatchGetCollectionVersionFilePaths(ctx, req)`
	`551`	`+ if err != nil {`
	`552`	`+ return nil, grpcutils.BuildInternalGrpcError(err.Error())`
	`553`	`+ }`
	`554`	`+ return res, nil`
	`555`	`+}`
Original file line number	Diff line number	Diff line change
`@@ -68,4 +68,5 @@ type ICollectionDb interface {`
`68`	`68`	`UpdateVersionRelatedFields(collectionID, existingVersionFileName, newVersionFileName string, oldestVersionTs time.Time, numActiveVersions int) (int64, error)`
`69`	`69`	`LockCollection(collectionID string) error`
`70`	`70`	`UpdateCollectionLineageFilePath(collectionID string, currentLineageFilePath string, newLineageFilePath string) error`
	`71`	`+ BatchGetCollectionVersionFilePaths(collectionIDs []string) (map[string]string, error)`
`71`	`72`	`}`