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

Add an option to skip the disk quota check for backup-restore import.

issue: https://github.com/milvus-io/milvus/issues/33775

Signed-off-by: bigsheeper <yihao.dai@zilliz.com>

internal/datacoord/import_util.go

@@ -272,6 +272,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

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"
 )
@@ -280,6 +281,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)

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
+}