enhance: Refresh proxy cache after restore rbac meta (#35636)

issue: #35443 pr: #35635 --------- Signed-off-by: Wei Liu <[email protected]>
milvus-io · Aug 22, 2024 · 421a00b · 421a00b
1 parent 9c67965
commit 421a00b
Show file tree

Hide file tree

Showing 4 changed files with 61 additions and 6 deletions.
diff --git a/internal/metastore/kv/rootcoord/kv_catalog.go b/internal/metastore/kv/rootcoord/kv_catalog.go
@@ -1343,6 +1343,7 @@ func (kc *Catalog) RestoreRBAC(ctx context.Context, tenant string, meta *milvusp
 
 	// restore grant
 	for _, grant := range meta.Grants {
+		grant.Grantor.Privilege.Name = util.PrivilegeNameForMetastore(grant.Grantor.Privilege.Name)
 		err = kc.AlterGrant(ctx, tenant, grant, milvuspb.OperatePrivilegeType_Grant)
 		if err != nil {
 			return err

diff --git a/internal/rootcoord/root_coord.go b/internal/rootcoord/root_coord.go
@@ -2240,7 +2240,7 @@ func (c *Core) DropRole(ctx context.Context, in *milvuspb.DropRoleRequest) (*com
 	})
 	if len(grantEntities) != 0 {
 		errMsg := "fail to drop the role that it has privileges. Use REVOKE API to revoke privileges"
-		ctxLog.Warn(errMsg, zap.Error(err))
+		ctxLog.Warn(errMsg, zap.Any("grants", grantEntities), zap.Error(err))
 		return merr.StatusWithErrorCode(errors.New(errMsg), commonpb.ErrorCode_DropRoleFailure), nil
 	}
 	redoTask := newBaseRedoTask(c.stepExecutor)
@@ -2730,15 +2730,36 @@ func (c *Core) RestoreRBAC(ctx context.Context, in *milvuspb.RestoreRBACMetaRequ
 	method := "RestoreRBAC"
 	metrics.RootCoordDDLReqCounter.WithLabelValues(method, metrics.TotalLabel).Inc()
 	tr := timerecord.NewTimeRecorder(method)
-	ctxLog := log.Ctx(ctx).With(zap.String("role", typeutil.RootCoordRole), zap.Any("in", in))
+	ctxLog := log.Ctx(ctx).With(zap.String("role", typeutil.RootCoordRole))
 	ctxLog.Debug(method)
 
 	if err := merr.CheckHealthy(c.GetStateCode()); err != nil {
 		return merr.Status(err), nil
 	}
 
-	if err := c.meta.RestoreRBAC(ctx, util.DefaultTenant, in.RBACMeta); err != nil {
-		return merr.Status(err), nil
+	redoTask := newBaseRedoTask(c.stepExecutor)
+	redoTask.AddSyncStep(NewSimpleStep("restore rbac meta data", func(ctx context.Context) ([]nestedStep, error) {
+		if err := c.meta.RestoreRBAC(ctx, util.DefaultTenant, in.RBACMeta); err != nil {
+			log.Warn("fail to restore rbac meta data", zap.Any("in", in), zap.Error(err))
+			return nil, err
+		}
+		return nil, nil
+	}))
+	redoTask.AddAsyncStep(NewSimpleStep("operate privilege cache", func(ctx context.Context) ([]nestedStep, error) {
+		if err := c.proxyClientManager.RefreshPolicyInfoCache(c.ctx, &proxypb.RefreshPolicyInfoCacheRequest{
+			OpType: int32(typeutil.CacheRefresh),
+		}); err != nil {
+			log.Warn("fail to refresh policy info cache", zap.Any("in", in), zap.Error(err))
+			return nil, err
+		}
+		return nil, nil
+	}))
+
+	err := redoTask.Execute(ctx)
+	if err != nil {
+		errMsg := "fail to execute task when restore rbac meta data"
+		log.Warn(errMsg, zap.Error(err))
+		return merr.StatusWithErrorCode(err, commonpb.ErrorCode_OperatePrivilegeFailure), nil
 	}
 
 	ctxLog.Debug(method + " success")

diff --git a/internal/rootcoord/root_coord_test.go b/internal/rootcoord/root_coord_test.go
@@ -40,6 +40,7 @@ import (
 	mockrootcoord "github.com/milvus-io/milvus/internal/rootcoord/mocks"
 	"github.com/milvus-io/milvus/internal/util/dependency"
 	kvfactory "github.com/milvus-io/milvus/internal/util/dependency/kv"
+	"github.com/milvus-io/milvus/internal/util/proxyutil"
 	"github.com/milvus-io/milvus/internal/util/sessionutil"
 	"github.com/milvus-io/milvus/pkg/util"
 	"github.com/milvus-io/milvus/pkg/util/etcd"
@@ -1990,6 +1991,9 @@ func TestCore_BackupRBAC(t *testing.T) {
 func TestCore_RestoreRBAC(t *testing.T) {
 	meta := mockrootcoord.NewIMetaTable(t)
 	c := newTestCore(withHealthyCode(), withMeta(meta))
+	mockProxyClientManager := proxyutil.NewMockProxyClientManager(t)
+	mockProxyClientManager.EXPECT().RefreshPolicyInfoCache(mock.Anything, mock.Anything).Return(nil)
+	c.proxyClientManager = mockProxyClientManager
 
 	meta.EXPECT().RestoreRBAC(mock.Anything, mock.Anything, mock.Anything).Return(nil)
 	resp, err := c.RestoreRBAC(context.Background(), &milvuspb.RestoreRBACMetaRequest{})

diff --git a/tests/integration/rbac/rbac_backup_test.go b/tests/integration/rbac/rbac_backup_test.go
@@ -87,7 +87,7 @@ func (s *RBACBackupTestSuite) TestBackup() {
 			DbName:     util.AnyWord,
 			Grantor: &milvuspb.GrantorEntity{
 				User:      &milvuspb.UserEntity{Name: util.UserRoot},
-				Privilege: &milvuspb.PrivilegeEntity{Name: util.AnyWord},
+				Privilege: &milvuspb.PrivilegeEntity{Name: "Search"},
 			},
 		},
 	})
@@ -131,7 +131,7 @@ func (s *RBACBackupTestSuite) TestBackup() {
 			DbName:     util.AnyWord,
 			Grantor: &milvuspb.GrantorEntity{
 				User:      &milvuspb.UserEntity{Name: util.UserRoot},
-				Privilege: &milvuspb.PrivilegeEntity{Name: util.AnyWord},
+				Privilege: &milvuspb.PrivilegeEntity{Name: "Search"},
 			},
 		},
 	})
@@ -159,6 +159,35 @@ func (s *RBACBackupTestSuite) TestBackup() {
 	s.NoError(err)
 	s.True(merr.Ok(resp11.GetStatus()))
 	s.Equal(resp11.GetRBACMeta().String(), resp5.GetRBACMeta().String())
+
+	// clean rbac meta
+	resp12, err := s.Cluster.Proxy.OperatePrivilege(ctx, &milvuspb.OperatePrivilegeRequest{
+		Type: milvuspb.OperatePrivilegeType_Revoke,
+		Entity: &milvuspb.GrantEntity{
+			Role:       &milvuspb.RoleEntity{Name: roleName},
+			Object:     &milvuspb.ObjectEntity{Name: commonpb.ObjectType_Collection.String()},
+			ObjectName: util.AnyWord,
+			DbName:     util.AnyWord,
+			Grantor: &milvuspb.GrantorEntity{
+				User:      &milvuspb.UserEntity{Name: util.UserRoot},
+				Privilege: &milvuspb.PrivilegeEntity{Name: "Search"},
+			},
+		},
+	})
+	s.NoError(err)
+	s.True(merr.Ok(resp12))
+
+	resp13, err := s.Cluster.Proxy.DropRole(ctx, &milvuspb.DropRoleRequest{
+		RoleName: roleName,
+	})
+	s.NoError(err)
+	s.True(merr.Ok(resp13))
+
+	resp14, err := s.Cluster.Proxy.DeleteCredential(ctx, &milvuspb.DeleteCredentialRequest{
+		Username: userName,
+	})
+	s.NoError(err)
+	s.True(merr.Ok(resp14))
 }
 
 func TestRBACBackup(t *testing.T) {