Support pass database in create and restore request

Signed-off-by: wayblink <anyang.wang@zilliz.com>

cmd/create.go

@@ -14,6 +14,7 @@ import (
 var (
 	backupName      string
 	collectionNames string
+	database        string
 )
 
 var createBackupCmd = &cobra.Command{
@@ -38,6 +39,7 @@ var createBackupCmd = &cobra.Command{
 		resp := backupContext.CreateBackup(context, &backuppb.CreateBackupRequest{
 			BackupName:      backupName,
 			CollectionNames: collectionNameArr,
+			Database:        database,
 		})
 		fmt.Println(resp.GetCode(), "\n", resp.GetMsg())
 	},
@@ -46,6 +48,7 @@ var createBackupCmd = &cobra.Command{
 func init() {
 	createBackupCmd.Flags().StringVarP(&backupName, "name", "n", "", "backup name, if unset will generate a name automatically")
 	createBackupCmd.Flags().StringVarP(&collectionNames, "colls", "", "", "collectionNames to backup, use ',' to connect multiple collections")
+	createBackupCmd.Flags().StringVarP(&database, "database", "d", "", "database to backup")
 
 	rootCmd.AddCommand(createBackupCmd)
 }

cmd/restore.go

@@ -19,6 +19,7 @@ var (
 	restoreCollectionNames string
 	renameSuffix           string
 	renameCollectionNames  string
+	restoreDatabase        string
 )
 
 var restoreBackupCmd = &cobra.Command{
@@ -55,6 +56,7 @@ var restoreBackupCmd = &cobra.Command{
 			CollectionNames:   collectionNameArr,
 			CollectionSuffix:  renameSuffix,
 			CollectionRenames: renameMap,
+			Database:          restoreDatabase,
 		})
 
 		fmt.Println(resp.GetCode(), "\n", resp.GetMsg())
@@ -66,6 +68,7 @@ func init() {
 	restoreBackupCmd.Flags().StringVarP(&restoreCollectionNames, "collections", "c", "", "collectionNames to restore")
 	restoreBackupCmd.Flags().StringVarP(&renameSuffix, "suffix", "s", "", "add a suffix to collection name to restore")
 	restoreBackupCmd.Flags().StringVarP(&renameCollectionNames, "rename", "r", "", "rename collections to new names")
+	restoreBackupCmd.Flags().StringVarP(&restoreDatabase, "database", "d", "", "database to restore, if not set, restore all databases")
 
 	rootCmd.AddCommand(restoreBackupCmd)
 }

core/backup_impl_create_backup.go

@@ -25,6 +25,7 @@ func (b BackupContext) CreateBackup(ctx context.Context, request *backuppb.Creat
 		zap.String("requestId", request.GetRequestId()),
 		zap.String("backupName", request.GetBackupName()),
 		zap.Strings("collections", request.GetCollectionNames()),
+		zap.String("database", request.GetDatabase()),
 		zap.Bool("async", request.GetAsync()))
 
 	resp := &backuppb.BackupInfoResponse{
@@ -135,27 +136,46 @@ func (b BackupContext) parseBackupCollections(request *backuppb.CreateBackupRequ
 		zap.Strings("request_collection_names", request.GetCollectionNames()),
 		zap.Int("length", len(request.GetCollectionNames())))
 	var toBackupCollections []collection
+
 	if request.GetCollectionNames() == nil || len(request.GetCollectionNames()) == 0 {
-		dbs, err := b.milvusClient.ListDatabases(b.ctx)
-		if err != nil {
-			log.Error("fail in ListDatabases", zap.Error(err))
-			return nil, err
+		dbs := make([]string, 0)
+		if request.GetDatabase() == "" {
+			dbEntities, err := b.milvusClient.ListDatabases(b.ctx)
+			if err != nil {
+				log.Error("fail in ListDatabases", zap.Error(err))
+				return nil, err
+			}
+			for _, db := range dbEntities {
+				dbs = append(dbs, db.Name)
+			}
+		} else {
+			dbs = []string{request.GetDatabase()}
 		}
+
 		for _, db := range dbs {
-			b.milvusClient.UsingDatabase(b.ctx, db.Name)
+			err := b.milvusClient.UsingDatabase(b.ctx, db)
+			if err != nil {
+				log.Error("fail to call SDK use database", zap.Error(err))
+				return nil, err
+			}
 			collections, err := b.milvusClient.ListCollections(b.ctx)
 			if err != nil {
 				log.Error("fail in ListCollections", zap.Error(err))
 				return nil, err
 			}
 			for _, coll := range collections {
-				toBackupCollections = append(toBackupCollections, collection{db.Name, coll.Name})
+				toBackupCollections = append(toBackupCollections, collection{db, coll.Name})
 			}
 		}
 		log.Debug(fmt.Sprintf("List %v collections", len(toBackupCollections)))
 	} else {
 		for _, collectionName := range request.GetCollectionNames() {
-			var dbName = "default"
+			var dbName string
+			if request.GetDatabase() == "" {
+				dbName = "default"
+			} else {
+				dbName = request.GetDatabase()
+			}
 			if strings.Contains(collectionName, ".") {
 				splits := strings.Split(collectionName, ".")
 				dbName = splits[0]

core/backup_impl_restore_backup.go

@@ -29,7 +29,8 @@ func (b BackupContext) RestoreBackup(ctx context.Context, request *backuppb.Rest
 		zap.Any("CollectionRenames", request.GetCollectionRenames()),
 		zap.Bool("async", request.GetAsync()),
 		zap.String("bucketName", request.GetBucketName()),
-		zap.String("path", request.GetPath()))
+		zap.String("path", request.GetPath()),
+		zap.String("database", request.GetDatabase()))
 
 	resp := &backuppb.RestoreBackupResponse{
 		RequestId: request.GetRequestId(),
@@ -100,7 +101,16 @@ func (b BackupContext) RestoreBackup(ctx context.Context, request *backuppb.Rest
 	// 2, initial restoreCollectionTasks
 	toRestoreCollectionBackups := make([]*backuppb.CollectionBackupInfo, 0)
 	if len(request.GetCollectionNames()) == 0 {
-		toRestoreCollectionBackups = backup.GetCollectionBackups()
+		if request.GetDatabase() != "" {
+			collBackups := backup.GetCollectionBackups()
+			for _, c := range collBackups {
+				if c.GetDbName() == request.GetDatabase() {
+					toRestoreCollectionBackups = append(toRestoreCollectionBackups, c)
+				}
+			}
+		} else {
+			toRestoreCollectionBackups = backup.GetCollectionBackups()
+		}
 	} else {
 		collectionNameDict := make(map[string]bool)
 		for _, collectionName := range request.GetCollectionNames() {
@@ -139,6 +149,7 @@ func (b BackupContext) RestoreBackup(ctx context.Context, request *backuppb.Rest
 			targetCollectionName = backupCollectionName
 		}
 
+		b.milvusClient.UsingDatabase(ctx, restoreCollection.DbName)
 		exist, err := b.milvusClient.HasCollection(ctx, targetCollectionName)
 		if err != nil {
 			errorMsg := fmt.Sprintf("fail to check whether the collection is exist, collection_name: %s, err: %s", targetCollectionName, err)

core/proto/backup.proto

@@ -141,6 +141,8 @@ message CreateBackupRequest {
   repeated string collection_names = 3;
   // async or not
   bool async = 4;
+  // database
+  string database = 5;
 }
 
 /**
@@ -239,6 +241,8 @@ message RestoreBackupRequest {
   string bucket_name = 7;
   // if bucket_name and path is set. will override bucket/path in config.
   string path = 8;
+  // database
+  string database = 9;
 }
 
 message RestorePartitionTask {