[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,15 @@ 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) {
+	numDeleted, err := s.catalog.FinishDatabaseDeletion(ctx, time.Unix(req.CutoffTime.Seconds, int64(req.CutoffTime.Nanos)))
+	if err != nil {
+		return nil, err
+	}
+
+	res := &coordinatorpb.FinishDatabaseDeletionResponse{
+		NumDeleted: numDeleted,
+	}
+	return res, nil
+}

go/pkg/sysdb/coordinator/table_catalog.go

@@ -184,27 +184,36 @@ 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
 	})
 }
 
-func (tc *Catalog) GetAllDatabases(ctx context.Context, ts types.Timestamp) ([]*model.Database, error) {
-	databases, err := tc.metaDomain.DatabaseDb(ctx).GetAllDatabases()
-	if err != nil {
-		log.Error("error getting all databases", zap.Error(err))
-		return nil, err
-	}
-	result := make([]*model.Database, 0, len(databases))
-	for _, database := range databases {
-		result = append(result, convertDatabaseToModel(database))
-	}
-	return result, nil
-}
-
 func (tc *Catalog) CreateTenant(ctx context.Context, createTenant *model.CreateTenant, ts types.Timestamp) (*model.Tenant, error) {
 	var result *model.Tenant
 
@@ -670,7 +679,7 @@ func (tc *Catalog) softDeleteCollection(ctx context.Context, deleteCollection *m
 
 		// Generate new name with timestamp and random number
 		oldName := *collections[0].Collection.Name
-		newName := fmt.Sprintf("_deleted_%s_%s", oldName, *types.FromUniqueID(deleteCollection.ID))
+		newName := fmt.Sprintf("_deleted_%s_%s", oldName, deleteCollection.ID.String())
 
 		dbCollection := &dbmodel.Collection{
 			ID:        deleteCollection.ID.String(),
@@ -2202,3 +2211,7 @@ func (tc *Catalog) GetVersionFileNamesForCollection(ctx context.Context, tenantI
 
 	return collectionEntry.VersionFileName, nil
 }
+
+func (tc *Catalog) FinishDatabaseDeletion(ctx context.Context, cutoffTime time.Time) (uint64, 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

@@ -14,6 +14,7 @@ import (
 	"github.com/chroma-core/chroma/go/pkg/sysdb/metastore/db/dbcore"
 	"github.com/chroma-core/chroma/go/pkg/sysdb/metastore/db/dbmodel"
 	s3metastore "github.com/chroma-core/chroma/go/pkg/sysdb/metastore/s3"
+	"github.com/chroma-core/chroma/go/pkg/types"
 	"github.com/google/uuid"
 	"github.com/pingcap/log"
 	"github.com/stretchr/testify/suite"
@@ -116,7 +117,7 @@ func (suite *TenantDatabaseServiceTestSuite) TestServer_DeleteDatabase() {
 	tenantName := "TestDeleteDatabase"
 	databaseName := "TestDeleteDatabase"
 	// Generate random uuid for db id
-	databaseeId := uuid.New().String()
+	databaseId := uuid.New().String()
 
 	_, err := suite.catalog.CreateTenant(context.Background(), &model.CreateTenant{
 		Name: tenantName,
@@ -127,29 +128,61 @@ func (suite *TenantDatabaseServiceTestSuite) TestServer_DeleteDatabase() {
 	_, err = suite.catalog.CreateDatabase(context.Background(), &model.CreateDatabase{
 		Tenant: tenantName,
 		Name:   databaseName,
-		ID:     databaseeId,
+		ID:     databaseId,
 		Ts:     time.Now().Unix(),
 	}, time.Now().Unix())
 	suite.NoError(err)
 
+	collectionID := types.NewUniqueID()
 	_, _, err = suite.catalog.CreateCollection(context.Background(), &model.CreateCollection{
+		ID:           collectionID,
 		TenantID:     tenantName,
 		DatabaseName: databaseName,
 		Name:         "TestCollection",
 	}, time.Now().Unix())
 	suite.NoError(err)
 
+	timeBeforeSoftDelete := time.Now()
+
 	err = suite.catalog.DeleteDatabase(context.Background(), &model.DeleteDatabase{
 		Tenant: tenantName,
 		Name:   databaseName,
 	})
 	suite.NoError(err)
 
-	// Check that associated collection was deleted
-	var count int64
+	// Check that associated collection was soft deleted
 	var collections []*dbmodel.Collection
-	suite.NoError(suite.db.Find(&collections).Count(&count).Error)
-	suite.Equal(int64(0), count)
+	suite.NoError(suite.db.Find(&collections).Error)
+	suite.Equal(1, len(collections))
+	suite.Equal(true, collections[0].IsDeleted)
+
+	// Database should not be eligible for hard deletion yet because it still has a (soft deleted) collection
+	numDeleted, err := suite.catalog.FinishDatabaseDeletion(context.Background(), time.Now())
+	suite.NoError(err)
+	suite.Equal(uint64(0), numDeleted)
+
+	// Hard delete associated collection
+	suite.NoError(err)
+	suite.NoError(suite.catalog.DeleteCollection(context.Background(), &model.DeleteCollection{
+		TenantID:     tenantName,
+		DatabaseName: databaseName,
+		ID:           collectionID,
+	}, false))
+
+	// Database should now be eligible for hard deletion, but first verify that database is not deleted if cutoff time is prior to soft delete
+	numDeleted, err = suite.catalog.FinishDatabaseDeletion(context.Background(), timeBeforeSoftDelete)
+	suite.NoError(err)
+	suite.Equal(uint64(0), numDeleted)
+
+	// Hard delete database
+	numDeleted, err = suite.catalog.FinishDatabaseDeletion(context.Background(), time.Now())
+	suite.NoError(err)
+	suite.Equal(uint64(1), numDeleted)
+
+	// Verify that database is hard deleted
+	var databases []*dbmodel.Database
+	suite.NoError(suite.db.Debug().Where("id = ?", databaseId).Find(&databases).Error)
+	suite.Equal(0, len(databases))
 }
 
 func TestTenantDatabaseServiceTestSuite(t *testing.T) {

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"
@@ -28,21 +29,12 @@ func (s *databaseDb) DeleteByTenantIdAndName(tenantId string, databaseName strin
 	return len(databases), err
 }
 
-func (s *databaseDb) GetAllDatabases() ([]*dbmodel.Database, error) {
-	var databases []*dbmodel.Database
-	query := s.db.Table("databases")
-
-	if err := query.Find(&databases).Error; err != nil {
-		return nil, err
-	}
-	return databases, nil
-}
-
 func (s *databaseDb) ListDatabases(limit *int32, offset *int32, tenantID string) ([]*dbmodel.Database, error) {
 	var databases []*dbmodel.Database
 	query := s.db.Table("databases").
 		Select("databases.id, databases.name, databases.tenant_id").
 		Where("databases.tenant_id = ?", tenantID).
+		Where("databases.is_deleted = ?", false).
 		Order("databases.created_at ASC")
 
 	if limit != nil {
@@ -65,7 +57,8 @@ func (s *databaseDb) GetDatabases(tenantID string, databaseName string) ([]*dbmo
 	query := s.db.Table("databases").
 		Select("databases.id, databases.name, databases.tenant_id").
 		Where("databases.name = ?", databaseName).
-		Where("databases.tenant_id = ?", tenantID)
+		Where("databases.tenant_id = ?", tenantID).
+		Where("databases.is_deleted = ?", false)
 
 	if err := query.Find(&databases).Error; err != nil {
 		log.Error("GetDatabases", zap.Error(err))
@@ -95,13 +88,13 @@ 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).
+			Update("updated_at", time.Now()).
+			Error; err != nil {
 			return err
 		}
 
@@ -113,11 +106,44 @@ func (s *databaseDb) GetDatabasesByTenantID(tenantID string) ([]*dbmodel.Databas
 	var databases []*dbmodel.Database
 	query := s.db.Table("databases").
 		Select("databases.id, databases.name, databases.tenant_id").
-		Where("databases.tenant_id = ?", tenantID)
+		Where("databases.tenant_id = ?", tenantID).
+		Where("databases.is_deleted = ?", false)
 
 	if err := query.Find(&databases).Error; err != nil {
 		log.Error("GetDatabasesByTenantID", zap.Error(err))
 		return nil, err
 	}
 	return databases, nil
 }
+
+func (s *databaseDb) FinishDatabaseDeletion(cutoffTime time.Time) (uint64, error) {
+	numDeleted := uint64(0)
+
+	for {
+		// Only hard delete databases that were soft deleted prior to the cutoff time 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 numDeleted, res.Error
+		}
+
+		numDeleted += uint64(res.RowsAffected)
+
+		if res.RowsAffected == 0 {
+			break
+		}
+	}
+
+	return numDeleted, nil
+}

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

@@ -22,10 +22,10 @@ func (v Database) TableName() string {
 
 //go:generate mockery --name=IDatabaseDb
 type IDatabaseDb interface {
-	GetAllDatabases() ([]*Database, error)
 	GetDatabases(tenantID string, databaseName string) ([]*Database, error)
 	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) (uint64, 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/sysdb/src/sysdb.rs

@@ -20,9 +20,9 @@ use chroma_types::{
 };
 use chroma_types::{
     BatchGetCollectionSoftDeleteStatusError, BatchGetCollectionVersionFilePathsError, Collection,
-    CollectionConversionError, CollectionUuid, CountForksError, FlushCompactionResponse,
-    FlushCompactionResponseConversionError, ForkCollectionError, Segment, SegmentConversionError,
-    SegmentScope, Tenant,
+    CollectionConversionError, CollectionUuid, CountForksError, FinishDatabaseDeletionError,
+    FlushCompactionResponse, FlushCompactionResponseConversionError, ForkCollectionError, Segment,
+    SegmentConversionError, SegmentScope, Tenant,
 };
 use std::collections::HashMap;
 use std::fmt::Debug;
@@ -126,6 +126,17 @@ impl SysDb {
         }
     }
 
+    pub async fn finish_database_deletion(
+        &mut self,
+        cutoff_time: SystemTime,
+    ) -> Result<usize, FinishDatabaseDeletionError> {
+        match self {
+            SysDb::Grpc(grpc) => grpc.finish_database_deletion(cutoff_time).await,
+            SysDb::Sqlite(_) => unimplemented!(),
+            SysDb::Test(_) => todo!(),
+        }
+    }
+
     pub async fn get_collections(
         &mut self,
         collection_id: Option<CollectionUuid>,
@@ -806,6 +817,22 @@ impl GrpcSysDb {
         }
     }
 
+    async fn finish_database_deletion(
+        &mut self,
+        cutoff_time: SystemTime,
+    ) -> Result<usize, FinishDatabaseDeletionError> {
+        let req = chroma_proto::FinishDatabaseDeletionRequest {
+            cutoff_time: Some(cutoff_time.into()),
+        };
+
+        let res = self
+            .client
+            .finish_database_deletion(req)
+            .await
+            .map_err(|e| TonicError(e).boxed())?;
+        Ok(res.into_inner().num_deleted as usize)
+    }
+
     async fn get_collections(
         &mut self,
         collection_id: Option<CollectionUuid>,