From d0ce3f602e0b9100eaef595138f1f5b449d96034 Mon Sep 17 00:00:00 2001
From: "yihao.dai" <yihao.dai@zilliz.com>
Date: Tue, 6 Aug 2024 15:54:17 +0800
Subject: [PATCH] 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: https://github.com/milvus-io/milvus/issues/33775

pr: https://github.com/milvus-io/milvus/pull/35274

Signed-off-by: bigsheeper <yihao.dai@zilliz.com>
---
 internal/datacoord/import_util.go      |  4 ++++
 internal/datacoord/import_util_test.go |  9 +++++++++
 internal/util/importutilv2/option.go   | 14 ++++++++++++++
 3 files changed, 27 insertions(+)

diff --git a/internal/datacoord/import_util.go b/internal/datacoord/import_util.go
index dd92fec3fb205..4633d2acf949a 100644
--- 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
index 5d3ae48d97cde..411bfb51ff2e5 100644
--- 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
index 652caeda2a14c..b8e958c19fdb2 100644
--- 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
+}