Repository: chroma-core/chroma
Pull request #4585: [BUG]: Miscellanous fixes for forking
Author: login: sanketkedia
Created: 2025-05-20T20:36:20Z | Merged: 2025-05-21T00:16:47Z

## Description of changes

_Summarize the changes made by this PR._

- Improvements & Bug fixes
  - We never hard delete a root now
  - Forking succeeds even if root collection is soft deleted
  - Do not retry if fork returns a not found error
  - Correctly retry in case of retryable errors
- New functionality
  - ...

## Test plan

_How are these changes tested?_
Added tests
- [x] Tests pass locally with `pytest` for python, `yarn test` for js, `cargo test` for rust

## Documentation Changes
None



Diff:
diff --git a/go/pkg/common/errors.go b/go/pkg/common/errors.go
index acd1a5c48e2..951875d914d 100644
--- a/go/pkg/common/errors.go
+++ b/go/pkg/common/errors.go
@@ -28,6 +28,8 @@ var (
 	ErrCollectionVersionFileNameStale        = errors.New("collection version file name stale")
 	ErrCollectionEntryIsStale                = errors.New("collection entry is stale - one of version, version_file_name, or log_position is stale")
 	ErrCollectionTooManyFork                 = errors.New("collection entry has too many forks")
+	ErrCollectionDeletedWithLocksHeld        = errors.New("collection got deleted concurrently even though select for update locks were held. Not possible unless corruption somehow")
+	ErrMissingLineageFileName                = errors.New("missing lineage file name in root collection entry")
 
 	// Collection metadata errors
 	ErrUnknownCollectionMetadataType = errors.New("collection metadata value type not supported")
@@ -43,4 +45,7 @@ var (
 
 	// Segment metadata errors
 	ErrUnknownSegmentMetadataType = errors.New("segment metadata value type not supported")
+
+	// Others
+	ErrCompactionOffsetSomehowAhead = errors.New("system invariant was violated. Compaction offset in sysdb should always be behind or equal to offset in log")
 )
diff --git a/go/pkg/sysdb/coordinator/table_catalog.go b/go/pkg/sysdb/coordinator/table_catalog.go
index 45473cfc214..84a11025dcf 100644
--- a/go/pkg/sysdb/coordinator/table_catalog.go
+++ b/go/pkg/sysdb/coordinator/table_catalog.go
@@ -402,7 +402,7 @@ func (tc *Catalog) CheckCollection(ctx context.Context, collectionID types.Uniqu
 		defer span.End()
 	}
 
-	collectionInfo, err := tc.metaDomain.CollectionDb(ctx).GetCollectionEntry(types.FromUniqueID(collectionID), nil)
+	collectionInfo, err := tc.metaDomain.CollectionDb(ctx).GetCollectionWithoutMetadata(types.FromUniqueID(collectionID), nil, nil)
 	if err != nil {
 		return false, err
 	}
@@ -555,7 +555,7 @@ func (tc *Catalog) hardDeleteCollection(ctx context.Context, deleteCollection *m
 	return tc.txImpl.Transaction(ctx, func(txCtx context.Context) error {
 		collectionID := deleteCollection.ID
 
-		collectionEntry, err := tc.metaDomain.CollectionDb(txCtx).GetCollectionEntry(types.FromUniqueID(collectionID), &deleteCollection.DatabaseName)
+		collectionEntry, err := tc.metaDomain.CollectionDb(txCtx).GetCollectionWithoutMetadata(types.FromUniqueID(collectionID), &deleteCollection.DatabaseName, nil)
 		if err != nil {
 			return err
 		}
@@ -647,12 +647,13 @@ func (tc *Catalog) GetSoftDeletedCollections(ctx context.Context, collectionID *
 	collectionList := make([]*model.Collection, 0, len(collections))
 	for _, dbCollection := range collections {
 		collection := &model.Collection{
-			ID:           types.MustParse(dbCollection.Collection.ID),
-			Name:         *dbCollection.Collection.Name,
-			DatabaseName: dbCollection.DatabaseName,
-			TenantID:     dbCollection.TenantID,
-			Ts:           types.Timestamp(dbCollection.Collection.Ts),
-			UpdatedAt:    types.Timestamp(dbCollection.Collection.UpdatedAt.Unix()),
+			ID:              types.MustParse(dbCollection.Collection.ID),
+			Name:            *dbCollection.Collection.Name,
+			DatabaseName:    dbCollection.DatabaseName,
+			TenantID:        dbCollection.TenantID,
+			Ts:              types.Timestamp(dbCollection.Collection.Ts),
+			UpdatedAt:       types.Timestamp(dbCollection.Collection.UpdatedAt.Unix()),
+			LineageFileName: dbCollection.Collection.LineageFileName,
 		}
 		collectionList = append(collectionList, collection)
 	}
@@ -864,15 +865,16 @@ func (tc *Catalog) UpdateCollection(ctx context.Context, updateCollection *model
 	return result, nil
 }
 
-func (tc *Catalog) getLineageFile(ctx context.Context, collection *model.Collection) (*coordinatorpb.CollectionLineageFile, error) {
-	if collection.LineageFileName == nil {
+func (tc *Catalog) getLineageFile(ctx context.Context, lineageFileName *string) (*coordinatorpb.CollectionLineageFile, error) {
+	if lineageFileName == nil {
 		// There is no lineage file for the given collection
 		return &coordinatorpb.CollectionLineageFile{
 			Dependencies: []*coordinatorpb.CollectionVersionDependency{},
 		}, nil
 	}
 
-	return tc.s3Store.GetLineageFile(*collection.LineageFileName)
+	// Safe to deref.
+	return tc.s3Store.GetLineageFile(*lineageFileName)
 }
 
 func (tc *Catalog) ForkCollection(ctx context.Context, forkCollection *model.ForkCollection) (*model.Collection, []*model.Segment, error) {
@@ -880,23 +882,29 @@ func (tc *Catalog) ForkCollection(ctx context.Context, forkCollection *model.For
 
 	err := tc.txImpl.Transaction(ctx, func(txCtx context.Context) error {
 		var err error
-		var rootCollection *model.Collection
 		var rootCollectionID types.UniqueID
 		var rootCollectionIDStr string
 		var sourceCollection *model.Collection
 		var sourceSegments []*model.Segment
 		var newLineageFileFullName string
+		var oldLineageFileName *string
+		var lineageFileTenantId string
 
 		ts := time.Now().UTC()
 
 		sourceCollectionIDStr := forkCollection.SourceCollectionID.String()
 
 		// NOTE: We need to retrieve the source collection to get root collection id, then acquire locks on source and root collections in order to avoid deadlock.
-		// This step is because root collection id is always populated when the collection is created and is never modified.
-		sourceCollectionDb, err := tc.metaDomain.CollectionDb(txCtx).GetCollectionEntry(&sourceCollectionIDStr, nil)
+		// This step is safe because root collection id is always populated when the collection is created and is never modified.
+		// If source collection is deleted then cannot fork.
+		isDeleted := false
+		sourceCollectionDb, err := tc.metaDomain.CollectionDb(txCtx).GetCollectionWithoutMetadata(&sourceCollectionIDStr, nil, &isDeleted)
 		if err != nil {
 			return err
 		}
+		if sourceCollectionDb == nil {
+			return common.ErrCollectionNotFound
+		}
 
 		if sourceCollectionDb.RootCollectionId != nil {
 			rootCollectionID, err = types.Parse(*sourceCollectionDb.RootCollectionId)
@@ -915,29 +923,47 @@ func (tc *Catalog) ForkCollection(ctx context.Context, forkCollection *model.For
 			slices.Sort(collectionsToLock)
 		}
 		for _, collectionID := range collectionsToLock {
-			err = tc.metaDomain.CollectionDb(txCtx).LockCollection(collectionID)
-			if err != nil {
-				return err
+			isDeleted, e := tc.metaDomain.CollectionDb(txCtx).LockCollection(collectionID)
+			if e != nil {
+				return e
+			}
+			// It's ok for the root collection to be deleted but not for the source collection.
+			// We disable hard delete for soft deleted collections that are root.
+			if collectionID == sourceCollectionIDStr && *isDeleted {
+				return common.ErrCollectionNotFound
 			}
 		}
 
 		// Get source and root collections after they are locked
+		// They can't get deleted (both soft as well as hard) concurrently since we have the locks
 		sourceCollection, sourceSegments, err = tc.GetCollectionWithSegments(txCtx, forkCollection.SourceCollectionID, false)
 		if err != nil {
 			return err
 		}
+		// this can't happen and implies something weird in the system since
+		// we have the locks on the source collection
+		if sourceCollection == nil || sourceSegments == nil {
+			return common.ErrCollectionDeletedWithLocksHeld
+		}
 		if rootCollectionID != forkCollection.SourceCollectionID {
-			limit := int32(1)
-			collections, err := tc.GetCollections(txCtx, rootCollectionID, nil, "", "", &limit, nil)
+			// Get root collection. It is ok to get soft deleted root collection entry here.
+			collection, err := tc.metaDomain.CollectionDb(txCtx).GetCollectionWithoutMetadata(&rootCollectionIDStr, nil, nil)
 			if err != nil {
 				return err
 			}
-			if len(collections) == 0 {
-				return common.ErrCollectionNotFound
+			// Implies a hard deleted and should not happen
+			if collection == nil {
+				return common.ErrCollectionDeletedWithLocksHeld
 			}
-			rootCollection = collections[0]
+			// Root should always have a lineage file
+			if collection.LineageFileName == nil {
+				return common.ErrMissingLineageFileName
+			}
+			lineageFileTenantId = collection.Tenant
+			oldLineageFileName = collection.LineageFileName
 		} else {
-			rootCollection = sourceCollection
+			lineageFileTenantId = sourceCollection.TenantID
+			oldLineageFileName = sourceCollection.LineageFileName
 		}
 		databases, err := tc.metaDomain.DatabaseDb(txCtx).GetDatabases(sourceCollection.TenantID, sourceCollection.DatabaseName)
 		if err != nil {
@@ -956,9 +982,12 @@ func (tc *Catalog) ForkCollection(ctx context.Context, forkCollection *model.For
 		// t1: User writes to source collection, compaction takes place, source collection log offset become [400, 500]
 		// t2: Fork source collection in sysdb, the latest source collection compaction offset is 400. If we add new logs, it will start after offset 300, and the data is lost after compaction.
 		latestSourceCompactionOffset := uint64(sourceCollection.LogPosition)
-		if forkCollection.SourceCollectionLogEnumerationOffset < latestSourceCompactionOffset || latestSourceCompactionOffset < forkCollection.SourceCollectionLogCompactionOffset {
+		if forkCollection.SourceCollectionLogEnumerationOffset < latestSourceCompactionOffset {
 			return common.ErrCollectionLogPositionStale
 		}
+		if latestSourceCompactionOffset < forkCollection.SourceCollectionLogCompactionOffset {
+			return common.ErrCompactionOffsetSomehowAhead
+		}
 
 		// Create the new collection with source collection information
 		createCollection := &model.CreateCollection{
@@ -1010,7 +1039,7 @@ func (tc *Catalog) ForkCollection(ctx context.Context, forkCollection *model.For
 		}
 
 		// Update the lineage file
-		lineageFile, err := tc.getLineageFile(txCtx, rootCollection)
+		lineageFile, err := tc.getLineageFile(txCtx, oldLineageFileName)
 		if err != nil {
 			return err
 		}
@@ -1025,12 +1054,12 @@ func (tc *Catalog) ForkCollection(ctx context.Context, forkCollection *model.For
 		})
 
 		newLineageFileBaseName := fmt.Sprintf("%s/%d/%s.binpb", sourceCollectionIDStr, sourceCollection.Version, forkCollection.TargetCollectionID)
-		newLineageFileFullName, err = tc.s3Store.PutLineageFile(rootCollection.TenantID, databaseID, rootCollectionIDStr, newLineageFileBaseName, lineageFile)
+		newLineageFileFullName, err = tc.s3Store.PutLineageFile(lineageFileTenantId, databaseID, rootCollectionIDStr, newLineageFileBaseName, lineageFile)
 		if err != nil {
 			return err
 		}
 
-		return tc.metaDomain.CollectionDb(txCtx).UpdateCollectionLineageFilePath(rootCollectionIDStr, rootCollection.LineageFileName, newLineageFileFullName)
+		return tc.metaDomain.CollectionDb(txCtx).UpdateCollectionLineageFilePath(rootCollectionIDStr, oldLineageFileName, newLineageFileFullName)
 	})
 	if err != nil {
 		return nil, nil, err
@@ -1043,7 +1072,8 @@ func (tc *Catalog) CountForks(ctx context.Context, sourceCollectionID types.Uniq
 	var rootCollectionID types.UniqueID
 
 	sourceCollectionIDStr := sourceCollectionID.String()
-	sourceCollectionDb, err := tc.metaDomain.CollectionDb(ctx).GetCollectionEntry(&sourceCollectionIDStr, nil)
+	isDeleted := false
+	sourceCollectionDb, err := tc.metaDomain.CollectionDb(ctx).GetCollectionWithoutMetadata(&sourceCollectionIDStr, nil, &isDeleted)
 	if err != nil {
 		return 0, err
 	}
@@ -1070,7 +1100,7 @@ func (tc *Catalog) CountForks(ctx context.Context, sourceCollectionID types.Uniq
 	}
 	rootCollection := collections[0]
 
-	lineageFile, err := tc.getLineageFile(ctx, rootCollection)
+	lineageFile, err := tc.getLineageFile(ctx, rootCollection.LineageFileName)
 	if err != nil {
 		return 0, err
 	}
@@ -1398,7 +1428,8 @@ func (tc *Catalog) ListCollectionVersions(ctx context.Context,
 	includeMarkedForDeletion bool,
 ) ([]*coordinatorpb.CollectionVersionInfo, error) {
 	// Get collection entry to get version file name
-	collectionEntry, err := tc.metaDomain.CollectionDb(ctx).GetCollectionEntry(types.FromUniqueID(collectionID), nil)
+	isDeleted := false
+	collectionEntry, err := tc.metaDomain.CollectionDb(ctx).GetCollectionWithoutMetadata(types.FromUniqueID(collectionID), nil, &isDeleted)
 	if err != nil {
 		log.Error("error getting collection entry", zap.Error(err))
 		return nil, err
@@ -1521,7 +1552,7 @@ func (tc *Catalog) FlushCollectionCompaction(ctx context.Context, flushCollectio
 
 	err := tc.txImpl.Transaction(ctx, func(txCtx context.Context) error {
 		// Check if collection exists.
-		collection, err := tc.metaDomain.CollectionDb(txCtx).GetCollectionEntry(types.FromUniqueID(flushCollectionCompaction.ID), nil)
+		collection, err := tc.metaDomain.CollectionDb(txCtx).GetCollectionWithoutMetadata(types.FromUniqueID(flushCollectionCompaction.ID), nil, nil)
 		if err != nil {
 			return err
 		}
@@ -1859,7 +1890,8 @@ func (tc *Catalog) markVersionForDeletionInSingleCollection(
 
 		// Read the existing version file.
 		collectionIDPtr := &collectionID
-		collectionEntry, err := tc.metaDomain.CollectionDb(ctx).GetCollectionEntry(collectionIDPtr, nil)
+		isDeleted := false
+		collectionEntry, err := tc.metaDomain.CollectionDb(ctx).GetCollectionWithoutMetadata(collectionIDPtr, nil, &isDeleted)
 		if err != nil {
 			return err
 		}
@@ -1984,7 +2016,8 @@ func (tc *Catalog) DeleteVersionEntriesForCollection(ctx context.Context, tenant
 
 		// Read the existing version file
 		collectionIDPtr := &collectionID
-		collectionEntry, err := tc.metaDomain.CollectionDb(ctx).GetCollectionEntry(collectionIDPtr, nil)
+		isDeleted := false
+		collectionEntry, err := tc.metaDomain.CollectionDb(ctx).GetCollectionWithoutMetadata(collectionIDPtr, nil, &isDeleted)
 		if err != nil {
 			return err
 		}
@@ -2063,7 +2096,8 @@ func (tc *Catalog) DeleteCollectionVersion(ctx context.Context, req *coordinator
 
 func (tc *Catalog) GetVersionFileNamesForCollection(ctx context.Context, tenantID string, collectionID string) (string, error) {
 	collectionIDPtr := &collectionID
-	collectionEntry, err := tc.metaDomain.CollectionDb(ctx).GetCollectionEntry(collectionIDPtr, nil)
+	isDeleted := false
+	collectionEntry, err := tc.metaDomain.CollectionDb(ctx).GetCollectionWithoutMetadata(collectionIDPtr, nil, &isDeleted)
 	if err != nil {
 		return "", err
 	}
diff --git a/go/pkg/sysdb/coordinator/table_catalog_test.go b/go/pkg/sysdb/coordinator/table_catalog_test.go
index 60762651080..02a4624e1d9 100644
--- a/go/pkg/sysdb/coordinator/table_catalog_test.go
+++ b/go/pkg/sysdb/coordinator/table_catalog_test.go
@@ -680,7 +680,7 @@ func TestCatalog_DeleteCollectionVersion(t *testing.T) {
 
 	// Setup mock behaviors
 	mockMetaDomain.On("CollectionDb", mock.Anything).Return(mockCollectionDb)
-	mockCollectionDb.On("GetCollectionEntry", &collectionID, mock.Anything).Return(mockCollectionEntry, nil)
+	mockCollectionDb.On("GetCollectionWithoutMetadata", &collectionID, mock.Anything, mock.Anything).Return(mockCollectionEntry, nil)
 	mockCollectionDb.On("UpdateVersionRelatedFields",
 		collectionID,
 		existingVersionFileName,
@@ -742,7 +742,7 @@ func TestCatalog_DeleteCollectionVersion_CollectionNotFound(t *testing.T) {
 
 	// Setup mock behaviors
 	mockMetaDomain.On("CollectionDb", mock.Anything).Return(mockCollectionDb)
-	mockCollectionDb.On("GetCollectionEntry", &collectionID, mock.Anything).Return(nil, nil)
+	mockCollectionDb.On("GetCollectionWithoutMetadata", &collectionID, mock.Anything, mock.Anything).Return(nil, nil)
 
 	// Create test request
 	req := &coordinatorpb.DeleteCollectionVersionRequest{
@@ -811,7 +811,7 @@ func TestCatalog_MarkVersionForDeletion(t *testing.T) {
 
 	// Setup mock behaviors
 	mockMetaDomain.On("CollectionDb", mock.Anything).Return(mockCollectionDb)
-	mockCollectionDb.On("GetCollectionEntry", &collectionID, mock.Anything).Return(mockCollectionEntry, nil)
+	mockCollectionDb.On("GetCollectionWithoutMetadata", &collectionID, mock.Anything, mock.Anything).Return(mockCollectionEntry, nil)
 	mockCollectionDb.On("UpdateVersionRelatedFields",
 		collectionID,
 		existingVersionFileName,
@@ -878,7 +878,7 @@ func TestCatalog_MarkVersionForDeletion_CollectionNotFound(t *testing.T) {
 
 	// Setup mock behaviors
 	mockMetaDomain.On("CollectionDb", mock.Anything).Return(mockCollectionDb)
-	mockCollectionDb.On("GetCollectionEntry", &collectionID, mock.Anything).Return(nil, nil)
+	mockCollectionDb.On("GetCollectionWithoutMetadata", &collectionID, mock.Anything, mock.Anything).Return(nil, nil)
 
 	// Create test request
 	req := &coordinatorpb.MarkVersionForDeletionRequest{
@@ -946,7 +946,7 @@ func TestCatalog_MarkVersionForDeletion_VersionNotFound(t *testing.T) {
 
 	// Setup mock behaviors
 	mockMetaDomain.On("CollectionDb", mock.Anything).Return(mockCollectionDb)
-	mockCollectionDb.On("GetCollectionEntry", &collectionID, mock.Anything).Return(mockCollectionEntry, nil)
+	mockCollectionDb.On("GetCollectionWithoutMetadata", &collectionID, mock.Anything, mock.Anything).Return(mockCollectionEntry, nil)
 
 	// Create test request
 	req := &coordinatorpb.MarkVersionForDeletionRequest{
diff --git a/go/pkg/sysdb/grpc/cleanup.go b/go/pkg/sysdb/grpc/cleanup.go
index 2715e3dfd69..43a92f1e59c 100644
--- a/go/pkg/sysdb/grpc/cleanup.go
+++ b/go/pkg/sysdb/grpc/cleanup.go
@@ -58,6 +58,11 @@ func (s *SoftDeleteCleaner) run() {
 		}
 		numDeleted := 0
 		for _, collection := range collections {
+			// Skip root collections.
+			// Only roots have lineage file name set.
+			if collection.LineageFileName != nil && *collection.LineageFileName != "" {
+				continue
+			}
 			timeSinceDelete := time.Since(time.Unix(collection.UpdatedAt, 0))
 			if timeSinceDelete > s.maxAge {
 				log.Info("Deleting soft deleted collection", zap.String("collection_id", collection.ID.String()), zap.Duration("time_since_delete", timeSinceDelete), zap.Time("last_updated_at", time.Unix(collection.UpdatedAt, 0)))
diff --git a/go/pkg/sysdb/grpc/cleaup_test.go b/go/pkg/sysdb/grpc/cleaup_test.go
index e287b91af61..d5cd0fe88ca 100644
--- a/go/pkg/sysdb/grpc/cleaup_test.go
+++ b/go/pkg/sysdb/grpc/cleaup_test.go
@@ -73,7 +73,7 @@ func (suite *CleanupTestSuite) TestSoftDeleteCleanup() {
 	collections := make([]string, 2)
 	for i := 0; i < 2; i++ {
 		collectionName := "cleanup_test_collection_" + strconv.Itoa(i)
-		collectionID, err := dao.CreateTestCollection(suite.db, collectionName, 128, suite.databaseId)
+		collectionID, err := dao.CreateTestCollection(suite.db, collectionName, 128, suite.databaseId, nil)
 		suite.NoError(err)
 		collections[i] = collectionID
 	}
@@ -109,7 +109,7 @@ func (suite *CleanupTestSuite) TestSoftDeleteCleanup() {
 
 	// Create a test collection
 	collectionName := "cleanup_test_collection_double_delete"
-	collectionID, err := dao.CreateTestCollection(suite.db, collectionName, 128, suite.databaseId)
+	collectionID, err := dao.CreateTestCollection(suite.db, collectionName, 128, suite.databaseId, nil)
 	suite.NoError(err)
 
 	// Hard delete it once
@@ -136,6 +136,54 @@ func (suite *CleanupTestSuite) TestSoftDeleteCleanup() {
 
 }
 
+func (suite *CleanupTestSuite) TestSoftDeleteCleanupForkedCollection() {
+	// Create 3 test collections
+	collections := make([]string, 3)
+	print("Creating collection")
+	collectionName := "cleanup_root_test_collection"
+	lineageFileName := "lineageFileName"
+	collectionID, err := dao.CreateTestCollection(suite.db, collectionName, 128, suite.databaseId, &lineageFileName)
+	suite.NoError(err)
+	collections[0] = collectionID
+
+	for i := 1; i < 3; i++ {
+		collectionName := "cleanup_non_root_test_collection_" + strconv.Itoa(i)
+		collectionID, err := dao.CreateTestCollection(suite.db, collectionName, 128, suite.databaseId, nil)
+		suite.NoError(err)
+		collections[i] = collectionID
+	}
+
+	// Soft delete the collections
+	for _, collectionID := range collections {
+		err := suite.s.coordinator.DeleteCollection(context.Background(), &model.DeleteCollection{
+			ID: types.MustParse(collectionID),
+		})
+		suite.NoError(err)
+	}
+
+	// Verify collections are soft deleted
+	softDeletedCollections, err := suite.s.coordinator.GetSoftDeletedCollections(context.Background(), nil, "", "", 10)
+	suite.NoError(err)
+	suite.Equal(3, len(softDeletedCollections))
+
+	// Start the cleaner.
+	suite.s.softDeleteCleaner.maxInitialJitter = 0 * time.Second
+	suite.s.softDeleteCleaner.Start()
+
+	// Wait for cleanup cycle
+	time.Sleep(3 * time.Second)
+
+	// Verify root collection is not hard deleted.
+	softDeletedCollections, err = suite.s.coordinator.GetSoftDeletedCollections(context.Background(), nil, "", "", 10)
+	suite.NoError(err)
+	log.Info("softDeletedCollections", zap.Any("softDeletedCollections", softDeletedCollections))
+	suite.Equal(1, len(softDeletedCollections))
+	suite.Equal(softDeletedCollections[0].ID, types.MustParse(collections[0]))
+
+	// Stop the cleaner.
+	suite.s.softDeleteCleaner.Stop()
+}
+
 func TestCleanupTestSuite(t *testing.T) {
 	testSuite := new(CleanupTestSuite)
 	suite.Run(t, testSuite)
diff --git a/go/pkg/sysdb/grpc/collection_service.go b/go/pkg/sysdb/grpc/collection_service.go
index f47e91a8d2d..cc70960c390 100644
--- a/go/pkg/sysdb/grpc/collection_service.go
+++ b/go/pkg/sysdb/grpc/collection_service.go
@@ -382,7 +382,7 @@ func (s *Server) ForkCollection(ctx context.Context, req *coordinatorpb.ForkColl
 		if err == common.ErrCollectionNotFound {
 			return res, grpcutils.BuildNotFoundGrpcError(err.Error())
 		}
-		if err == common.ErrCollectionEntryIsStale {
+		if err == common.ErrCollectionLogPositionStale {
 			return res, grpcutils.BuildFailedPreconditionGrpcError(err.Error())
 		}
 		if err == common.ErrCollectionUniqueConstraintViolation {
diff --git a/go/pkg/sysdb/grpc/collection_service_test.go b/go/pkg/sysdb/grpc/collection_service_test.go
index bfbe6bbe297..74554d36b2b 100644
--- a/go/pkg/sysdb/grpc/collection_service_test.go
+++ b/go/pkg/sysdb/grpc/collection_service_test.go
@@ -348,7 +348,7 @@ func (suite *CollectionServiceTestSuite) TestCreateCollection() {
 func (suite *CollectionServiceTestSuite) TestServer_GetCollection() {
 	// Create a test collection
 	collectionName := "test_get_collection"
-	collectionID, err := dao.CreateTestCollection(suite.db, collectionName, 128, suite.databaseId)
+	collectionID, err := dao.CreateTestCollection(suite.db, collectionName, 128, suite.databaseId, nil)
 	suite.NoError(err)
 
 	// Enable soft delete mode
@@ -382,7 +382,7 @@ func (suite *CollectionServiceTestSuite) TestServer_FlushCollectionCompaction()
 	log.Info("TestServer_FlushCollectionCompaction")
 	// create test collection
 	collectionName := "collection_service_test_flush_collection_compaction"
-	collectionID, err := dao.CreateTestCollection(suite.db, collectionName, 128, suite.databaseId)
+	collectionID, err := dao.CreateTestCollection(suite.db, collectionName, 128, suite.databaseId, nil)
 	suite.NoError(err)
 
 	// flush collection compaction
@@ -533,7 +533,7 @@ func (suite *CollectionServiceTestSuite) TestServer_FlushCollectionCompaction()
 	// It should fail with a failed precondition error.
 	// Create collection and soft-delete it.
 	suite.s.coordinator.SetDeleteMode(coordinator.SoftDelete)
-	collectionID, err = dao.CreateTestCollection(suite.db, "test_flush_collection_compaction_soft_delete", 128, suite.databaseId)
+	collectionID, err = dao.CreateTestCollection(suite.db, "test_flush_collection_compaction_soft_delete", 128, suite.databaseId, nil)
 	suite.NoError(err)
 	suite.s.coordinator.DeleteCollection(context.Background(), &model.DeleteCollection{
 		ID:           types.MustParse(collectionID),
@@ -559,7 +559,7 @@ func (suite *CollectionServiceTestSuite) TestServer_FlushCollectionCompaction()
 
 func (suite *CollectionServiceTestSuite) TestGetCollectionSize() {
 	collectionName := "collection_service_test_get_collection_size"
-	collectionID, err := dao.CreateTestCollection(suite.db, collectionName, 128, suite.databaseId)
+	collectionID, err := dao.CreateTestCollection(suite.db, collectionName, 128, suite.databaseId, nil)
 	suite.NoError(err)
 
 	req := coordinatorpb.GetCollectionSizeRequest{
@@ -575,7 +575,7 @@ func (suite *CollectionServiceTestSuite) TestGetCollectionSize() {
 
 func (suite *CollectionServiceTestSuite) TestCountForks() {
 	collectionName := "collection_service_test_count_forks"
-	collectionID, err := dao.CreateTestCollection(suite.db, collectionName, 128, suite.databaseId)
+	collectionID, err := dao.CreateTestCollection(suite.db, collectionName, 128, suite.databaseId, nil)
 	suite.NoError(err)
 
 	req := coordinatorpb.CountForksRequest{
@@ -618,6 +618,88 @@ func (suite *CollectionServiceTestSuite) TestCountForks() {
 	suite.NoError(err)
 }
 
+func (suite *CollectionServiceTestSuite) TestFork() {
+	collectionName := "collection_service_test_forks"
+	collectionID, err := dao.CreateTestCollection(suite.db, collectionName, 128, suite.databaseId, nil)
+	suite.NoError(err)
+	targetCollectionID := types.NewUniqueID()
+
+	req := coordinatorpb.ForkCollectionRequest{
+		SourceCollectionId:                   collectionID,
+		SourceCollectionLogEnumerationOffset: 0,
+		SourceCollectionLogCompactionOffset:  0,
+		TargetCollectionId:                   targetCollectionID.String(),
+		TargetCollectionName:                 "test_fork_collection",
+	}
+	res, err := suite.s.ForkCollection(context.Background(), &req)
+	suite.NoError(err)
+	suite.Equal(res.Collection.Id, targetCollectionID.String())
+	suite.Equal(len(res.Segments), 2)
+
+	fork2CollectionId := types.NewUniqueID()
+	// Create fork of fork
+	forkCollectionReq := &coordinatorpb.ForkCollectionRequest{
+		SourceCollectionId:                   targetCollectionID.String(),
+		SourceCollectionLogCompactionOffset:  0,
+		SourceCollectionLogEnumerationOffset: 0,
+		TargetCollectionId:                   fork2CollectionId.String(),
+		TargetCollectionName:                 "test_fork_collection_fork",
+	}
+	forkedCollection2, err := suite.s.ForkCollection(context.Background(), forkCollectionReq)
+	suite.NoError(err)
+	suite.Equal(forkedCollection2.Collection.Id, fork2CollectionId.String())
+	suite.Equal(len(forkedCollection2.Segments), 2)
+
+	// Enable soft delete mode
+	suite.s.coordinator.SetDeleteMode(coordinator.SoftDelete)
+
+	// Delete the root.
+	deleteReq := model.DeleteCollection{
+		ID:           types.MustParse(collectionID),
+		TenantID:     suite.tenantName,
+		DatabaseName: suite.databaseName,
+		Ts:           time.Now().Unix(),
+	}
+	err = suite.s.coordinator.DeleteCollection(context.Background(), &deleteReq)
+	suite.NoError(err)
+
+	// Fork should still succeed.
+	fork3CollectionId := types.NewUniqueID()
+	fork3CollectionReq := &coordinatorpb.ForkCollectionRequest{
+		SourceCollectionId:                   fork2CollectionId.String(),
+		SourceCollectionLogCompactionOffset:  0,
+		SourceCollectionLogEnumerationOffset: 0,
+		TargetCollectionId:                   fork3CollectionId.String(),
+		TargetCollectionName:                 "test_fork_collection_fork_fork",
+	}
+	forkedCollection3, err := suite.s.ForkCollection(context.Background(), fork3CollectionReq)
+	suite.NoError(err)
+	suite.Equal(forkedCollection3.Collection.Id, fork3CollectionId.String())
+	suite.Equal(len(forkedCollection2.Segments), 2)
+
+	// Deleting the source and fork should not succeed.
+	deleteReq2 := model.DeleteCollection{
+		ID:           fork3CollectionId,
+		TenantID:     suite.tenantName,
+		DatabaseName: suite.databaseName,
+		Ts:           time.Now().Unix(),
+	}
+	err = suite.s.coordinator.DeleteCollection(context.Background(), &deleteReq2)
+	suite.NoError(err)
+
+	// Fork should not succeed.
+	fork4CollectionId := types.NewUniqueID()
+	fork4CollectionReq := &coordinatorpb.ForkCollectionRequest{
+		SourceCollectionId:                   fork3CollectionId.String(),
+		SourceCollectionLogCompactionOffset:  0,
+		SourceCollectionLogEnumerationOffset: 0,
+		TargetCollectionId:                   fork4CollectionId.String(),
+		TargetCollectionName:                 "test_fork_collection_fork_fork_fork",
+	}
+	_, err = suite.s.ForkCollection(context.Background(), fork4CollectionReq)
+	suite.Error(err)
+}
+
 func TestCollectionServiceTestSuite(t *testing.T) {
 	testSuite := new(CollectionServiceTestSuite)
 	suite.Run(t, testSuite)
diff --git a/go/pkg/sysdb/metastore/db/dao/collection.go b/go/pkg/sysdb/metastore/db/dao/collection.go
index 65bf6b2124e..c3de4fadaf7 100644
--- a/go/pkg/sysdb/metastore/db/dao/collection.go
+++ b/go/pkg/sysdb/metastore/db/dao/collection.go
@@ -27,7 +27,7 @@ func (s *collectionDb) DeleteAll() error {
 	return s.db.Where("1 = 1").Delete(&dbmodel.Collection{}).Error
 }
 
-func (s *collectionDb) GetCollectionEntry(collectionID *string, databaseName *string) (*dbmodel.Collection, error) {
+func (s *collectionDb) GetCollectionWithoutMetadata(collectionID *string, databaseName *string, softDeletedFlag *bool) (*dbmodel.Collection, error) {
 	var collections []*dbmodel.Collection
 	query := s.db.Table("collections").
 		Select("collections.id, collections.name, collections.database_id, collections.is_deleted, collections.tenant, collections.version, collections.version_file_name, collections.root_collection_id, NULLIF(collections.lineage_file_name, '') AS lineage_file_name").
@@ -38,6 +38,10 @@ func (s *collectionDb) GetCollectionEntry(collectionID *string, databaseName *st
 		query = query.Where("databases.name = ?", databaseName)
 	}
 
+	if softDeletedFlag != nil {
+		query = query.Where("collections.is_deleted = ?", *softDeletedFlag)
+	}
+
 	err := query.Find(&collections).Error
 	if err != nil {
 		return nil, err
@@ -53,8 +57,8 @@ func (s *collectionDb) GetCollectionEntries(id *string, name *string, tenantID s
 }
 
 func (s *collectionDb) GetCollections(id *string, name *string, tenantID string, databaseName string, limit *int32, offset *int32) ([]*dbmodel.CollectionAndMetadata, error) {
-	is_deleted := false
-	return s.getCollections(id, name, tenantID, databaseName, limit, offset, &is_deleted)
+	isDeleted := false
+	return s.getCollections(id, name, tenantID, databaseName, limit, offset, &isDeleted)
 }
 
 func (s *collectionDb) ListCollectionsToGc(cutoffTimeSecs *uint64, limit *uint64, tenantID *string) ([]*dbmodel.CollectionToGc, error) {
@@ -322,8 +326,8 @@ func (s *collectionDb) GetCollectionSize(id string) (uint64, error) {
 }
 
 func (s *collectionDb) GetSoftDeletedCollections(collectionID *string, tenantID string, databaseName string, limit int32) ([]*dbmodel.CollectionAndMetadata, error) {
-	is_deleted := true
-	return s.getCollections(collectionID, nil, tenantID, databaseName, &limit, nil, &is_deleted)
+	isDeleted := true
+	return s.getCollections(collectionID, nil, tenantID, databaseName, &limit, nil, &isDeleted)
 }
 
 // NOTE: This is the only method to do a hard delete of a single collection.
@@ -500,17 +504,17 @@ func (s *collectionDb) UpdateVersionRelatedFields(collectionID, existingVersionF
 	return result.RowsAffected, nil
 }
 
-func (s *collectionDb) LockCollection(collectionID string) error {
+func (s *collectionDb) LockCollection(collectionID string) (*bool, error) {
 	var collections []dbmodel.Collection
 	err := s.db.Model(&dbmodel.Collection{}).
 		Where("collections.id = ?", collectionID).Clauses(clause.Locking{
 		Strength: "UPDATE",
 	}).Find(&collections).Error
 	if err != nil {
-		return err
+		return nil, err
 	}
 	if len(collections) == 0 {
-		return common.ErrCollectionNotFound
+		return nil, common.ErrCollectionNotFound
 	}
 
 	err = s.db.Model(&dbmodel.CollectionMetadata{}).
@@ -518,7 +522,7 @@ func (s *collectionDb) LockCollection(collectionID string) error {
 		Strength: "UPDATE",
 	}).Find(nil).Error
 	if err != nil {
-		return err
+		return nil, err
 	}
 
 	var segments []*dbmodel.Segment
@@ -527,7 +531,7 @@ func (s *collectionDb) LockCollection(collectionID string) error {
 		Strength: "UPDATE",
 	}).Find(&segments).Error
 	if err != nil {
-		return err
+		return nil, err
 	}
 
 	var segmentIDs []*string
@@ -540,10 +544,10 @@ func (s *collectionDb) LockCollection(collectionID string) error {
 		Strength: "UPDATE",
 	}).Find(nil).Error
 	if err != nil {
-		return err
+		return nil, err
 	}
 
-	return nil
+	return &collections[0].IsDeleted, nil
 }
 
 func (s *collectionDb) UpdateCollectionLineageFilePath(collectionID string, currentLineageFileName *string, newLineageFileName string) error {
diff --git a/go/pkg/sysdb/metastore/db/dao/collection_test.go b/go/pkg/sysdb/metastore/db/dao/collection_test.go
index 8827a365ea8..ca39573bacc 100644
--- a/go/pkg/sysdb/metastore/db/dao/collection_test.go
+++ b/go/pkg/sysdb/metastore/db/dao/collection_test.go
@@ -48,7 +48,7 @@ func (suite *CollectionDbTestSuite) TearDownSuite() {
 func (suite *CollectionDbTestSuite) TestCollectionDb_GetCollections() {
 	collectionName := "test_collection_get_collections"
 	dim := int32(128)
-	collectionID, err := CreateTestCollection(suite.db, collectionName, dim, suite.databaseId)
+	collectionID, err := CreateTestCollection(suite.db, collectionName, dim, suite.databaseId, nil)
 	suite.NoError(err)
 
 	testKey := "test"
@@ -103,7 +103,7 @@ func (suite *CollectionDbTestSuite) TestCollectionDb_GetCollections() {
 	suite.Len(collections, 1)
 	suite.Equal(collectionID, collections[0].Collection.ID)
 
-	collectionID2, err := CreateTestCollection(suite.db, "test_collection_get_collections2", 128, suite.databaseId)
+	collectionID2, err := CreateTestCollection(suite.db, "test_collection_get_collections2", 128, suite.databaseId, nil)
 	suite.NoError(err)
 
 	// Test order by. Collections are ordered by create time so collectionID2 should be second
@@ -137,10 +137,10 @@ func (suite *CollectionDbTestSuite) TestCollectionDb_GetCollections() {
 	suite.NoError(err)
 
 	// Create two collections in the new database.
-	collectionID3, err := CreateTestCollection(suite.db, "test_collection_get_collections3", 128, DbId)
+	collectionID3, err := CreateTestCollection(suite.db, "test_collection_get_collections3", 128, DbId, nil)
 	suite.NoError(err)
 
-	collectionID4, err := CreateTestCollection(suite.db, "test_collection_get_collections4", 128, DbId)
+	collectionID4, err := CreateTestCollection(suite.db, "test_collection_get_collections4", 128, DbId, nil)
 	suite.NoError(err)
 
 	// Test count collections
@@ -174,7 +174,7 @@ func (suite *CollectionDbTestSuite) TestCollectionDb_GetCollections() {
 
 func (suite *CollectionDbTestSuite) TestCollectionDb_UpdateLogPositionVersionTotalRecordsAndLogicalSize() {
 	collectionName := "test_collection_get_collections"
-	collectionID, _ := CreateTestCollection(suite.db, collectionName, 128, suite.databaseId)
+	collectionID, _ := CreateTestCollection(suite.db, collectionName, 128, suite.databaseId, nil)
 	// verify default values
 	collections, err := suite.collectionDb.GetCollections(&collectionID, nil, "", "", nil, nil)
 	suite.NoError(err)
@@ -227,9 +227,9 @@ func (suite *CollectionDbTestSuite) TestCollectionDb_SoftDelete() {
 	// Create 2 collections.
 	collectionName1 := "test_collection_soft_delete1"
 	collectionName2 := "test_collection_soft_delete2"
-	collectionID1, err := CreateTestCollection(suite.db, collectionName1, 128, suite.databaseId)
+	collectionID1, err := CreateTestCollection(suite.db, collectionName1, 128, suite.databaseId, nil)
 	suite.NoError(err)
-	collectionID2, err := CreateTestCollection(suite.db, collectionName2, 128, suite.databaseId)
+	collectionID2, err := CreateTestCollection(suite.db, collectionName2, 128, suite.databaseId, nil)
 	suite.NoError(err)
 
 	// Soft delete collection 1 by Updating the is_deleted column
@@ -264,7 +264,7 @@ func (suite *CollectionDbTestSuite) TestCollectionDb_SoftDelete() {
 
 func (suite *CollectionDbTestSuite) TestCollectionDb_GetCollectionSize() {
 	collectionName := "test_collection_get_collection_size"
-	collectionID, err := CreateTestCollection(suite.db, collectionName, 128, suite.databaseId)
+	collectionID, err := CreateTestCollection(suite.db, collectionName, 128, suite.databaseId, nil)
 	suite.NoError(err)
 
 	total_records_post_compaction, err := suite.collectionDb.GetCollectionSize(collectionID)
diff --git a/go/pkg/sysdb/metastore/db/dao/segment_test.go b/go/pkg/sysdb/metastore/db/dao/segment_test.go
index e19d52e80fd..3fff87f9b47 100644
--- a/go/pkg/sysdb/metastore/db/dao/segment_test.go
+++ b/go/pkg/sysdb/metastore/db/dao/segment_test.go
@@ -84,7 +84,7 @@ func (suite *SegmentDbTestSuite) TestSegmentDb_RegisterFilePath() {
 	// create a collection for testing
 	databaseId := types.NewUniqueID().String()
 	collectionName := "test_segment_register_file_paths"
-	collectionID, err := CreateTestCollection(suite.db, collectionName, 128, databaseId)
+	collectionID, err := CreateTestCollection(suite.db, collectionName, 128, databaseId, nil)
 	suite.NoError(err)
 
 	segments, err := suite.segmentDb.GetSegments(types.NilUniqueID(), nil, nil, types.MustParse(collectionID))
diff --git a/go/pkg/sysdb/metastore/db/dao/test_utils.go b/go/pkg/sysdb/metastore/db/dao/test_utils.go
index 580702020aa..4aeeffe0e46 100644
--- a/go/pkg/sysdb/metastore/db/dao/test_utils.go
+++ b/go/pkg/sysdb/metastore/db/dao/test_utils.go
@@ -125,7 +125,7 @@ func CleanUpTestTenant(db *gorm.DB, tenantName string) error {
 	return nil
 }
 
-func CreateTestCollection(db *gorm.DB, collectionName string, dimension int32, databaseID string) (string, error) {
+func CreateTestCollection(db *gorm.DB, collectionName string, dimension int32, databaseID string, lineageFileName *string) (string, error) {
 	log.Info("create test collection", zap.String("collectionName", collectionName), zap.Int32("dimension", dimension), zap.String("databaseID", databaseID))
 	collectionDb := &collectionDb{
 		db: db,
@@ -147,6 +147,7 @@ func CreateTestCollection(db *gorm.DB, collectionName string, dimension int32, d
 		SizeBytesPostCompaction:    uint64(500000),
 		LastCompactionTimeSecs:     uint64(1741037006),
 		Tenant:                     "test_tenant",
+		LineageFileName:            lineageFileName,
 	})
 	if err != nil {
 		return "", err
diff --git a/go/pkg/sysdb/metastore/db/dbmodel/collection.go b/go/pkg/sysdb/metastore/db/dbmodel/collection.go
index be12580cd00..6ac5d8221a1 100644
--- a/go/pkg/sysdb/metastore/db/dbmodel/collection.go
+++ b/go/pkg/sysdb/metastore/db/dbmodel/collection.go
@@ -63,10 +63,10 @@ type ICollectionDb interface {
 	UpdateLogPositionVersionTotalRecordsAndLogicalSize(collectionID string, logPosition int64, currentCollectionVersion int32, totalRecordsPostCompaction uint64, sizeBytesPostCompaction uint64, lastCompactionTimeSecs uint64, tenant string) (int32, error)
 	UpdateLogPositionAndVersionInfo(collectionID string, logPosition int64, currentCollectionVersion int32, currentVersionFilePath string, newCollectionVersion int32, newVersionFilePath string, totalRecordsPostCompaction uint64,
 		sizeBytesPostCompaction uint64, lastCompactionTimeSecs uint64) (int64, error)
-	GetCollectionEntry(collectionID *string, databaseName *string) (*Collection, error)
+	GetCollectionWithoutMetadata(collectionID *string, databaseName *string, softDeletedFlag *bool) (*Collection, error)
 	GetCollectionSize(collectionID string) (uint64, error)
 	ListCollectionsToGc(cutoffTimeSecs *uint64, limit *uint64, tenantID *string) ([]*CollectionToGc, error)
 	UpdateVersionRelatedFields(collectionID, existingVersionFileName, newVersionFileName string, oldestVersionTs *time.Time, numActiveVersions *int) (int64, error)
-	LockCollection(collectionID string) error
+	LockCollection(collectionID string) (*bool, error)
 	UpdateCollectionLineageFilePath(collectionID string, currentLineageFilePath *string, newLineageFilePath string) error
 }
diff --git a/go/pkg/sysdb/metastore/db/dbmodel/mocks/ICollectionDb.go b/go/pkg/sysdb/metastore/db/dbmodel/mocks/ICollectionDb.go
index b1a8842119c..c5ac7bb69aa 100644
--- a/go/pkg/sysdb/metastore/db/dbmodel/mocks/ICollectionDb.go
+++ b/go/pkg/sysdb/metastore/db/dbmodel/mocks/ICollectionDb.go
@@ -1,4 +1,4 @@
-// Code generated by mockery v2.42.1. DO NOT EDIT.
+// Code generated by mockery v2.46.2. DO NOT EDIT.
 
 package mocks
 
@@ -118,29 +118,27 @@ func (_m *ICollectionDb) GetCollectionEntries(id *string, name *string, tenantID
 	return r0, r1
 }
 
-// GetCollectionEntry provides a mock function with given fields: collectionID, databaseName
-func (_m *ICollectionDb) GetCollectionEntry(collectionID *string, databaseName *string) (*dbmodel.Collection, error) {
-	ret := _m.Called(collectionID, databaseName)
+// GetCollectionSize provides a mock function with given fields: collectionID
+func (_m *ICollectionDb) GetCollectionSize(collectionID string) (uint64, error) {
+	ret := _m.Called(collectionID)
 
 	if len(ret) == 0 {
-		panic("no return value specified for GetCollectionEntry")
+		panic("no return value specified for GetCollectionSize")
 	}
 
-	var r0 *dbmodel.Collection
+	var r0 uint64
 	var r1 error
-	if rf, ok := ret.Get(0).(func(*string, *string) (*dbmodel.Collection, error)); ok {
-		return rf(collectionID, databaseName)
+	if rf, ok := ret.Get(0).(func(string) (uint64, error)); ok {
+		return rf(collectionID)
 	}
-	if rf, ok := ret.Get(0).(func(*string, *string) *dbmodel.Collection); ok {
-		r0 = rf(collectionID, databaseName)
+	if rf, ok := ret.Get(0).(func(string) uint64); ok {
+		r0 = rf(collectionID)
 	} else {
-		if ret.Get(0) != nil {
-			r0 = ret.Get(0).(*dbmodel.Collection)
-		}
+		r0 = ret.Get(0).(uint64)
 	}
 
-	if rf, ok := ret.Get(1).(func(*string, *string) error); ok {
-		r1 = rf(collectionID, databaseName)
+	if rf, ok := ret.Get(1).(func(string) error); ok {
+		r1 = rf(collectionID)
 	} else {
 		r1 = ret.Error(1)
 	}
@@ -148,27 +146,29 @@ func (_m *ICollectionDb) GetCollectionEntry(collectionID *string, databaseName *
 	return r0, r1
 }
 
-// GetCollectionSize provides a mock function with given fields: collectionID
-func (_m *ICollectionDb) GetCollectionSize(collectionID string) (uint64, error) {
-	ret := _m.Called(collectionID)
+// GetCollectionWithoutMetadata provides a mock function with given fields: collectionID, databaseName, softDeletedFlag
+func (_m *ICollectionDb) GetCollectionWithoutMetadata(collectionID *string, databaseName *string, softDeletedFlag *bool) (*dbmodel.Collection, error) {
+	ret := _m.Called(collectionID, databaseName, softDeletedFlag)
 
 	if len(ret) == 0 {
-		panic("no return value specified for GetCollectionSize")
+		panic("no return value specified for GetCollectionWithoutMetadata")
 	}
 
-	var r0 uint64
+	var r0 *dbmodel.Collection
 	var r1 error
-	if rf, ok := ret.Get(0).(func(string) (uint64, error)); ok {
-		return rf(collectionID)
+	if rf, ok := ret.Get(0).(func(*string, *string, *bool) (*dbmodel.Collection, error)); ok {
+		return rf(collectionID, databaseName, softDeletedFlag)
 	}
-	if rf, ok := ret.Get(0).(func(string) uint64); ok {
-		r0 = rf(collectionID)
+	if rf, ok := ret.Get(0).(func(*string, *string, *bool) *dbmodel.Collection); ok {
+		r0 = rf(collectionID, databaseName, softDeletedFlag)
 	} else {
-		r0 = ret.Get(0).(uint64)
+		if ret.Get(0) != nil {
+			r0 = ret.Get(0).(*dbmodel.Collection)
+		}
 	}
 
-	if rf, ok := ret.Get(1).(func(string) error); ok {
-		r1 = rf(collectionID)
+	if rf, ok := ret.Get(1).(func(*string, *string, *bool) error); ok {
+		r1 = rf(collectionID, databaseName, softDeletedFlag)
 	} else {
 		r1 = ret.Error(1)
 	}
@@ -313,21 +313,33 @@ func (_m *ICollectionDb) ListCollectionsToGc(cutoffTimeSecs *uint64, limit *uint
 }
 
 // LockCollection provides a mock function with given fields: collectionID
-func (_m *ICollectionDb) LockCollection(collectionID string) error {
+func (_m *ICollectionDb) LockCollection(collectionID string) (*bool, error) {
 	ret := _m.Called(collectionID)
 
 	if len(ret) == 0 {
 		panic("no return value specified for LockCollection")
 	}
 
-	var r0 error
-	if rf, ok := ret.Get(0).(func(string) error); ok {
+	var r0 *bool
+	var r1 error
+	if rf, ok := ret.Get(0).(func(string) (*bool, error)); ok {
+		return rf(collectionID)
+	}
+	if rf, ok := ret.Get(0).(func(string) *bool); ok {
 		r0 = rf(collectionID)
 	} else {
-		r0 = ret.Error(0)
+		if ret.Get(0) != nil {
+			r0 = ret.Get(0).(*bool)
+		}
 	}
 
-	return r0
+	if rf, ok := ret.Get(1).(func(string) error); ok {
+		r1 = rf(collectionID)
+	} else {
+		r1 = ret.Error(1)
+	}
+
+	return r0, r1
 }
 
 // Update provides a mock function with given fields: in
diff --git a/rust/frontend/src/impls/service_based_frontend.rs b/rust/frontend/src/impls/service_based_frontend.rs
index 0d8f541c367..38109f2f843 100644
--- a/rust/frontend/src/impls/service_based_frontend.rs
+++ b/rust/frontend/src/impls/service_based_frontend.rs
@@ -634,7 +634,7 @@ impl ServiceBasedFrontend {
             .when(|e| {
                 matches!(
                     e.code(),
-                    ErrorCodes::FailedPrecondition | ErrorCodes::NotFound | ErrorCodes::Unknown
+                    ErrorCodes::FailedPrecondition | ErrorCodes::Unknown
                 )
             })
             .notify(|_, _| {


Reviews:
user: login: Sicheng-Pan
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-20T21:45:21Z
user: login: Sicheng-Pan
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-20T22:10:27Z
user: login: sanketkedia
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-20T22:11:17Z
user: login: sanketkedia
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-20T22:12:12Z
user: login: Sicheng-Pan
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-20T22:42:05Z
user: login: Sicheng-Pan
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-20T22:46:18Z
user: login: sanketkedia
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-20T23:17:43Z
user: login: sanketkedia
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-20T23:18:18Z
user: login: Sicheng-Pan
state: APPROVED
author_association: CONTRIBUTOR
submitted_at: 2025-05-20T23:27:45Z

Review comments:
diff_hunk: @@ -915,29 +922,45 @@ func (tc *Catalog) ForkCollection(ctx context.Context, forkCollection *model.For
 			slices.Sort(collectionsToLock)
 		}
 		for _, collectionID := range collectionsToLock {
-			err = tc.metaDomain.CollectionDb(txCtx).LockCollection(collectionID)
-			if err != nil {
+			isDeleted, e := tc.metaDomain.CollectionDb(txCtx).LockCollection(collectionID)
+			if e != nil {
 				return err
 			}
path: go/pkg/sysdb/coordinator/table_catalog.go
user: login: Sicheng-Pan
body: why return `err` if `e` is not null?
created_at: 2025-05-20T21:45:21Z
author_association: CONTRIBUTOR
diff_hunk: @@ -1025,12 +1048,12 @@ func (tc *Catalog) ForkCollection(ctx context.Context, forkCollection *model.For
 		})
 
 		newLineageFileBaseName := fmt.Sprintf("%s/%d/%s.binpb", sourceCollectionIDStr, sourceCollection.Version, forkCollection.TargetCollectionID)
-		newLineageFileFullName, err = tc.s3Store.PutLineageFile(rootCollection.TenantID, databaseID, rootCollectionIDStr, newLineageFileBaseName, lineageFile)
+		newLineageFileFullName, err = tc.s3Store.PutLineageFile(sourceCollection.TenantID, databaseID, rootCollectionIDStr, newLineageFileBaseName, lineageFile)
path: go/pkg/sysdb/coordinator/table_catalog.go
user: login: Sicheng-Pan
body: Shouldn't we always put the lineage file under the root collection, even if it is deleted?
created_at: 2025-05-20T22:10:27Z
author_association: CONTRIBUTOR
diff_hunk: @@ -915,29 +922,45 @@ func (tc *Catalog) ForkCollection(ctx context.Context, forkCollection *model.For
 			slices.Sort(collectionsToLock)
 		}
 		for _, collectionID := range collectionsToLock {
-			err = tc.metaDomain.CollectionDb(txCtx).LockCollection(collectionID)
-			if err != nil {
+			isDeleted, e := tc.metaDomain.CollectionDb(txCtx).LockCollection(collectionID)
+			if e != nil {
 				return err
 			}
path: go/pkg/sysdb/coordinator/table_catalog.go
user: login: sanketkedia
body: copy pasta, thanks for correcting
created_at: 2025-05-20T22:11:17Z
author_association: CONTRIBUTOR
diff_hunk: @@ -1025,12 +1048,12 @@ func (tc *Catalog) ForkCollection(ctx context.Context, forkCollection *model.For
 		})
 
 		newLineageFileBaseName := fmt.Sprintf("%s/%d/%s.binpb", sourceCollectionIDStr, sourceCollection.Version, forkCollection.TargetCollectionID)
-		newLineageFileFullName, err = tc.s3Store.PutLineageFile(rootCollection.TenantID, databaseID, rootCollectionIDStr, newLineageFileBaseName, lineageFile)
+		newLineageFileFullName, err = tc.s3Store.PutLineageFile(sourceCollection.TenantID, databaseID, rootCollectionIDStr, newLineageFileBaseName, lineageFile)
path: go/pkg/sysdb/coordinator/table_catalog.go
user: login: sanketkedia
body: why will the tenant ids change?
created_at: 2025-05-20T22:12:12Z
author_association: CONTRIBUTOR
diff_hunk: @@ -1025,12 +1048,12 @@ func (tc *Catalog) ForkCollection(ctx context.Context, forkCollection *model.For
 		})
 
 		newLineageFileBaseName := fmt.Sprintf("%s/%d/%s.binpb", sourceCollectionIDStr, sourceCollection.Version, forkCollection.TargetCollectionID)
-		newLineageFileFullName, err = tc.s3Store.PutLineageFile(rootCollection.TenantID, databaseID, rootCollectionIDStr, newLineageFileBaseName, lineageFile)
+		newLineageFileFullName, err = tc.s3Store.PutLineageFile(sourceCollection.TenantID, databaseID, rootCollectionIDStr, newLineageFileBaseName, lineageFile)
path: go/pkg/sysdb/coordinator/table_catalog.go
user: login: Sicheng-Pan
body: Not sure if we will support collection ownership transfer in the future, and if that is the case the source collection tenant might not be the same as the root collection tenant
created_at: 2025-05-20T22:42:05Z
author_association: CONTRIBUTOR
diff_hunk: @@ -27,7 +27,7 @@ func (s *collectionDb) DeleteAll() error {
 	return s.db.Where("1 = 1").Delete(&dbmodel.Collection{}).Error
 }
 
-func (s *collectionDb) GetCollectionEntry(collectionID *string, databaseName *string) (*dbmodel.Collection, error) {
+func (s *collectionDb) GetCollectionWithoutMetadata(collectionID *string, databaseName *string, softDeletedFlag *bool) (*dbmodel.Collection, error) {
path: go/pkg/sysdb/metastore/db/dao/collection.go
user: login: Sicheng-Pan
body: is there a use case for `softDeletedFlag = &true`? if not maybe it's cleaner to replace it with `includeSoftDeletedCollection bool`
created_at: 2025-05-20T22:46:18Z
author_association: CONTRIBUTOR
diff_hunk: @@ -1025,12 +1048,12 @@ func (tc *Catalog) ForkCollection(ctx context.Context, forkCollection *model.For
 		})
 
 		newLineageFileBaseName := fmt.Sprintf("%s/%d/%s.binpb", sourceCollectionIDStr, sourceCollection.Version, forkCollection.TargetCollectionID)
-		newLineageFileFullName, err = tc.s3Store.PutLineageFile(rootCollection.TenantID, databaseID, rootCollectionIDStr, newLineageFileBaseName, lineageFile)
+		newLineageFileFullName, err = tc.s3Store.PutLineageFile(sourceCollection.TenantID, databaseID, rootCollectionIDStr, newLineageFileBaseName, lineageFile)
path: go/pkg/sysdb/coordinator/table_catalog.go
user: login: sanketkedia
body: will change thanks but I don't think we want to support that in future.
created_at: 2025-05-20T23:17:43Z
author_association: CONTRIBUTOR
diff_hunk: @@ -27,7 +27,7 @@ func (s *collectionDb) DeleteAll() error {
 	return s.db.Where("1 = 1").Delete(&dbmodel.Collection{}).Error
 }
 
-func (s *collectionDb) GetCollectionEntry(collectionID *string, databaseName *string) (*dbmodel.Collection, error) {
+func (s *collectionDb) GetCollectionWithoutMetadata(collectionID *string, databaseName *string, softDeletedFlag *bool) (*dbmodel.Collection, error) {
path: go/pkg/sysdb/metastore/db/dao/collection.go
user: login: sanketkedia
body: In future if you want to check if a particular collection is soft deleted. Would prefer this as an API as it is a bit more flexible. Not strongly opinionated though, can change if you want
created_at: 2025-05-20T23:18:17Z
author_association: CONTRIBUTOR

PR comments:
user: login: github-actions[bot]
created_at: 2025-05-20T20:36:29Z
body: # Reviewer Checklist
  Please leverage this checklist to ensure your code review is thorough before approving
  ## Testing, Bugs, Errors, Logs, Documentation
  - [ ] Can you think of any use case in which the code does not behave as intended? Have they been tested?
  - [ ] Can you think of any inputs or external events that could break the code? Is user input validated and safe? Have they been tested?
  - [ ] If appropriate, are there adequate property based tests?
  - [ ] If appropriate, are there adequate unit tests?
  - [ ] Should any logging, debugging, tracing information be added or removed?
  - [ ] Are error messages user-friendly?
  - [ ] Have all documentation changes needed been made?
  - [ ] Have all non-obvious changes been commented?
  ## System Compatibility
  - [ ] Are there any potential impacts on other parts of the system or backward compatibility?
  - [ ] Does this change intersect with any items on our roadmap, and if so, is there a plan for fitting them together?
  ## Quality
  - [ ] Is this code of a unexpectedly high quality (Readability, Modularity, Intuitiveness)
author_association: NONE
user: login: sanketkedia
created_at: 2025-05-20T20:36:35Z
body: * **#4585** <a href="https://app.graphite.dev/github/pr/chroma-core/chroma/4585?utm_source=stack-comment-icon" target="_blank"><img src="https://static.graphite.dev/graphite-32x32-black.png" alt="Graphite" width="10px" height="10px"/></a> 👈 <a href="https://app.graphite.dev/github/pr/chroma-core/chroma/4585?utm_source=stack-comment-view-in-graphite" target="_blank">(View in Graphite)</a>
* `main`


<h2></h2>

This stack of pull requests is managed by <a href="https://graphite.dev?utm-source=stack-comment"><b>Graphite</b></a>. Learn more about <a href="https://stacking.dev/?utm_source=stack-comment">stacking</a>.
<!-- Current dependencies on/for this PR: -->
author_association: CONTRIBUTOR
user: login: propel-code-bot[bot]
created_at: 2025-05-20T20:38:37Z
body: <!-- Summary by @propel-code-bot -->

**Robust Forking and Cleanup Logic: Root Collection Handling, Soft Deletes, and Retrying Improvements**

This PR overhauls and hardens the forking and deletion behavior in the system's collection management, particularly around soft and hard deletes of root collections. It updates large portions of the coordinator and DAO logic to ensure root collections are never hard deleted, enables forks to succeed even if the root collection has been soft deleted, improves retry handling for forking (including not retrying on 'not found' errors), better distinguishes between retryable and non-retryable errors, and addresses several related edge cases. It introduces new and updated tests to provide thorough coverage and explicitly modifies the soft delete cleaner to avoid hard deleting root collections, along with a new error handling structure.

**Key Changes:**
• Refactored forking logic to allow forking from soft-deleted root collections, and prevent hard deletion of root collections.
• Introduced GetCollectionWithoutMetadata in data access APIs and replaced legacy calls throughout (DB, mocks, services, coordinator).
• Restructured deletion logic and error handling to separate soft and hard delete scenarios, including new errors (e.g., ErrCollectionDeletedWithLocksHeld, ErrMissingLineageFileName, ErrCompactionOffsetSomehowAhead).
• Altered cleanup mechanism to skip hard deletion of root collections in the soft delete cleanup cycle.
• Improved handling in ForkCollection for correct lock acquisition and detection of deleted/soft-deleted collections; added fine-grained error/edge case detection (invalid lineage files, concurrent deletion, log offset invariants, etc).
• Updated and expanded testing suite: added tests for forking edge cases, soft/hard deletion, cleanup, and retry logic.
• Extensive method and signature changes across coordinator, DAO, models, and mocks for better consistency and correctness.


**Affected Areas:**
• Coordinator logic (table_catalog.go)
• Collection DAO and DB models (collection.go, test_utils.go, dbmodel/*)
• Soft delete cleanup (cleanup.go)
• Unit/integration tests (collection_service_test.go, cleaup_test.go, dao/collection_test.go, dao/segment_test.go, dao/test_utils.go)
• Mocks (ICollectionDb.go)
• Error codes/handling (common/errors.go)
• Rust frontend retry logic


**Potential Impact:**

**Functionality**: Significantly improves correctness and predictability of collection forking and deletion, especially around soft/hard deleted roots; prevents invalid state transitions and unexpected errors during concurrent operations and cleanup.

**Performance**: Minimal direct impact; slightly deeper checks and increased memory for added error cases, but should be negligible.

**Security**: No direct changes to authentication or authorization; improved error handling may reduce risk of inconsistent states.

**Scalability**: No negative impact observed; robust logic should prevent rare edge cases that could otherwise impact cluster health under scale.

**Review Focus:**
• Correctness of lock handling and race condition prevention in ForkCollection and cleanup logic.
• Error path handling—ensure all new errors are logically triggered and all edge cases are covered.
• Dependencies on new or changed interfaces (especially GetCollectionWithoutMetadata or LockCollection return semantics).
• Test robustness: validate all new/modified tests truly check introduced behaviors and edge cases.
• Impact/compatibility for downstream consumers of these APIs/methods, especially in Rust frontend changes.


<details>
<summary><strong>Testing Needed</strong></summary>

• Run all collection operations (create, fork, soft-delete, hard-delete) across normal and edge-case scenarios.
• Specifically test forking with soft/hard deleted roots, concurrent deletion attempts, and retry cases.
• Assert soft delete cleanup does not hard delete roots but cleans up non-root collections.
• Run or review expanded test suite for full functional coverage.

</details>

<details>
<summary><strong>Code Quality Assessment</strong></summary>

**go/pkg/common/errors.go**: Good: new error conditions add precision to edge case handling.

**test files**: Expanded for good coverage—tests are explicit about root/non-root behavior.

**rust/frontend/src/impls/service_based_frontend.rs**: Minor change to retry logic aligning with new error handling.

**go/pkg/sysdb/coordinator/table_catalog.go**: Large, carefully-reworked with more robust transactional logic, but high complexity warrants extra review for race conditions.

**go/pkg/sysdb/metastore/db/dao/collection.go**: Refactored for clarity and extensibility; increased flexibility with softDeletedFlag but complexity added—double-check null/boolean pathways.

**go/pkg/sysdb/metastore/db/dbmodel/mocks/ICollectionDb.go**: Regenerated/refreshed—should match new interfaces; verify for test consistency.

**go/pkg/sysdb/grpc/cleanup.go**: Logic is clear—now skips root in deletions; ensure comment/code match intent.


</details>

<details>
<summary><strong>Best Practices</strong></summary>

**Code Organization**:
• Factored retrieval by logical purpose (with/without metadata; `softDeletedFlag`)
• Reused/refactored data access code

**Error Handling**:
• Explicit error types for all major edge cases
• No generic panic or untyped errors in main logic

**Transactional Consistency**:
• Locks acquired before observable changes are made
• No window for race condition on critical resources

**Test Coverage**:
• Extensive unit/integration tests for edge cases
• All new behaviors validated


</details>

<details>
<summary><strong>Potential Issues</strong></summary>

• Potential subtleties in lock acquisition and error propagation—race conditions, null dereferences, or missed transactional failures if not all paths are consistent.
• Downstream consumers may assume deleted roots can be hard deleted or forked differently, so API/behavioral contracts must be verified.
• Some interfaces now return booleans/pointers for deleted state; double-check all use sites for correct handling.

</details>

*This summary was automatically generated by @propel-code-bot*
author_association: CONTRIBUTOR