enhance: Add import option to skip disk quota check (#35274) (#35275)

Add an option to skip the disk quota check for backup-restore import. issue: #33775 pr: #35274 Signed-off-by: bigsheeper <[email protected]>
milvus-io · Aug 6, 2024 · d0ce3f6 · d0ce3f6
1 parent 6e8083d
commit d0ce3f6
Show file tree

Hide file tree

Showing 3 changed files with 27 additions and 0 deletions.
diff --git a/internal/datacoord/import_util.go b/internal/datacoord/import_util.go
@@ -273,6 +273,10 @@ func CheckDiskQuota(job ImportJob, meta *meta, imeta ImportMeta) (int64, error)
 	if !Params.QuotaConfig.DiskProtectionEnabled.GetAsBool() {
 		return 0, nil
 	}
+	if importutilv2.SkipDiskQuotaCheck(job.GetOptions()) {
+		log.Info("skip disk quota check for import", zap.Int64("jobID", job.GetJobID()))
+		return 0, nil
+	}
 
 	var (
 		requestedTotal       int64

diff --git a/internal/datacoord/import_util_test.go b/internal/datacoord/import_util_test.go
@@ -35,6 +35,7 @@ import (
 	"github.com/milvus-io/milvus/internal/proto/datapb"
 	"github.com/milvus-io/milvus/internal/proto/internalpb"
 	"github.com/milvus-io/milvus/internal/storage"
+	"github.com/milvus-io/milvus/internal/util/importutilv2"
 	"github.com/milvus-io/milvus/pkg/util/merr"
 	"github.com/milvus-io/milvus/pkg/util/paramtable"
 )
@@ -281,6 +282,14 @@ func TestImportUtil_CheckDiskQuota(t *testing.T) {
 	assert.NoError(t, err)
 
 	Params.Save(Params.QuotaConfig.DiskProtectionEnabled.Key, "true")
+	job.Options = []*commonpb.KeyValuePair{
+		{Key: importutilv2.BackupFlag, Value: "true"},
+		{Key: importutilv2.SkipDQC, Value: "true"},
+	}
+	_, err = CheckDiskQuota(job, meta, imeta)
+	assert.NoError(t, err)
+
+	job.Options = nil
 	Params.Save(Params.QuotaConfig.DiskQuota.Key, "10000")
 	Params.Save(Params.QuotaConfig.DiskQuotaPerCollection.Key, "10000")
 	defer Params.Reset(Params.QuotaConfig.DiskQuota.Key)

diff --git a/internal/util/importutilv2/option.go b/internal/util/importutilv2/option.go
@@ -35,6 +35,7 @@ const (
 	EndTs2     = "endTs"
 	BackupFlag = "backup"
 	L0Import   = "l0_import"
+	SkipDQC    = "skip_disk_quota_check"
 )
 
 type Options []*commonpb.KeyValuePair
@@ -85,3 +86,16 @@ func IsL0Import(options Options) bool {
 	}
 	return true
 }
+
+// SkipDiskQuotaCheck indicates whether the import skips the disk quota check.
+// This option should only be enabled during backup restoration.
+func SkipDiskQuotaCheck(options Options) bool {
+	if !IsBackup(options) {
+		return false
+	}
+	skip, err := funcutil.GetAttrByKeyFromRepeatedKV(SkipDQC, options)
+	if err != nil || strings.ToLower(skip) != "true" {
+		return false
+	}
+	return true
+}