Fix enabling both drop_exist_collection and skip_create_collection at…

… the same time is incorrect (#348) Signed-off-by: Syang1997 <[email protected]> Co-authored-by: Syang1997 <[email protected]>
zilliztech · May 30, 2024 · acdfde1 · acdfde1
1 parent 98aa172
commit acdfde1
Show file tree

Hide file tree

Showing 2 changed files with 18 additions and 7 deletions.
diff --git a/cmd/restore.go b/cmd/restore.go
@@ -89,7 +89,7 @@ var restoreBackupCmd = &cobra.Command{
 			MetaOnly:             restoreMetaOnly,
 			RestoreIndex:         restoreRestoreIndex,
 			UseAutoIndex:         restoreUseAutoIndex,
-			DropExistCollection:  restoreDropExistIndex,
+			DropExistCollection:  restoreDropExistCollection,
 			DropExistIndex:       restoreDropExistIndex,
 			SkipCreateCollection: restoreSkipCreateCollection,
 		})

diff --git a/core/backup_impl_restore_backup.go b/core/backup_impl_restore_backup.go
@@ -419,16 +419,27 @@ func (b *BackupContext) executeRestoreCollectionTask(ctx context.Context, backup
 	}
 
 	if task.GetDropExistCollection() {
-		err := b.milvusClient.DropCollection(ctx, targetDBName, targetCollectionName)
+		//check if the collection exist, if collection exist, will drop it
+		exist, err := b.milvusClient.HasCollection(ctx, targetDBName, targetCollectionName)
 		if err != nil {
-			errorMsg := fmt.Sprintf("fail to drop collection, CollectionName: %s.%s err: %s", targetDBName, targetCollectionName, err)
+			errorMsg := fmt.Sprintf("fail to check whether the collection is exist, collection_name: %s, err: %s", targetCollectionName, err)
 			log.Error(errorMsg)
-			task.StateCode = backuppb.RestoreTaskStateCode_FAIL
-			task.ErrorMessage = errorMsg
-			return task, err
+		}
+		if exist {
+			err := b.milvusClient.DropCollection(ctx, targetDBName, targetCollectionName)
+			if err != nil {
+				errorMsg := fmt.Sprintf("fail to drop collection, CollectionName: %s.%s err: %s", targetDBName, targetCollectionName, err)
+				log.Error(errorMsg)
+				task.StateCode = backuppb.RestoreTaskStateCode_FAIL
+				task.ErrorMessage = errorMsg
+				return task, err
+			}
 		}
 	}
-	if !task.GetSkipCreateCollection() {
+	//in function RestoreBackup, before executing the restore,
+	//the SkipCreateCollection has been checked,
+	//so here it is necessary to be compatible with the situation where SkipCreateCollection and DropExistCollection are enabled at the same time.
+	if !task.GetSkipCreateCollection() || task.GetDropExistCollection() {
 		err := retry.Do(ctx, func() error {
 			if hasPartitionKey {
 				partitionNum := len(task.GetCollBackup().GetPartitionBackups())