Skip to content

Commit a01849e

Browse files
huanghaoyuanhhyhuanghaoyuanhhy
authored
*: support send restore rbac message (#884)
Signed-off-by: huanghaoyuanhhy <haoyuan.huang@zilliz.com>
1 parent 1a4fd47 commit a01849e

File tree

7 files changed

+345
-339
lines changed

7 files changed

+345
-339
lines changed

‎core/restore/conv/conv.go‎

Lines changed: 126 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -4,7 +4,9 @@ import (
44
"encoding/base64"
55
"fmt"
66

7+
"github.com/milvus-io/milvus-proto/go-api/v2/milvuspb"
78
"github.com/milvus-io/milvus-proto/go-api/v2/schemapb"
9+
"github.com/samber/lo"
810
"google.golang.org/protobuf/proto"
911

1012
"github.com/zilliztech/milvus-backup/core/proto/backuppb"
@@ -42,15 +44,15 @@ func Functions(bakFuncs []*backuppb.FunctionSchema) []*schemapb.FunctionSchema {
4244
funcs := make([]*schemapb.FunctionSchema, 0, len(bakFuncs))
4345
for _, bakFunc := range bakFuncs {
4446
fun := &schemapb.FunctionSchema{
45-
Name: bakFunc.Name,
46-
Id: bakFunc.Id,
47-
Description: bakFunc.Description,
48-
Type: schemapb.FunctionType(bakFunc.Type),
49-
InputFieldNames: bakFunc.InputFieldNames,
50-
InputFieldIds: bakFunc.InputFieldIds,
51-
OutputFieldNames: bakFunc.OutputFieldNames,
52-
OutputFieldIds: bakFunc.OutputFieldIds,
53-
Params: pbconv.BakKVToMilvusKV(bakFunc.Params),
47+
Name: bakFunc.GetName(),
48+
Id: bakFunc.GetId(),
49+
Description: bakFunc.GetDescription(),
50+
Type: schemapb.FunctionType(bakFunc.GetType()),
51+
InputFieldNames: bakFunc.GetInputFieldNames(),
52+
InputFieldIds: bakFunc.GetInputFieldIds(),
53+
OutputFieldNames: bakFunc.GetOutputFieldNames(),
54+
OutputFieldIds: bakFunc.GetOutputFieldIds(),
55+
Params: pbconv.BakKVToMilvusKV(bakFunc.GetParams()),
5456
}
5557
funcs = append(funcs, fun)
5658
}
@@ -138,3 +140,118 @@ func Schema(bakSchema *backuppb.CollectionSchema) (*schemapb.CollectionSchema, e
138140

139141
return schema, nil
140142
}
143+
144+
func Users(bacUsers []*backuppb.UserInfo, curUsers []*milvuspb.UserInfo) []*milvuspb.UserInfo {
145+
um := lo.SliceToMap(curUsers, func(user *milvuspb.UserInfo) (string, struct{}) {
146+
return user.User, struct{}{}
147+
})
148+
149+
users := make([]*milvuspb.UserInfo, 0, len(bacUsers))
150+
for _, user := range bacUsers {
151+
// skip if user already exist
152+
if _, ok := um[user.GetUser()]; ok {
153+
continue
154+
}
155+
156+
ur := lo.Map(user.GetRoles(), func(role *backuppb.RoleEntity, index int) *milvuspb.RoleEntity {
157+
return &milvuspb.RoleEntity{Name: role.Name}
158+
})
159+
userEntity := &milvuspb.UserInfo{
160+
User: user.User,
161+
Password: user.Password,
162+
Roles: ur,
163+
}
164+
users = append(users, userEntity)
165+
}
166+
167+
return users
168+
}
169+
170+
func Roles(bakRoles []*backuppb.RoleEntity, curRoles []*milvuspb.RoleEntity) []*milvuspb.RoleEntity {
171+
rm := lo.SliceToMap(curRoles, func(role *milvuspb.RoleEntity) (string, struct{}) {
172+
return role.Name, struct{}{}
173+
})
174+
175+
roles := make([]*milvuspb.RoleEntity, 0, len(bakRoles))
176+
for _, role := range bakRoles {
177+
// skip if role already exist
178+
if _, ok := rm[role.GetName()]; ok {
179+
continue
180+
}
181+
182+
roleEntity := &milvuspb.RoleEntity{
183+
Name: role.GetName(),
184+
}
185+
roles = append(roles, roleEntity)
186+
}
187+
188+
return roles
189+
}
190+
191+
func backupGrantKey(grant *backuppb.GrantEntity) string {
192+
return fmt.Sprintf("%s/%s/%s/%s/%s/%s",
193+
grant.GetObject().GetName(),
194+
grant.GetObjectName(),
195+
grant.GetRole().GetName(),
196+
grant.GetGrantor().GetUser().GetName(),
197+
grant.GetGrantor().GetPrivilege().GetName(),
198+
grant.GetDbName())
199+
}
200+
201+
func milvusGrantKey(grant *milvuspb.GrantEntity) string {
202+
return fmt.Sprintf("%s/%s/%s/%s/%s/%s",
203+
grant.GetObject().GetName(),
204+
grant.GetObjectName(),
205+
grant.GetRole().GetName(),
206+
grant.GetGrantor().GetUser().GetName(),
207+
grant.GetGrantor().GetPrivilege().GetName(),
208+
grant.GetDbName())
209+
}
210+
211+
func Grants(bakGrants []*backuppb.GrantEntity, curGrants []*milvuspb.GrantEntity) []*milvuspb.GrantEntity {
212+
gm := lo.SliceToMap(curGrants, func(grant *milvuspb.GrantEntity) (string, struct{}) {
213+
return milvusGrantKey(grant), struct{}{}
214+
})
215+
grants := make([]*milvuspb.GrantEntity, 0, len(bakGrants))
216+
for _, grant := range bakGrants {
217+
// skip if grant already exist
218+
if _, ok := gm[backupGrantKey(grant)]; ok {
219+
continue
220+
}
221+
222+
grantEntity := &milvuspb.GrantEntity{
223+
Role: &milvuspb.RoleEntity{Name: grant.GetRole().GetName()},
224+
Object: &milvuspb.ObjectEntity{Name: grant.GetObject().GetName()},
225+
ObjectName: grant.GetObjectName(),
226+
Grantor: &milvuspb.GrantorEntity{
227+
User: &milvuspb.UserEntity{Name: grant.GetGrantor().GetUser().GetName()},
228+
Privilege: &milvuspb.PrivilegeEntity{Name: grant.GetGrantor().GetPrivilege().GetName()},
229+
},
230+
DbName: grant.GetDbName(),
231+
}
232+
grants = append(grants, grantEntity)
233+
}
234+
235+
return grants
236+
}
237+
238+
func PrivilegeGroups(bakGroups []*backuppb.PrivilegeGroupInfo, curGroups []*milvuspb.PrivilegeGroupInfo) []*milvuspb.PrivilegeGroupInfo {
239+
curGroupName := lo.SliceToMap(curGroups, func(group *milvuspb.PrivilegeGroupInfo) (string, struct{}) {
240+
return group.GetGroupName(), struct{}{}
241+
})
242+
243+
groups := make([]*milvuspb.PrivilegeGroupInfo, 0, len(bakGroups))
244+
for _, group := range bakGroups {
245+
if _, ok := curGroupName[group.GetGroupName()]; ok {
246+
continue
247+
}
248+
249+
privileges := lo.Map(group.GetPrivileges(), func(privilege *backuppb.PrivilegeEntity, _ int) *milvuspb.PrivilegeEntity {
250+
return &milvuspb.PrivilegeEntity{Name: privilege.GetName()}
251+
})
252+
g := &milvuspb.PrivilegeGroupInfo{GroupName: group.GetGroupName(), Privileges: privileges}
253+
groups = append(groups, g)
254+
}
255+
256+
return groups
257+
}

‎core/restore/conv/conv_test.go‎

Lines changed: 137 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -4,6 +4,7 @@ import (
44
"encoding/base64"
55
"testing"
66

7+
"github.com/milvus-io/milvus-proto/go-api/v2/milvuspb"
78
"github.com/milvus-io/milvus-proto/go-api/v2/schemapb"
89
"github.com/stretchr/testify/assert"
910
"google.golang.org/protobuf/proto"
@@ -42,3 +43,139 @@ func TestDefaultValue(t *testing.T) {
4243
assert.Nil(t, val)
4344
})
4445
}
46+
47+
func TestBackupGrantKey(t *testing.T) {
48+
grant := &backuppb.GrantEntity{
49+
Role: &backuppb.RoleEntity{Name: "role.name"},
50+
Object: &backuppb.ObjectEntity{Name: "object.name"},
51+
ObjectName: "objectName",
52+
Grantor: &backuppb.GrantorEntity{
53+
User: &backuppb.UserEntity{Name: "grantor.user.name"},
54+
Privilege: &backuppb.PrivilegeEntity{Name: "grantor.privilege.name"},
55+
},
56+
DbName: "dbName",
57+
}
58+
59+
grantKey := backupGrantKey(grant)
60+
expected := "object.name/objectName/role.name/grantor.user.name/grantor.privilege.name/dbName"
61+
assert.Equal(t, expected, grantKey)
62+
}
63+
64+
func TestMilvusGrantKey(t *testing.T) {
65+
grant := &milvuspb.GrantEntity{
66+
Role: &milvuspb.RoleEntity{Name: "role.name"},
67+
Object: &milvuspb.ObjectEntity{Name: "object.name"},
68+
ObjectName: "objectName",
69+
Grantor: &milvuspb.GrantorEntity{
70+
User: &milvuspb.UserEntity{Name: "grantor.user.name"},
71+
Privilege: &milvuspb.PrivilegeEntity{Name: "grantor.privilege.name"},
72+
},
73+
DbName: "dbName",
74+
}
75+
76+
grantKey := milvusGrantKey(grant)
77+
expected := "object.name/objectName/role.name/grantor.user.name/grantor.privilege.name/dbName"
78+
assert.Equal(t, expected, grantKey)
79+
}
80+
81+
func TestGrants(t *testing.T) {
82+
bakGrants := []*backuppb.GrantEntity{
83+
{
84+
Role: &backuppb.RoleEntity{Name: "role.name"},
85+
Object: &backuppb.ObjectEntity{Name: "object.name"},
86+
ObjectName: "objectName",
87+
Grantor: &backuppb.GrantorEntity{
88+
User: &backuppb.UserEntity{Name: "grantor.user.name"},
89+
Privilege: &backuppb.PrivilegeEntity{Name: "grantor.privilege.name"},
90+
},
91+
DbName: "dbName",
92+
},
93+
{
94+
Role: &backuppb.RoleEntity{Name: "role1.name"}, // different role
95+
Object: &backuppb.ObjectEntity{Name: "object.name"},
96+
ObjectName: "objectName",
97+
Grantor: &backuppb.GrantorEntity{
98+
User: &backuppb.UserEntity{Name: "grantor.user.name"},
99+
Privilege: &backuppb.PrivilegeEntity{Name: "grantor.privilege.name"},
100+
},
101+
DbName: "dbName",
102+
},
103+
}
104+
105+
curGrants := []*milvuspb.GrantEntity{
106+
{
107+
Role: &milvuspb.RoleEntity{Name: "role.name"},
108+
Object: &milvuspb.ObjectEntity{Name: "object.name"},
109+
ObjectName: "objectName",
110+
Grantor: &milvuspb.GrantorEntity{
111+
User: &milvuspb.UserEntity{Name: "grantor.user.name"},
112+
Privilege: &milvuspb.PrivilegeEntity{Name: "grantor.privilege.name"},
113+
},
114+
DbName: "dbName",
115+
},
116+
}
117+
118+
restoreGrants := Grants(bakGrants, curGrants)
119+
assert.Len(t, restoreGrants, 1)
120+
}
121+
122+
func TestRoles(t *testing.T) {
123+
bakRoles := []*backuppb.RoleEntity{
124+
{Name: "role.name"},
125+
{Name: "role1.name"},
126+
}
127+
128+
curRoles := []*milvuspb.RoleEntity{
129+
{Name: "role.name"},
130+
}
131+
132+
restoreRoles := Roles(bakRoles, curRoles)
133+
assert.Len(t, restoreRoles, 1)
134+
assert.Equal(t, "role1.name", restoreRoles[0].Name)
135+
}
136+
137+
func TestUsers(t *testing.T) {
138+
bakUsers := []*backuppb.UserInfo{
139+
{User: "user.user"},
140+
{User: "user1.user"},
141+
}
142+
143+
curUsers := []*milvuspb.UserInfo{
144+
{User: "user.user"},
145+
}
146+
147+
restoreUsers := Users(bakUsers, curUsers)
148+
assert.Len(t, restoreUsers, 1)
149+
assert.Equal(t, "user1.user", restoreUsers[0].User)
150+
}
151+
152+
func TestPrivilegeGroups(t *testing.T) {
153+
bakPG := []*backuppb.PrivilegeGroupInfo{{
154+
GroupName: "group1",
155+
Privileges: []*backuppb.PrivilegeEntity{
156+
{Name: "privilege1"},
157+
{Name: "privilege2"},
158+
},
159+
}, {
160+
GroupName: "group2",
161+
Privileges: []*backuppb.PrivilegeEntity{
162+
{Name: "privilege3"},
163+
{Name: "privilege4"},
164+
},
165+
}}
166+
167+
curPrivilegeGroups := []*milvuspb.PrivilegeGroupInfo{{
168+
GroupName: "group1",
169+
Privileges: []*milvuspb.PrivilegeEntity{
170+
{Name: "privilege1"},
171+
{Name: "privilege2"},
172+
},
173+
}}
174+
175+
restorePG := PrivilegeGroups(bakPG, curPrivilegeGroups)
176+
assert.Len(t, restorePG, 1)
177+
assert.Equal(t, "group2", restorePG[0].GroupName)
178+
assert.Equal(t, 2, len(restorePG[0].Privileges))
179+
assert.Equal(t, "privilege3", restorePG[0].Privileges[0].Name)
180+
assert.Equal(t, "privilege4", restorePG[0].Privileges[1].Name)
181+
}

0 commit comments

Comments
 (0)