enhance: refine rootcoord/metatable interfaces to ensure that each me…

…thod includes a ctx parameter (#37846) issue: #35917 Before enhancing log trace information, it's necessary to pass the context to the method entry point. This PR first refine the rootcoord/metatable interfaces to ensure that each method includes a ctx parameter. Signed-off-by: tinswzy <[email protected]>
milvus-io · Nov 21, 2024 · e247ff9 · e247ff9
1 parent 965bda6
commit e247ff9
Show file tree

Hide file tree

Showing 16 changed files with 658 additions and 627 deletions.
diff --git a/internal/rootcoord/alter_collection_task.go b/internal/rootcoord/alter_collection_task.go
@@ -84,7 +84,7 @@ func (a *alterCollectionTask) Execute(ctx context.Context) error {
 	})
 
 	// properties needs to be refreshed in the cache
-	aliases := a.core.meta.ListAliasesByID(oldColl.CollectionID)
+	aliases := a.core.meta.ListAliasesByID(ctx, oldColl.CollectionID)
 	redoTask.AddSyncStep(&expireCacheStep{
 		baseStep:        baseStep{core: a.core},
 		dbName:          a.Req.GetDbName(),

diff --git a/internal/rootcoord/alter_collection_task_test.go b/internal/rootcoord/alter_collection_task_test.go
@@ -92,7 +92,7 @@ func Test_alterCollectionTask_Execute(t *testing.T) {
 			mock.Anything,
 			mock.Anything,
 		).Return(errors.New("err"))
-		meta.On("ListAliasesByID", mock.Anything).Return([]string{})
+		meta.On("ListAliasesByID", mock.Anything, mock.Anything).Return([]string{})
 
 		core := newTestCore(withValidProxyManager(), withMeta(meta))
 		task := &alterCollectionTask{
@@ -122,7 +122,7 @@ func Test_alterCollectionTask_Execute(t *testing.T) {
 			mock.Anything,
 			mock.Anything,
 		).Return(nil)
-		meta.On("ListAliasesByID", mock.Anything).Return([]string{})
+		meta.On("ListAliasesByID", mock.Anything, mock.Anything).Return([]string{})
 
 		broker := newMockBroker()
 		broker.BroadcastAlteredCollectionFunc = func(ctx context.Context, req *milvuspb.AlterCollectionRequest) error {
@@ -157,7 +157,7 @@ func Test_alterCollectionTask_Execute(t *testing.T) {
 			mock.Anything,
 			mock.Anything,
 		).Return(nil)
-		meta.On("ListAliasesByID", mock.Anything).Return([]string{})
+		meta.On("ListAliasesByID", mock.Anything, mock.Anything).Return([]string{})
 
 		broker := newMockBroker()
 		broker.BroadcastAlteredCollectionFunc = func(ctx context.Context, req *milvuspb.AlterCollectionRequest) error {
@@ -231,7 +231,7 @@ func Test_alterCollectionTask_Execute(t *testing.T) {
 			mock.Anything,
 			mock.Anything,
 		).Return(nil)
-		meta.On("ListAliasesByID", mock.Anything).Return([]string{})
+		meta.On("ListAliasesByID", mock.Anything, mock.Anything).Return([]string{})
 
 		broker := newMockBroker()
 		broker.BroadcastAlteredCollectionFunc = func(ctx context.Context, req *milvuspb.AlterCollectionRequest) error {

diff --git a/internal/rootcoord/describe_collection_task.go b/internal/rootcoord/describe_collection_task.go
@@ -45,7 +45,7 @@ func (t *describeCollectionTask) Execute(ctx context.Context) (err error) {
 		return err
 	}
 
-	aliases := t.core.meta.ListAliasesByID(coll.CollectionID)
+	aliases := t.core.meta.ListAliasesByID(ctx, coll.CollectionID)
 	db, err := t.core.meta.GetDatabaseByID(ctx, coll.DBID, t.GetTs())
 	if err != nil {
 		return err

diff --git a/internal/rootcoord/describe_collection_task_test.go b/internal/rootcoord/describe_collection_task_test.go
@@ -105,6 +105,7 @@ func Test_describeCollectionTask_Execute(t *testing.T) {
 		}, nil)
 		meta.On("ListAliasesByID",
 			mock.Anything,
+			mock.Anything,
 		).Return([]string{alias1, alias2})
 		meta.EXPECT().GetDatabaseByID(mock.Anything, mock.Anything, mock.Anything).Return(&model.Database{
 			ID:   1,

diff --git a/internal/rootcoord/drop_collection_task.go b/internal/rootcoord/drop_collection_task.go
@@ -37,18 +37,18 @@ type dropCollectionTask struct {
 	Req *milvuspb.DropCollectionRequest
 }
 
-func (t *dropCollectionTask) validate() error {
+func (t *dropCollectionTask) validate(ctx context.Context) error {
 	if err := CheckMsgType(t.Req.GetBase().GetMsgType(), commonpb.MsgType_DropCollection); err != nil {
 		return err
 	}
-	if t.core.meta.IsAlias(t.Req.GetDbName(), t.Req.GetCollectionName()) {
+	if t.core.meta.IsAlias(ctx, t.Req.GetDbName(), t.Req.GetCollectionName()) {
 		return fmt.Errorf("cannot drop the collection via alias = %s", t.Req.CollectionName)
 	}
 	return nil
 }
 
 func (t *dropCollectionTask) Prepare(ctx context.Context) error {
-	return t.validate()
+	return t.validate(ctx)
 }
 
 func (t *dropCollectionTask) Execute(ctx context.Context) error {
@@ -68,7 +68,7 @@ func (t *dropCollectionTask) Execute(ctx context.Context) error {
 	}
 
 	// meta cache of all aliases should also be cleaned.
-	aliases := t.core.meta.ListAliasesByID(collMeta.CollectionID)
+	aliases := t.core.meta.ListAliasesByID(ctx, collMeta.CollectionID)
 
 	ts := t.GetTs()
 

diff --git a/internal/rootcoord/drop_collection_task_test.go b/internal/rootcoord/drop_collection_task_test.go
@@ -51,6 +51,7 @@ func Test_dropCollectionTask_Prepare(t *testing.T) {
 		meta.On("IsAlias",
 			mock.Anything,
 			mock.Anything,
+			mock.Anything,
 		).Return(true)
 
 		core := newTestCore(withMeta(meta))
@@ -72,6 +73,7 @@ func Test_dropCollectionTask_Prepare(t *testing.T) {
 		meta.On("IsAlias",
 			mock.Anything,
 			mock.Anything,
+			mock.Anything,
 		).Return(false)
 
 		core := newTestCore(withMeta(meta))
@@ -129,6 +131,7 @@ func Test_dropCollectionTask_Execute(t *testing.T) {
 			mock.Anything,
 		).Return(coll.Clone(), nil)
 		meta.On("ListAliasesByID",
+			mock.Anything,
 			mock.AnythingOfType("int64"),
 		).Return([]string{})
 
@@ -163,6 +166,7 @@ func Test_dropCollectionTask_Execute(t *testing.T) {
 		).Return(errors.New("error mock ChangeCollectionState"))
 		meta.On("ListAliasesByID",
 			mock.Anything,
+			mock.Anything,
 		).Return([]string{})
 
 		core := newTestCore(withValidProxyManager(), withMeta(meta))
@@ -207,6 +211,7 @@ func Test_dropCollectionTask_Execute(t *testing.T) {
 		).Return(nil)
 		meta.On("ListAliasesByID",
 			mock.Anything,
+			mock.Anything,
 		).Return([]string{})
 		removeCollectionMetaCalled := false
 		removeCollectionMetaChan := make(chan struct{}, 1)

diff --git a/internal/rootcoord/list_db_task.go b/internal/rootcoord/list_db_task.go
@@ -58,7 +58,7 @@ func (t *listDatabaseTask) Execute(ctx context.Context) error {
 			privilegeDBs.Insert(util.AnyWord)
 			return privilegeDBs, nil
 		}
-		userRoles, err := t.core.meta.SelectUser("", &milvuspb.UserEntity{
+		userRoles, err := t.core.meta.SelectUser(ctx, "", &milvuspb.UserEntity{
 			Name: curUser,
 		}, true)
 		if err != nil {
@@ -72,7 +72,7 @@ func (t *listDatabaseTask) Execute(ctx context.Context) error {
 				privilegeDBs.Insert(util.AnyWord)
 				return privilegeDBs, nil
 			}
-			entities, err := t.core.meta.SelectGrant("", &milvuspb.GrantEntity{
+			entities, err := t.core.meta.SelectGrant(ctx, "", &milvuspb.GrantEntity{
 				Role:   role,
 				DbName: util.AnyWord,
 			})

diff --git a/internal/rootcoord/list_db_task_test.go b/internal/rootcoord/list_db_task_test.go
@@ -132,7 +132,7 @@ func Test_ListDBTask(t *testing.T) {
 
 		{
 			// select role fail
-			meta.EXPECT().SelectUser(mock.Anything, mock.Anything, mock.Anything).
+			meta.EXPECT().SelectUser(mock.Anything, mock.Anything, mock.Anything, mock.Anything).
 				Return(nil, errors.New("mock select user error")).Once()
 			ctx := GetContext(context.Background(), "foo:root")
 			task := getTask()
@@ -142,7 +142,7 @@ func Test_ListDBTask(t *testing.T) {
 
 		{
 			// select role, empty result
-			meta.EXPECT().SelectUser(mock.Anything, mock.Anything, mock.Anything).
+			meta.EXPECT().SelectUser(mock.Anything, mock.Anything, mock.Anything, mock.Anything).
 				Return([]*milvuspb.UserResult{}, nil).Once()
 			ctx := GetContext(context.Background(), "foo:root")
 			task := getTask()
@@ -153,7 +153,7 @@ func Test_ListDBTask(t *testing.T) {
 
 		{
 			// select role, the user is added to admin role
-			meta.EXPECT().SelectUser(mock.Anything, mock.Anything, mock.Anything).
+			meta.EXPECT().SelectUser(mock.Anything, mock.Anything, mock.Anything, mock.Anything).
 				Return([]*milvuspb.UserResult{
 					{
 						User: &milvuspb.UserEntity{
@@ -176,7 +176,7 @@ func Test_ListDBTask(t *testing.T) {
 
 		{
 			// select grant fail
-			meta.EXPECT().SelectUser(mock.Anything, mock.Anything, mock.Anything).
+			meta.EXPECT().SelectUser(mock.Anything, mock.Anything, mock.Anything, mock.Anything).
 				Return([]*milvuspb.UserResult{
 					{
 						User: &milvuspb.UserEntity{
@@ -189,7 +189,7 @@ func Test_ListDBTask(t *testing.T) {
 						},
 					},
 				}, nil).Once()
-			meta.EXPECT().SelectGrant(mock.Anything, mock.Anything).
+			meta.EXPECT().SelectGrant(mock.Anything, mock.Anything, mock.Anything).
 				Return(nil, errors.New("mock select grant error")).Once()
 			ctx := GetContext(context.Background(), "foo:root")
 			task := getTask()
@@ -199,7 +199,7 @@ func Test_ListDBTask(t *testing.T) {
 
 		{
 			// normal user
-			meta.EXPECT().SelectUser(mock.Anything, mock.Anything, mock.Anything).
+			meta.EXPECT().SelectUser(mock.Anything, mock.Anything, mock.Anything, mock.Anything).
 				Return([]*milvuspb.UserResult{
 					{
 						User: &milvuspb.UserEntity{
@@ -220,7 +220,7 @@ func Test_ListDBTask(t *testing.T) {
 					Name: "default",
 				},
 			}, nil).Once()
-			meta.EXPECT().SelectGrant(mock.Anything, mock.Anything).
+			meta.EXPECT().SelectGrant(mock.Anything, mock.Anything, mock.Anything).
 				Return([]*milvuspb.GrantEntity{
 					{
 						DbName: "fooDB",
@@ -236,7 +236,7 @@ func Test_ListDBTask(t *testing.T) {
 
 		{
 			// normal user with any db privilege
-			meta.EXPECT().SelectUser(mock.Anything, mock.Anything, mock.Anything).
+			meta.EXPECT().SelectUser(mock.Anything, mock.Anything, mock.Anything, mock.Anything).
 				Return([]*milvuspb.UserResult{
 					{
 						User: &milvuspb.UserEntity{
@@ -257,7 +257,7 @@ func Test_ListDBTask(t *testing.T) {
 					Name: "default",
 				},
 			}, nil).Once()
-			meta.EXPECT().SelectGrant(mock.Anything, mock.Anything).
+			meta.EXPECT().SelectGrant(mock.Anything, mock.Anything, mock.Anything).
 				Return([]*milvuspb.GrantEntity{
 					{
 						DbName: "*",