[ENH]: soft delete databases, hard delete from garbage collector

Author: codetheweb

go/pkg/sysdb/coordinator/coordinator.go

@@ -2,6 +2,7 @@ package coordinator
 
 import (
 	"context"
+	"time"
 
 	"github.com/chroma-core/chroma/go/pkg/common"
 	"github.com/chroma-core/chroma/go/pkg/proto/coordinatorpb"
@@ -255,3 +256,13 @@ func (s *Coordinator) BatchGetCollectionVersionFilePaths(ctx context.Context, re
 func (s *Coordinator) BatchGetCollectionSoftDeleteStatus(ctx context.Context, req *coordinatorpb.BatchGetCollectionSoftDeleteStatusRequest) (*coordinatorpb.BatchGetCollectionSoftDeleteStatusResponse, error) {
 	return s.catalog.BatchGetCollectionSoftDeleteStatus(ctx, req.CollectionIds)
 }
+
+func (s *Coordinator) FinishDatabaseDeletion(ctx context.Context, req *coordinatorpb.FinishDatabaseDeletionRequest) (*coordinatorpb.FinishDatabaseDeletionResponse, error) {
+	err := s.catalog.FinishDatabaseDeletion(ctx, time.Unix(req.CutoffTime.Seconds, int64(req.CutoffTime.Nanos)))
+	if err != nil {
+		return nil, err
+	}
+
+	res := &coordinatorpb.FinishDatabaseDeletionResponse{}
+	return res, nil
+}

go/pkg/sysdb/coordinator/table_catalog.go

@@ -184,10 +184,32 @@ func (tc *Catalog) DeleteDatabase(ctx context.Context, deleteDatabase *model.Del
 		if len(databases) == 0 {
 			return common.ErrDatabaseNotFound
 		}
-		err = tc.metaDomain.DatabaseDb(txCtx).Delete(databases[0].ID)
+		err = tc.metaDomain.DatabaseDb(txCtx).SoftDelete(databases[0].ID)
 		if err != nil {
 			return err
 		}
+
+		collections, err := tc.metaDomain.CollectionDb(txCtx).GetCollections(nil, nil, deleteDatabase.Tenant, deleteDatabase.Name, nil, nil)
+		if err != nil {
+			return err
+		}
+
+		for _, collection := range collections {
+			collectionID, err := types.Parse(collection.Collection.ID)
+			if err != nil {
+				return err
+			}
+
+			err = tc.softDeleteCollection(txCtx, &model.DeleteCollection{
+				ID:           collectionID,
+				TenantID:     deleteDatabase.Tenant,
+				DatabaseName: deleteDatabase.Name,
+			})
+			if err != nil {
+				return err
+			}
+		}
+
 		return nil
 	})
 }
@@ -2196,3 +2218,7 @@ func (tc *Catalog) GetVersionFileNamesForCollection(ctx context.Context, tenantI
 
 	return collectionEntry.VersionFileName, nil
 }
+
+func (tc *Catalog) FinishDatabaseDeletion(ctx context.Context, cutoffTime time.Time) error {
+	return tc.metaDomain.DatabaseDb(ctx).FinishDatabaseDeletion(cutoffTime)
+}

go/pkg/sysdb/grpc/tenant_database_service.go

@@ -158,3 +158,11 @@ func (s *Server) GetLastCompactionTimeForTenant(ctx context.Context, req *coordi
 	}
 	return res, nil
 }
+
+func (s *Server) FinishDatabaseDeletion(ctx context.Context, req *coordinatorpb.FinishDatabaseDeletionRequest) (*coordinatorpb.FinishDatabaseDeletionResponse, error) {
+	res, err := s.coordinator.FinishDatabaseDeletion(ctx, req)
+	if err != nil {
+		return nil, grpcutils.BuildInternalGrpcError(err.Error())
+	}
+	return res, nil
+}

go/pkg/sysdb/grpc/tenant_database_service_test.go

@@ -145,7 +145,7 @@ func (suite *TenantDatabaseServiceTestSuite) TestServer_DeleteDatabase() {
 	})
 	suite.NoError(err)
 
-	// Check that associated collection was deleted
+	// Check that associated collection was soft deleted
 	var count int64
 	var collections []*dbmodel.Collection
 	suite.NoError(suite.db.Find(&collections).Count(&count).Error)

go/pkg/sysdb/metastore/db/dao/database.go

@@ -2,6 +2,7 @@ package dao
 
 import (
 	"errors"
+	"time"
 
 	"github.com/chroma-core/chroma/go/pkg/common"
 	"github.com/chroma-core/chroma/go/pkg/sysdb/metastore/db/dbmodel"
@@ -95,13 +96,9 @@ func (s *databaseDb) Insert(database *dbmodel.Database) error {
 	return err
 }
 
-func (s *databaseDb) Delete(databaseID string) error {
+func (s *databaseDb) SoftDelete(databaseID string) error {
 	return s.db.Transaction(func(tx *gorm.DB) error {
-		if err := tx.Where("id = ?", databaseID).Delete(&dbmodel.Database{}).Error; err != nil {
-			return err
-		}
-
-		if err := tx.Where("database_id = ?", databaseID).Delete(&dbmodel.Collection{}).Error; err != nil {
+		if err := tx.Table("databases").Where("id = ?", databaseID).Update("is_deleted", true).Error; err != nil {
 			return err
 		}
 
@@ -121,3 +118,30 @@ func (s *databaseDb) GetDatabasesByTenantID(tenantID string) ([]*dbmodel.Databas
 	}
 	return databases, nil
 }
+
+func (s *databaseDb) FinishDatabaseDeletion(cutoffTime time.Time) error {
+	for {
+		// Only hard delete databases that are soft deleted and have no collections
+		databasesSubQuery := s.db.
+			Table("databases d").
+			Select("d.id").
+			Joins("LEFT JOIN collections c ON c.database_id = d.id").
+			Where("d.is_deleted = ?", true).
+			Where("d.updated_at < ?", cutoffTime).
+			Group("d.id").
+			Having("COUNT(c.id) = 0").
+			Limit(1000)
+
+		res := s.db.Table("databases").
+			Where("id IN (?)", databasesSubQuery).
+			Delete(&dbmodel.Database{})
+		if res.Error != nil {
+			return res.Error
+		}
+		if res.RowsAffected == 0 {
+			break
+		}
+	}
+
+	return nil
+}

go/pkg/sysdb/metastore/db/dbmodel/database.go

@@ -27,5 +27,6 @@ type IDatabaseDb interface {
 	ListDatabases(limit *int32, offset *int32, tenantID string) ([]*Database, error)
 	Insert(in *Database) error
 	DeleteAll() error
-	Delete(databaseID string) error
+	SoftDelete(databaseID string) error
+	FinishDatabaseDeletion(cutoffTime time.Time) error
 }

idl/chromadb/proto/coordinator.proto

@@ -46,6 +46,14 @@ message DeleteDatabaseRequest {
 
 message DeleteDatabaseResponse {}
 
+message FinishDatabaseDeletionRequest {
+  google.protobuf.Timestamp cutoff_time = 1;
+}
+
+message FinishDatabaseDeletionResponse {
+  uint64 num_deleted = 1;
+}
+
 message CreateTenantRequest {
   string name = 2; // Names are globally unique
 }
@@ -496,6 +504,7 @@ service SysDB {
   rpc GetDatabase(GetDatabaseRequest) returns (GetDatabaseResponse) {}
   rpc ListDatabases(ListDatabasesRequest) returns (ListDatabasesResponse) {}
   rpc DeleteDatabase(DeleteDatabaseRequest) returns (DeleteDatabaseResponse) {}
+  rpc FinishDatabaseDeletion(FinishDatabaseDeletionRequest) returns (FinishDatabaseDeletionResponse) {}
   rpc CreateTenant(CreateTenantRequest) returns (CreateTenantResponse) {}
   rpc GetTenant(GetTenantRequest) returns (GetTenantResponse) {}
   rpc CreateSegment(CreateSegmentRequest) returns (CreateSegmentResponse) {}

rust/garbage_collector/src/garbage_collector_orchestrator_v2.rs

@@ -34,6 +34,7 @@ use chroma_system::{
 use chroma_types::chroma_proto::{CollectionVersionFile, VersionListForCollection};
 use chroma_types::{
     BatchGetCollectionSoftDeleteStatusError, CollectionUuid, DeleteCollectionError,
+    DeleteDatabaseError,
 };
 use chrono::{DateTime, Utc};
 use std::collections::{HashMap, HashSet};
@@ -158,6 +159,8 @@ pub enum GarbageCollectorError {
     UnparsableUuid(#[from] uuid::Error),
     #[error("Collection deletion failed: {0}")]
     CollectionDeletionFailed(#[from] DeleteCollectionError),
+    #[error("Database deletion failed: {0}")]
+    DeleteDatabaseFailed(#[from] DeleteDatabaseError),
 }
 
 impl ChromaError for GarbageCollectorError {
@@ -716,6 +719,19 @@ impl GarbageCollectorOrchestrator {
 
         self.num_pending_tasks -= 1;
         if self.num_pending_tasks == 0 {
+            let tenant = self
+                .tenant
+                .clone()
+                .ok_or(GarbageCollectorError::InvariantViolation(
+                    "Expected tenant to be set".to_string(),
+                ))?;
+            let database_name =
+                self.database_name
+                    .clone()
+                    .ok_or(GarbageCollectorError::InvariantViolation(
+                        "Expected database to be set".to_string(),
+                    ))?;
+
             for collection_id in self.soft_deleted_collections_to_gc.iter() {
                 let graph =
                     self.graph
@@ -758,16 +774,8 @@ impl GarbageCollectorOrchestrator {
                 if are_all_children_in_fork_tree_also_soft_deleted {
                     self.sysdb_client
                         .finish_collection_deletion(
-                            self.tenant.clone().ok_or(
-                                GarbageCollectorError::InvariantViolation(
-                                    "Expected tenant to be set".to_string(),
-                                ),
-                            )?,
-                            self.database_name.clone().ok_or(
-                                GarbageCollectorError::InvariantViolation(
-                                    "Expected database to be set".to_string(),
-                                ),
-                            )?,
+                            tenant.clone(),
+                            database_name.clone(),
                             *collection_id,
                         )
                         .await?;

rust/garbage_collector/tests/proptest_helpers/garbage_collector_reference.rs

@@ -232,10 +232,6 @@ impl ReferenceStateMachine for ReferenceGarbageCollector {
                 segments: segment_group,
             });
 
-        let _delete_collection_transition = alive_collection_id_strategy
-            .clone()
-            .prop_map(Transition::DeleteCollection);
-
         let fork_collection_transition =
             alive_collection_id_strategy
                 .clone()
@@ -277,11 +273,11 @@ impl ReferenceStateMachine for ReferenceGarbageCollector {
 
         if alive_collection_ids.is_empty() {
             if state.root_collection_id.is_some() {
-                // If all collections are deleted, we cannot create a new collection, so there is nothing further to do
+                // If all collections are deleted, we cannot create a new collection
                 return Just(Transition::NoOp).boxed();
             }
 
-            return prop_oneof![create_collection_transition,].boxed();
+            return create_collection_transition.boxed();
         }
 
         // While the garbage collector can technically run on any collection in a fork tree, we always run it on the root collection as the test fixture will call `ListCollectionsToGc()` which only returns the root collection.
@@ -293,10 +289,10 @@ impl ReferenceStateMachine for ReferenceGarbageCollector {
                 });
 
             return prop_oneof![
-                2 => fork_collection_transition,
-                3 => increment_collection_version_transition,
-                2 => garbage_collect_transition,
-                1 => delete_collection_transition,
+                5 => increment_collection_version_transition,
+                4 => fork_collection_transition,
+                4 => garbage_collect_transition,
+                3 => delete_collection_transition,
             ]
             .boxed();
         }

rust/garbage_collector/tests/proptest_helpers/garbage_collector_under_test.rs

@@ -32,7 +32,7 @@ define_thread_local_stats!(STATS);
 pub struct GarbageCollectorUnderTest {
     runtime: Arc<tokio::runtime::Runtime>,
     tenant: String,
-    database: String,
+    database_name: String,
     sysdb: SysDb,
     storage: Storage,
     root_manager: RootManager,
@@ -125,7 +125,7 @@ impl StateMachineTest for GarbageCollectorUnderTest {
             Self {
                 runtime: ref_state.runtime.clone(),
                 tenant: tenant_name,
-                database: database_name,
+                database_name,
                 sysdb,
                 storage: storage.clone(),
                 root_manager,
@@ -192,7 +192,7 @@ impl StateMachineTest for GarbageCollectorUnderTest {
                         .sysdb
                         .create_collection(
                             state.tenant.clone(),
-                            state.database.clone(),
+                            state.database_name.clone(),
                             collection_id,
                             format!("Collection {}", collection_id),
                             segments,
@@ -211,7 +211,7 @@ impl StateMachineTest for GarbageCollectorUnderTest {
                     .runtime
                     .block_on(state.sysdb.delete_collection(
                         state.tenant.clone(),
-                        state.database.clone(),
+                        state.database_name.clone(),
                         collection_id,
                         vec![],
                     ))