Skip to content

Commit 893b88e

Browse files
SimFGSimFG
committed
enhance: [2.5] add privilege group message types and corresponding tests
Signed-off-by: SimFG <[email protected]>
1 parent 65dd398 commit 893b88e

File tree

4 files changed

+296
-0
lines changed

4 files changed

+296
-0
lines changed

‎internal/proxy/util.go

Lines changed: 15 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2257,6 +2257,21 @@ func SendReplicateMessagePack(ctx context.Context, replicateMsgStream msgstream.
22572257
BaseMsg: getBaseMsg(ctx, ts),
22582258
OperatePrivilegeV2Request: r,
22592259
}
2260+
case *milvuspb.CreatePrivilegeGroupRequest:
2261+
tsMsg = &msgstream.CreatePrivilegeGroupMsg{
2262+
BaseMsg: getBaseMsg(ctx, ts),
2263+
CreatePrivilegeGroupRequest: r,
2264+
}
2265+
case *milvuspb.DropPrivilegeGroupRequest:
2266+
tsMsg = &msgstream.DropPrivilegeGroupMsg{
2267+
BaseMsg: getBaseMsg(ctx, ts),
2268+
DropPrivilegeGroupRequest: r,
2269+
}
2270+
case *milvuspb.OperatePrivilegeGroupRequest:
2271+
tsMsg = &msgstream.OperatePrivilegeGroupMsg{
2272+
BaseMsg: getBaseMsg(ctx, ts),
2273+
OperatePrivilegeGroupRequest: r,
2274+
}
22602275
case *milvuspb.CreateAliasRequest:
22612276
tsMsg = &msgstream.CreateAliasMsg{
22622277
BaseMsg: getBaseMsg(ctx, ts),

‎pkg/mq/msgstream/msg_for_user_role.go

Lines changed: 159 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -447,3 +447,162 @@ func (c *OperatePrivilegeV2Msg) Unmarshal(input MarshalType) (TsMsg, error) {
447447
func (c *OperatePrivilegeV2Msg) Size() int {
448448
return proto.Size(c.OperatePrivilegeV2Request)
449449
}
450+
451+
type CreatePrivilegeGroupMsg struct {
452+
BaseMsg
453+
*milvuspb.CreatePrivilegeGroupRequest
454+
}
455+
456+
var _ TsMsg = &CreatePrivilegeGroupMsg{}
457+
458+
func (c *CreatePrivilegeGroupMsg) ID() UniqueID {
459+
return c.Base.MsgID
460+
}
461+
462+
func (c *CreatePrivilegeGroupMsg) SetID(id UniqueID) {
463+
c.Base.MsgID = id
464+
}
465+
466+
func (c *CreatePrivilegeGroupMsg) Type() MsgType {
467+
return c.Base.MsgType
468+
}
469+
470+
func (c *CreatePrivilegeGroupMsg) SourceID() int64 {
471+
return c.Base.SourceID
472+
}
473+
474+
func (c *CreatePrivilegeGroupMsg) Marshal(input TsMsg) (MarshalType, error) {
475+
createPrivilegeGroupMsg := input.(*CreatePrivilegeGroupMsg)
476+
createPrivilegeGroupRequest := createPrivilegeGroupMsg.CreatePrivilegeGroupRequest
477+
mb, err := proto.Marshal(createPrivilegeGroupRequest)
478+
if err != nil {
479+
return nil, err
480+
}
481+
return mb, nil
482+
}
483+
484+
func (c *CreatePrivilegeGroupMsg) Unmarshal(input MarshalType) (TsMsg, error) {
485+
createPrivilegeGroupRequest := &milvuspb.CreatePrivilegeGroupRequest{}
486+
in, err := convertToByteArray(input)
487+
if err != nil {
488+
return nil, err
489+
}
490+
err = proto.Unmarshal(in, createPrivilegeGroupRequest)
491+
if err != nil {
492+
return nil, err
493+
}
494+
createPrivilegeGroupMsg := &CreatePrivilegeGroupMsg{CreatePrivilegeGroupRequest: createPrivilegeGroupRequest}
495+
createPrivilegeGroupMsg.BeginTimestamp = createPrivilegeGroupMsg.GetBase().GetTimestamp()
496+
createPrivilegeGroupMsg.EndTimestamp = createPrivilegeGroupMsg.GetBase().GetTimestamp()
497+
return createPrivilegeGroupMsg, nil
498+
}
499+
500+
func (c *CreatePrivilegeGroupMsg) Size() int {
501+
return proto.Size(c.CreatePrivilegeGroupRequest)
502+
}
503+
504+
type DropPrivilegeGroupMsg struct {
505+
BaseMsg
506+
*milvuspb.DropPrivilegeGroupRequest
507+
}
508+
509+
var _ TsMsg = &DropPrivilegeGroupMsg{}
510+
511+
func (c *DropPrivilegeGroupMsg) ID() UniqueID {
512+
return c.Base.MsgID
513+
}
514+
515+
func (c *DropPrivilegeGroupMsg) SetID(id UniqueID) {
516+
c.Base.MsgID = id
517+
}
518+
519+
func (c *DropPrivilegeGroupMsg) Type() MsgType {
520+
return c.Base.MsgType
521+
}
522+
523+
func (c *DropPrivilegeGroupMsg) SourceID() int64 {
524+
return c.Base.SourceID
525+
}
526+
527+
func (c *DropPrivilegeGroupMsg) Marshal(input TsMsg) (MarshalType, error) {
528+
dropPrivilegeGroupMsg := input.(*DropPrivilegeGroupMsg)
529+
dropPrivilegeGroupRequest := dropPrivilegeGroupMsg.DropPrivilegeGroupRequest
530+
mb, err := proto.Marshal(dropPrivilegeGroupRequest)
531+
if err != nil {
532+
return nil, err
533+
}
534+
return mb, nil
535+
}
536+
537+
func (c *DropPrivilegeGroupMsg) Unmarshal(input MarshalType) (TsMsg, error) {
538+
dropPrivilegeGroupRequest := &milvuspb.DropPrivilegeGroupRequest{}
539+
in, err := convertToByteArray(input)
540+
if err != nil {
541+
return nil, err
542+
}
543+
err = proto.Unmarshal(in, dropPrivilegeGroupRequest)
544+
if err != nil {
545+
return nil, err
546+
}
547+
dropPrivilegeGroupMsg := &DropPrivilegeGroupMsg{DropPrivilegeGroupRequest: dropPrivilegeGroupRequest}
548+
dropPrivilegeGroupMsg.BeginTimestamp = dropPrivilegeGroupMsg.GetBase().GetTimestamp()
549+
dropPrivilegeGroupMsg.EndTimestamp = dropPrivilegeGroupMsg.GetBase().GetTimestamp()
550+
return dropPrivilegeGroupMsg, nil
551+
}
552+
553+
func (c *DropPrivilegeGroupMsg) Size() int {
554+
return proto.Size(c.DropPrivilegeGroupRequest)
555+
}
556+
557+
type OperatePrivilegeGroupMsg struct {
558+
BaseMsg
559+
*milvuspb.OperatePrivilegeGroupRequest
560+
}
561+
562+
var _ TsMsg = &OperatePrivilegeGroupMsg{}
563+
564+
func (c *OperatePrivilegeGroupMsg) ID() UniqueID {
565+
return c.Base.MsgID
566+
}
567+
568+
func (c *OperatePrivilegeGroupMsg) SetID(id UniqueID) {
569+
c.Base.MsgID = id
570+
}
571+
572+
func (c *OperatePrivilegeGroupMsg) Type() MsgType {
573+
return c.Base.MsgType
574+
}
575+
576+
func (c *OperatePrivilegeGroupMsg) SourceID() int64 {
577+
return c.Base.SourceID
578+
}
579+
580+
func (c *OperatePrivilegeGroupMsg) Marshal(input TsMsg) (MarshalType, error) {
581+
operatePrivilegeGroupMsg := input.(*OperatePrivilegeGroupMsg)
582+
operatePrivilegeGroupRequest := operatePrivilegeGroupMsg.OperatePrivilegeGroupRequest
583+
mb, err := proto.Marshal(operatePrivilegeGroupRequest)
584+
if err != nil {
585+
return nil, err
586+
}
587+
return mb, nil
588+
}
589+
590+
func (c *OperatePrivilegeGroupMsg) Unmarshal(input MarshalType) (TsMsg, error) {
591+
operatePrivilegeGroupRequest := &milvuspb.OperatePrivilegeGroupRequest{}
592+
in, err := convertToByteArray(input)
593+
if err != nil {
594+
return nil, err
595+
}
596+
err = proto.Unmarshal(in, operatePrivilegeGroupRequest)
597+
if err != nil {
598+
return nil, err
599+
}
600+
operatePrivilegeGroupMsg := &OperatePrivilegeGroupMsg{OperatePrivilegeGroupRequest: operatePrivilegeGroupRequest}
601+
operatePrivilegeGroupMsg.BeginTimestamp = operatePrivilegeGroupMsg.GetBase().GetTimestamp()
602+
operatePrivilegeGroupMsg.EndTimestamp = operatePrivilegeGroupMsg.GetBase().GetTimestamp()
603+
return operatePrivilegeGroupMsg, nil
604+
}
605+
606+
func (c *OperatePrivilegeGroupMsg) Size() int {
607+
return proto.Size(c.OperatePrivilegeGroupRequest)
608+
}

‎pkg/mq/msgstream/msg_for_user_role_test.go

Lines changed: 116 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -354,3 +354,119 @@ func TestOperatePrivilegeV2(t *testing.T) {
354354
assert.EqualValues(t, "unit_user", newMsg.(*OperatePrivilegeV2Msg).GetGrantor().GetUser().GetName())
355355
assert.EqualValues(t, "unit_privilege", newMsg.(*OperatePrivilegeV2Msg).GetGrantor().GetPrivilege().GetName())
356356
}
357+
358+
func TestCreatePrivilegeGroup(t *testing.T) {
359+
var msg TsMsg = &CreatePrivilegeGroupMsg{
360+
CreatePrivilegeGroupRequest: &milvuspb.CreatePrivilegeGroupRequest{
361+
Base: &commonpb.MsgBase{
362+
MsgType: commonpb.MsgType_CreatePrivilegeGroup,
363+
MsgID: 100,
364+
Timestamp: 1000,
365+
SourceID: 10000,
366+
TargetID: 100000,
367+
ReplicateInfo: nil,
368+
},
369+
GroupName: "unit_group",
370+
},
371+
}
372+
assert.EqualValues(t, 100, msg.ID())
373+
msg.SetID(200)
374+
assert.EqualValues(t, 200, msg.ID())
375+
assert.Equal(t, commonpb.MsgType_CreatePrivilegeGroup, msg.Type())
376+
assert.EqualValues(t, 10000, msg.SourceID())
377+
378+
msgBytes, err := msg.Marshal(msg)
379+
assert.NoError(t, err)
380+
381+
var newMsg TsMsg = &CreatePrivilegeGroupMsg{}
382+
_, err = newMsg.Unmarshal("1")
383+
assert.Error(t, err)
384+
385+
newMsg, err = newMsg.Unmarshal(msgBytes)
386+
assert.NoError(t, err)
387+
assert.EqualValues(t, 200, newMsg.ID())
388+
assert.EqualValues(t, 1000, newMsg.BeginTs())
389+
assert.EqualValues(t, 1000, newMsg.EndTs())
390+
assert.EqualValues(t, "unit_group", newMsg.(*CreatePrivilegeGroupMsg).GetGroupName())
391+
assert.EqualValues(t, commonpb.MsgType_CreatePrivilegeGroup, newMsg.Type())
392+
393+
assert.True(t, msg.Size() > 0)
394+
}
395+
396+
func TestDropPrivilegeGroup(t *testing.T) {
397+
var msg TsMsg = &DropPrivilegeGroupMsg{
398+
DropPrivilegeGroupRequest: &milvuspb.DropPrivilegeGroupRequest{
399+
Base: &commonpb.MsgBase{
400+
MsgType: commonpb.MsgType_DropPrivilegeGroup,
401+
MsgID: 100,
402+
Timestamp: 1000,
403+
SourceID: 10000,
404+
TargetID: 100000,
405+
ReplicateInfo: nil,
406+
},
407+
GroupName: "unit_group",
408+
},
409+
}
410+
assert.EqualValues(t, 100, msg.ID())
411+
msg.SetID(200)
412+
assert.EqualValues(t, 200, msg.ID())
413+
assert.Equal(t, commonpb.MsgType_DropPrivilegeGroup, msg.Type())
414+
assert.EqualValues(t, 10000, msg.SourceID())
415+
416+
msgBytes, err := msg.Marshal(msg)
417+
assert.NoError(t, err)
418+
419+
var newMsg TsMsg = &DropPrivilegeGroupMsg{}
420+
_, err = newMsg.Unmarshal("1")
421+
assert.Error(t, err)
422+
423+
newMsg, err = newMsg.Unmarshal(msgBytes)
424+
assert.NoError(t, err)
425+
assert.EqualValues(t, 200, newMsg.ID())
426+
assert.EqualValues(t, 1000, newMsg.BeginTs())
427+
assert.EqualValues(t, 1000, newMsg.EndTs())
428+
assert.EqualValues(t, "unit_group", newMsg.(*DropPrivilegeGroupMsg).GetGroupName())
429+
430+
assert.True(t, msg.Size() > 0)
431+
}
432+
433+
func TestOperatePrivilegeGroup(t *testing.T) {
434+
var msg TsMsg = &OperatePrivilegeGroupMsg{
435+
OperatePrivilegeGroupRequest: &milvuspb.OperatePrivilegeGroupRequest{
436+
Base: &commonpb.MsgBase{
437+
MsgType: commonpb.MsgType_OperatePrivilegeGroup,
438+
MsgID: 100,
439+
Timestamp: 1000,
440+
SourceID: 10000,
441+
TargetID: 100000,
442+
ReplicateInfo: nil,
443+
},
444+
GroupName: "unit_group",
445+
Type: milvuspb.OperatePrivilegeGroupType_AddPrivilegesToGroup,
446+
Privileges: []*milvuspb.PrivilegeEntity{
447+
{Name: "unit_privilege"},
448+
},
449+
},
450+
}
451+
assert.EqualValues(t, 100, msg.ID())
452+
msg.SetID(200)
453+
assert.EqualValues(t, 200, msg.ID())
454+
assert.Equal(t, commonpb.MsgType_OperatePrivilegeGroup, msg.Type())
455+
assert.EqualValues(t, 10000, msg.SourceID())
456+
457+
msgBytes, err := msg.Marshal(msg)
458+
assert.NoError(t, err)
459+
460+
var newMsg TsMsg = &OperatePrivilegeGroupMsg{}
461+
_, err = newMsg.Unmarshal("1")
462+
assert.Error(t, err)
463+
464+
newMsg, err = newMsg.Unmarshal(msgBytes)
465+
assert.NoError(t, err)
466+
assert.EqualValues(t, 200, newMsg.ID())
467+
assert.EqualValues(t, 1000, newMsg.BeginTs())
468+
assert.EqualValues(t, 1000, newMsg.EndTs())
469+
assert.EqualValues(t, "unit_group", newMsg.(*OperatePrivilegeGroupMsg).GetGroupName())
470+
assert.EqualValues(t, milvuspb.OperatePrivilegeGroupType_AddPrivilegesToGroup, newMsg.(*OperatePrivilegeGroupMsg).GetType())
471+
assert.EqualValues(t, "unit_privilege", newMsg.(*OperatePrivilegeGroupMsg).GetPrivileges()[0].GetName())
472+
}

‎pkg/mq/msgstream/unmarshal.go

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -89,6 +89,9 @@ func (pudf *ProtoUDFactory) NewUnmarshalDispatcher() *ProtoUnmarshalDispatcher {
8989
operateUserRoleMsg := OperateUserRoleMsg{}
9090
operatePrivilegeMsg := OperatePrivilegeMsg{}
9191
operatePrivilegeV2Msg := OperatePrivilegeV2Msg{}
92+
createPrivilegeGroupMsg := CreatePrivilegeGroupMsg{}
93+
dropPrivilegeGroupMsg := DropPrivilegeGroupMsg{}
94+
operatePrivilegeGroupMsg := OperatePrivilegeGroupMsg{}
9295
replicateMsg := ReplicateMsg{}
9396
importMsg := ImportMsg{}
9497

@@ -128,6 +131,9 @@ func (pudf *ProtoUDFactory) NewUnmarshalDispatcher() *ProtoUnmarshalDispatcher {
128131
p.TempMap[commonpb.MsgType_OperateUserRole] = operateUserRoleMsg.Unmarshal
129132
p.TempMap[commonpb.MsgType_OperatePrivilege] = operatePrivilegeMsg.Unmarshal
130133
p.TempMap[commonpb.MsgType_OperatePrivilegeV2] = operatePrivilegeV2Msg.Unmarshal
134+
p.TempMap[commonpb.MsgType_CreatePrivilegeGroup] = createPrivilegeGroupMsg.Unmarshal
135+
p.TempMap[commonpb.MsgType_DropPrivilegeGroup] = dropPrivilegeGroupMsg.Unmarshal
136+
p.TempMap[commonpb.MsgType_OperatePrivilegeGroup] = operatePrivilegeGroupMsg.Unmarshal
131137
p.TempMap[commonpb.MsgType_Replicate] = replicateMsg.Unmarshal
132138
p.TempMap[commonpb.MsgType_Import] = importMsg.Unmarshal
133139
p.TempMap[commonpb.MsgType_CreateAlias] = createAliasMsg.Unmarshal

0 commit comments

Comments
 (0)