Skip to content

Commit 0d136b9

Browse files
codethewebcodetheweb
committed
[ENH]: add batch get version file paths method to Sysdb
1 parent 3f2d8b1 commit 0d136b9

File tree

10 files changed

+176
-3
lines changed

10 files changed

+176
-3
lines changed

‎go/pkg/sysdb/coordinator/coordinator.go

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -263,6 +263,10 @@ func (s *Coordinator) DeleteCollectionVersion(ctx context.Context, req *coordina
263263
return s.catalog.DeleteCollectionVersion(ctx, req)
264264
}
265265

266+
func (s *Coordinator) BatchGetCollectionVersionFilePaths(ctx context.Context, req *coordinatorpb.BatchGetCollectionVersionFilePathsRequest) (*coordinatorpb.BatchGetCollectionVersionFilePathsResponse, error) {
267+
return s.catalog.BatchGetCollectionVersionFilePaths(ctx, req)
268+
}
269+
266270
// SetDeleteMode sets the delete mode for testing
267271
func (c *Coordinator) SetDeleteMode(mode DeleteMode) {
268272
c.deleteMode = mode

‎go/pkg/sysdb/coordinator/coordinator_test.go

Lines changed: 32 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -9,6 +9,7 @@ import (
99
"testing"
1010
"time"
1111

12+
"github.com/chroma-core/chroma/go/pkg/proto/coordinatorpb"
1213
"github.com/chroma-core/chroma/go/pkg/sysdb/metastore/db/dao"
1314
s3metastore "github.com/chroma-core/chroma/go/pkg/sysdb/metastore/s3"
1415
"github.com/pingcap/log"
@@ -1484,6 +1485,37 @@ func (suite *APIsTestSuite) TestForkCollection() {
14841485
suite.Empty(collections)
14851486
}
14861487

1488+
func (suite *APIsTestSuite) TestBatchGetCollectionVersionFilePaths() {
1489+
ctx := context.Background()
1490+
1491+
// Create a new collection
1492+
newCollection := &model.CreateCollection{
1493+
ID: types.NewUniqueID(),
1494+
Name: "test_batch_get_collection_version_file_paths",
1495+
TenantID: suite.tenantName,
1496+
DatabaseName: suite.databaseName,
1497+
}
1498+
1499+
newSegments := []*model.CreateSegment{}
1500+
1501+
// Create the collection
1502+
suite.coordinator.catalog.versionFileEnabled = true
1503+
_, _, err := suite.coordinator.CreateCollectionAndSegments(ctx, newCollection, newSegments)
1504+
suite.NoError(err)
1505+
1506+
// Get the version file paths for the collection
1507+
versionFilePaths, err := suite.coordinator.BatchGetCollectionVersionFilePaths(ctx, &coordinatorpb.BatchGetCollectionVersionFilePathsRequest{
1508+
CollectionIds: []string{newCollection.ID.String()},
1509+
})
1510+
suite.NoError(err)
1511+
suite.Len(versionFilePaths.CollectionIdToVersionFilePath, 1)
1512+
1513+
// Verify version file exists in S3
1514+
exists, err := suite.s3MetaStore.HasObjectWithPrefix(ctx, versionFilePaths.CollectionIdToVersionFilePath[newCollection.ID.String()])
1515+
suite.NoError(err)
1516+
suite.True(exists, "Version file should exist in S3")
1517+
}
1518+
14871519
func (suite *APIsTestSuite) TestCountForks() {
14881520
ctx := context.Background()
14891521

‎go/pkg/sysdb/coordinator/table_catalog.go

Lines changed: 14 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1983,6 +1983,20 @@ func (tc *Catalog) DeleteCollectionVersion(ctx context.Context, req *coordinator
19831983
return &result, nil
19841984
}
19851985

1986+
func (tc *Catalog) BatchGetCollectionVersionFilePaths(ctx context.Context, req *coordinatorpb.BatchGetCollectionVersionFilePathsRequest) (*coordinatorpb.BatchGetCollectionVersionFilePathsResponse, error) {
1987+
result := coordinatorpb.BatchGetCollectionVersionFilePathsResponse{
1988+
CollectionIdToVersionFilePath: make(map[string]string),
1989+
}
1990+
1991+
paths, err := tc.metaDomain.CollectionDb(ctx).BatchGetCollectionVersionFilePaths(req.CollectionIds)
1992+
if err != nil {
1993+
return nil, err
1994+
}
1995+
result.CollectionIdToVersionFilePath = paths
1996+
1997+
return &result, nil
1998+
}
1999+
19862000
func (tc *Catalog) GetVersionFileNamesForCollection(ctx context.Context, tenantID string, collectionID string) (string, error) {
19872001
collectionIDPtr := &collectionID
19882002
collectionEntry, err := tc.metaDomain.CollectionDb(ctx).GetCollectionEntry(collectionIDPtr, nil)

‎go/pkg/sysdb/grpc/collection_service.go

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -538,3 +538,11 @@ func (s *Server) DeleteCollectionVersion(ctx context.Context, req *coordinatorpb
538538
}
539539
return res, nil
540540
}
541+
542+
func (s *Server) BatchGetCollectionVersionFilePaths(ctx context.Context, req *coordinatorpb.BatchGetCollectionVersionFilePathsRequest) (*coordinatorpb.BatchGetCollectionVersionFilePathsResponse, error) {
543+
res, err := s.coordinator.BatchGetCollectionVersionFilePaths(ctx, req)
544+
if err != nil {
545+
return nil, grpcutils.BuildInternalGrpcError(err.Error())
546+
}
547+
return res, nil
548+
}

‎go/pkg/sysdb/metastore/db/dao/collection.go

Lines changed: 17 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -514,3 +514,20 @@ func (s *collectionDb) UpdateCollectionLineageFilePath(collectionID string, curr
514514
}).Error
515515

516516
}
517+
518+
func (s *collectionDb) BatchGetCollectionVersionFilePaths(collectionIDs []string) (map[string]string, error) {
519+
var collections []dbmodel.Collection
520+
err := s.read_db.Model(&dbmodel.Collection{}).
521+
Select("id, version_file_name").
522+
Where("id IN ?", collectionIDs).
523+
Find(&collections).Error
524+
if err != nil {
525+
return nil, err
526+
}
527+
528+
result := make(map[string]string)
529+
for _, collection := range collections {
530+
result[collection.ID] = collection.VersionFileName
531+
}
532+
return result, nil
533+
}

‎go/pkg/sysdb/metastore/db/dbmodel/collection.go

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -69,4 +69,5 @@ type ICollectionDb interface {
6969
UpdateVersionRelatedFields(collectionID, existingVersionFileName, newVersionFileName string, oldestVersionTs *time.Time, numActiveVersions *int) (int64, error)
7070
LockCollection(collectionID string) error
7171
UpdateCollectionLineageFilePath(collectionID string, currentLineageFilePath string, newLineageFilePath string) error
72+
BatchGetCollectionVersionFilePaths(collectionIDs []string) (map[string]string, error)
7273
}

‎go/pkg/sysdb/metastore/db/dbmodel/mocks/ICollectionDb.go

Lines changed: 30 additions & 0 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

‎idl/chromadb/proto/coordinator.proto

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -464,6 +464,14 @@ message DeleteCollectionVersionResponse {
464464
map<string, bool> collection_id_to_success = 1;
465465
}
466466

467+
message BatchGetCollectionVersionFilePathsRequest {
468+
repeated string collection_ids = 1;
469+
}
470+
471+
message BatchGetCollectionVersionFilePathsResponse {
472+
map<string, string> collection_id_to_version_file_path = 1;
473+
}
474+
467475
service SysDB {
468476
rpc CreateDatabase(CreateDatabaseRequest) returns (CreateDatabaseResponse) {}
469477
rpc GetDatabase(GetDatabaseRequest) returns (GetDatabaseResponse) {}
@@ -495,4 +503,5 @@ service SysDB {
495503
rpc ListCollectionsToGc(ListCollectionsToGcRequest) returns (ListCollectionsToGcResponse) {}
496504
rpc MarkVersionForDeletion(MarkVersionForDeletionRequest) returns (MarkVersionForDeletionResponse) {}
497505
rpc DeleteCollectionVersion(DeleteCollectionVersionRequest) returns (DeleteCollectionVersionResponse) {}
506+
rpc BatchGetCollectionVersionFilePaths(BatchGetCollectionVersionFilePathsRequest) returns (BatchGetCollectionVersionFilePathsResponse) {}
498507
}

‎rust/sysdb/src/sysdb.rs

Lines changed: 44 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -19,9 +19,9 @@ use chroma_types::{
1919
UpdateCollectionConfiguration, UpdateCollectionError, VectorIndexConfiguration,
2020
};
2121
use chroma_types::{
22-
Collection, CollectionConversionError, CollectionUuid, CountForksError,
23-
FlushCompactionResponse, FlushCompactionResponseConversionError, ForkCollectionError, Segment,
24-
SegmentConversionError, SegmentScope, Tenant,
22+
BatchGetCollectionVersionFilePathsError, Collection, CollectionConversionError, CollectionUuid,
23+
CountForksError, FlushCompactionResponse, FlushCompactionResponseConversionError,
24+
ForkCollectionError, Segment, SegmentConversionError, SegmentScope, Tenant,
2525
};
2626
use std::collections::HashMap;
2727
use std::fmt::Debug;
@@ -396,6 +396,20 @@ impl SysDb {
396396
}
397397
}
398398

399+
pub async fn batch_get_collection_version_file_paths(
400+
&mut self,
401+
collection_ids: Vec<CollectionUuid>,
402+
) -> Result<HashMap<CollectionUuid, String>, BatchGetCollectionVersionFilePathsError> {
403+
match self {
404+
SysDb::Grpc(grpc) => {
405+
grpc.batch_get_collection_version_file_paths(collection_ids)
406+
.await
407+
}
408+
SysDb::Sqlite(_) => todo!(),
409+
SysDb::Test(_) => todo!(),
410+
}
411+
}
412+
399413
pub async fn get_last_compaction_time(
400414
&mut self,
401415
tanant_ids: Vec<String>,
@@ -1107,6 +1121,33 @@ impl GrpcSysDb {
11071121
})
11081122
}
11091123

1124+
async fn batch_get_collection_version_file_paths(
1125+
&mut self,
1126+
collection_ids: Vec<CollectionUuid>,
1127+
) -> Result<HashMap<CollectionUuid, String>, BatchGetCollectionVersionFilePathsError> {
1128+
let res = self
1129+
.client
1130+
.batch_get_collection_version_file_paths(
1131+
chroma_proto::BatchGetCollectionVersionFilePathsRequest {
1132+
collection_ids: collection_ids
1133+
.into_iter()
1134+
.map(|id| id.0.to_string())
1135+
.collect(),
1136+
},
1137+
)
1138+
.await?;
1139+
let collection_id_to_path = res.into_inner().collection_id_to_version_file_path;
1140+
let mut result = HashMap::new();
1141+
for (key, value) in collection_id_to_path {
1142+
let collection_id = CollectionUuid(
1143+
Uuid::try_parse(&key)
1144+
.map_err(|err| BatchGetCollectionVersionFilePathsError::Uuid(err, key))?,
1145+
);
1146+
result.insert(collection_id, value);
1147+
}
1148+
Ok(result)
1149+
}
1150+
11101151
async fn get_last_compaction_time(
11111152
&mut self,
11121153
tenant_ids: Vec<String>,

‎rust/types/src/api_types.rs

Lines changed: 17 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -97,6 +97,23 @@ impl ChromaError for GetCollectionWithSegmentsError {
9797
}
9898
}
9999

100+
#[derive(Debug, Error)]
101+
pub enum BatchGetCollectionVersionFilePathsError {
102+
#[error("Grpc error: {0}")]
103+
Grpc(#[from] Status),
104+
#[error("Could not parse UUID from string {1}: {0}")]
105+
Uuid(uuid::Error, String),
106+
}
107+
108+
impl ChromaError for BatchGetCollectionVersionFilePathsError {
109+
fn code(&self) -> ErrorCodes {
110+
match self {
111+
BatchGetCollectionVersionFilePathsError::Grpc(status) => status.code().into(),
112+
BatchGetCollectionVersionFilePathsError::Uuid(_, _) => ErrorCodes::InvalidArgument,
113+
}
114+
}
115+
}
116+
100117
#[derive(Serialize, ToSchema)]
101118
pub struct ResetResponse {}
102119

0 commit comments

Comments
 (0)