From d9ca06eb4e5c3dcddf8ab485a84f8ab57589cba7 Mon Sep 17 00:00:00 2001
From: Jiquan Long <jiquan.long@zilliz.com>
Date: Thu, 16 Feb 2023 16:34:38 +0800
Subject: [PATCH] Fix max partition number don't work (#22216)

Signed-off-by: longjiquan <jiquan.long@zilliz.com>
---
 internal/rootcoord/create_partition_task.go     |  7 +++++++
 .../rootcoord/create_partition_task_test.go     | 17 +++++++++++++++++
 2 files changed, 24 insertions(+)

diff --git a/internal/rootcoord/create_partition_task.go b/internal/rootcoord/create_partition_task.go
index 5c608398fd..59bdd4d2b3 100644
--- a/internal/rootcoord/create_partition_task.go
+++ b/internal/rootcoord/create_partition_task.go
@@ -2,6 +2,7 @@ package rootcoord
 
 import (
 	"context"
+	"fmt"
 
 	pb "github.com/milvus-io/milvus/internal/proto/etcdpb"
 
@@ -40,6 +41,12 @@ func (t *createPartitionTask) Execute(ctx context.Context) error {
 		}
 	}
 
+	cfgMaxPartitionNum := Params.RootCoordCfg.MaxPartitionNum.GetAsInt()
+	if len(t.collMeta.Partitions) >= cfgMaxPartitionNum {
+		return fmt.Errorf("partition number (%d) exceeds max configuration (%d), collection: %s",
+			len(t.collMeta.Partitions), cfgMaxPartitionNum, t.collMeta.Name)
+	}
+
 	partID, err := t.core.idAllocator.AllocOne()
 	if err != nil {
 		return err
diff --git a/internal/rootcoord/create_partition_task_test.go b/internal/rootcoord/create_partition_task_test.go
index 34a1d9c73f..342d81c3d2 100644
--- a/internal/rootcoord/create_partition_task_test.go
+++ b/internal/rootcoord/create_partition_task_test.go
@@ -64,6 +64,23 @@ func Test_createPartitionTask_Execute(t *testing.T) {
 		assert.NoError(t, err)
 	})
 
+	t.Run("create too many partitions", func(t *testing.T) {
+		cfgMaxPartitionNum := Params.RootCoordCfg.MaxPartitionNum.GetAsInt()
+		partitions := make([]*model.Partition, 0, cfgMaxPartitionNum)
+		for i := 0; i < cfgMaxPartitionNum; i++ {
+			partitions = append(partitions, &model.Partition{})
+		}
+		collectionName := funcutil.GenRandomStr()
+		partitionName := funcutil.GenRandomStr()
+		coll := &model.Collection{Name: collectionName, Partitions: partitions}
+		task := &createPartitionTask{
+			collMeta: coll,
+			Req:      &milvuspb.CreatePartitionRequest{CollectionName: collectionName, PartitionName: partitionName},
+		}
+		err := task.Execute(context.Background())
+		assert.Error(t, err)
+	})
+
 	t.Run("failed to allocate partition id", func(t *testing.T) {
 		collectionName := funcutil.GenRandomStr()
 		partitionName := funcutil.GenRandomStr()