RBAC custom privilege group

Signed-off-by: shaoting-huang <[email protected]>
milvus-io · Oct 23, 2024 · 7da2c9f · 7da2c9f
1 parent 6853e8c
commit 7da2c9f
Show file tree

Hide file tree

Showing 2 changed files with 113 additions and 6 deletions.
diff --git a/internal/rootcoord/meta_table.go b/internal/rootcoord/meta_table.go
@@ -99,9 +99,9 @@ type IMetaTable interface {
 	ListUserRole(tenant string) ([]string, error)
 	BackupRBAC(ctx context.Context, tenant string) (*milvuspb.RBACMeta, error)
 	RestoreRBAC(ctx context.Context, tenant string, meta *milvuspb.RBACMeta) error
-	CreatePrivilegeGroup(ctx context.Context, entity *milvuspb.PrivilegeGroupEntity) error
+	CreatePrivilegeGroup(ctx context.Context, groupName string, privileges []*milvuspb.PrivilegeEntity) error
 	DropPrivilegeGroup(ctx context.Context, groupName string) error
-	ListPrivilegeGroups(ctx context.Context) (*milvuspb.ListPrivilegeGroupsResponse, error)
+	ListPrivilegeGroups(ctx context.Context) ([]*milvuspb.PrivilegeGroupEntity, error)
 	AddPrivilegesToGroup(ctx context.Context, groupName string, privileges []*milvuspb.PrivilegeEntity) error
 	DropPrivilegesFromGroup(ctx context.Context, groupName string, privileges []*milvuspb.PrivilegeEntity) error
 }

diff --git a/internal/rootcoord/root_coord.go b/internal/rootcoord/root_coord.go
@@ -2932,17 +2932,124 @@ func (c *Core) CreatePrivilegeGroup(ctx context.Context, in *milvuspb.CreatePriv
 		return merr.Status(err), nil
 	}
 
-	err := c.meta.CreateRole(util.DefaultTenant, &milvuspb.RoleEntity{Name: entity.Name})
+	err := c.meta.CreatePrivilegeGroup(ctx, in.GroupName.Name, in.Privileges)
 	if err != nil {
-		errMsg := "fail to create role"
+		errMsg := "fail to create privilege group"
 		ctxLog.Warn(errMsg, zap.Error(err))
-		return merr.StatusWithErrorCode(err, commonpb.ErrorCode_CreateRoleFailure), nil
+		return merr.StatusWithErrorCode(err, commonpb.ErrorCode_CreatePrivilegeGroupFailure), nil
+	}
+
+	ctxLog.Debug(method + " success")
+	metrics.RootCoordDDLReqCounter.WithLabelValues(method, metrics.SuccessLabel).Inc()
+	metrics.RootCoordDDLReqLatency.WithLabelValues(method).Observe(float64(tr.ElapseSpan().Milliseconds()))
+	metrics.RootCoordNumOfPrivilegeGroups.Inc()
+
+	return merr.Success(), nil
+}
+
+func (c *Core) DropPrivilegeGroup(ctx context.Context, in *milvuspb.DropPrivilegeGroupRequest) (*commonpb.Status, error) {
+	method := "DropPrivilegeGroup"
+	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.Debug(method + " begin")
+
+	if err := merr.CheckHealthy(c.GetStateCode()); err != nil {
+		return merr.Status(err), nil
+	}
+
+	err := c.meta.DropPrivilegeGroup(ctx, in.GroupName.Name)
+	if err != nil {
+		errMsg := "fail to drop privilege group"
+		ctxLog.Warn(errMsg, zap.Error(err))
+		return merr.StatusWithErrorCode(err, commonpb.ErrorCode_CreatePrivilegeGroupFailure), nil
+	}
+
+	ctxLog.Debug(method + " success")
+	metrics.RootCoordDDLReqCounter.WithLabelValues(method, metrics.SuccessLabel).Inc()
+	metrics.RootCoordDDLReqLatency.WithLabelValues(method).Observe(float64(tr.ElapseSpan().Milliseconds()))
+	metrics.RootCoordNumOfPrivilegeGroups.Desc()
+
+	return merr.Success(), nil
+}
+
+func (c *Core) ListPrivilegeGroup(ctx context.Context, in *milvuspb.ListPrivilegeGroupsRequest) (*milvuspb.ListPrivilegeGroupsResponse, error) {
+	method := "ListPrivilegeGroup"
+	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.Debug(method + " begin")
+
+	if err := merr.CheckHealthy(c.GetStateCode()); err != nil {
+		return &milvuspb.ListPrivilegeGroupsResponse{
+			Status: merr.Status(err),
+		}, nil
+	}
+
+	groups, err := c.meta.ListPrivilegeGroups(ctx)
+	if err != nil {
+		errMsg := "fail to list privilege group"
+		ctxLog.Warn(errMsg, zap.Error(err))
+		return &milvuspb.ListPrivilegeGroupsResponse{
+			Status: merr.StatusWithErrorCode(err, commonpb.ErrorCode_ListPrivilegeGroupsFailure),
+		}, nil
+	}
+
+	ctxLog.Debug(method + " success")
+	metrics.RootCoordDDLReqCounter.WithLabelValues(method, metrics.SuccessLabel).Inc()
+	metrics.RootCoordDDLReqLatency.WithLabelValues(method).Observe(float64(tr.ElapseSpan().Milliseconds()))
+
+	return &milvuspb.ListPrivilegeGroupsResponse{
+		Status: merr.Success(),
+		Groups: groups,
+	}, nil
+}
+
+func (c *Core) AddPrivilegesToGroup(ctx context.Context, in *milvuspb.AddPrivilegesToGroupRequest) (*commonpb.Status, error) {
+	method := "AddPrivilegesToGroup"
+	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.Debug(method + " begin")
+
+	if err := merr.CheckHealthy(c.GetStateCode()); err != nil {
+		return merr.Status(err), nil
+	}
+
+	err := c.meta.AddPrivilegesToGroup(ctx, in.GroupName.Name, in.Privileges)
+	if err != nil {
+		errMsg := "fail to add privileges to group"
+		ctxLog.Warn(errMsg, zap.Error(err))
+		return merr.StatusWithErrorCode(err, commonpb.ErrorCode_AddPrivilegesToGroupFailure), nil
+	}
+
+	ctxLog.Debug(method + " success")
+	metrics.RootCoordDDLReqCounter.WithLabelValues(method, metrics.SuccessLabel).Inc()
+	metrics.RootCoordDDLReqLatency.WithLabelValues(method).Observe(float64(tr.ElapseSpan().Milliseconds()))
+
+	return merr.Success(), nil
+}
+
+func (c *Core) DropPrivilegesFromGroup(ctx context.Context, in *milvuspb.DropPrivilegesFromGroupRequest) (*commonpb.Status, error) {
+	method := "DropPrivilegesFromGroup"
+	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.Debug(method + " begin")
+	if err := merr.CheckHealthy(c.GetStateCode()); err != nil {
+		return merr.Status(err), nil
+	}
+
+	err := c.meta.DropPrivilegesFromGroup(ctx, in.GroupName.Name, in.Privileges)
+	if err != nil {
+		errMsg := "fail to drop privileges from group"
+		ctxLog.Warn(errMsg, zap.Error(err))
+		return merr.StatusWithErrorCode(err, commonpb.ErrorCode_DropPrivilegesFromGroupFailure), nil
 	}
 
 	ctxLog.Debug(method + " success")
 	metrics.RootCoordDDLReqCounter.WithLabelValues(method, metrics.SuccessLabel).Inc()
 	metrics.RootCoordDDLReqLatency.WithLabelValues(method).Observe(float64(tr.ElapseSpan().Milliseconds()))
-	metrics.RootCoordNumOfRoles.Inc()
 
 	return merr.Success(), nil
 }