Skip to content

Commit 9a11573

Browse files
authored
Add resetConnection for updatePassword(), add isL0 for compact() (#1868)
Signed-off-by: yhmo <yihua.mo@zilliz.com>
1 parent 4061e20 commit 9a11573

5 files changed

Lines changed: 316 additions & 1 deletion

File tree

Lines changed: 273 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,273 @@
1+
/*
2+
* Licensed to the Apache Software Foundation (ASF) under one
3+
* or more contributor license agreements. See the NOTICE file
4+
* distributed with this work for additional information
5+
* regarding copyright ownership. The ASF licenses this file
6+
* to you under the Apache License, Version 2.0 (the
7+
* "License"); you may not use this file except in compliance
8+
* with the License. You may obtain a copy of the License at
9+
*
10+
* http://www.apache.org/licenses/LICENSE-2.0
11+
*
12+
* Unless required by applicable law or agreed to in writing,
13+
* software distributed under the License is distributed on an
14+
* "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
15+
* KIND, either express or implied. See the License for the
16+
* specific language governing permissions and limitations
17+
* under the License.
18+
*/
19+
20+
package io.milvus.v2;
21+
22+
import io.milvus.v2.client.ConnectConfig;
23+
import io.milvus.v2.client.MilvusClientV2;
24+
import io.milvus.v2.service.rbac.PrivilegeGroup;
25+
import io.milvus.v2.service.rbac.request.AddPrivilegesToGroupReq;
26+
import io.milvus.v2.service.rbac.request.CreatePrivilegeGroupReq;
27+
import io.milvus.v2.service.rbac.request.CreateRoleReq;
28+
import io.milvus.v2.service.rbac.request.CreateUserReq;
29+
import io.milvus.v2.service.rbac.request.DescribeRoleReq;
30+
import io.milvus.v2.service.rbac.request.DescribeUserReq;
31+
import io.milvus.v2.service.rbac.request.DropPrivilegeGroupReq;
32+
import io.milvus.v2.service.rbac.request.DropRoleReq;
33+
import io.milvus.v2.service.rbac.request.DropUserReq;
34+
import io.milvus.v2.service.rbac.request.GrantPrivilegeReq;
35+
import io.milvus.v2.service.rbac.request.GrantPrivilegeReqV2;
36+
import io.milvus.v2.service.rbac.request.GrantRoleReq;
37+
import io.milvus.v2.service.rbac.request.ListPrivilegeGroupsReq;
38+
import io.milvus.v2.service.rbac.request.RemovePrivilegesFromGroupReq;
39+
import io.milvus.v2.service.rbac.request.RevokePrivilegeReq;
40+
import io.milvus.v2.service.rbac.request.RevokePrivilegeReqV2;
41+
import io.milvus.v2.service.rbac.request.RevokeRoleReq;
42+
import io.milvus.v2.service.rbac.request.UpdatePasswordReq;
43+
import io.milvus.v2.service.rbac.response.DescribeRoleResp;
44+
import io.milvus.v2.service.rbac.response.DescribeUserResp;
45+
import io.milvus.v2.service.rbac.response.ListPrivilegeGroupsResp;
46+
47+
import java.util.Arrays;
48+
import java.util.List;
49+
50+
// Before running this example, make sure you have enabled RBAC in Milvus and have the root user credentials ready.
51+
// Read this doc to know how to enable RBAC in milvus: https://milvus.io/docs/authenticate.md?tab=docker
52+
public class RBACExample {
53+
private static final String URI = "http://localhost:19530";
54+
private static final String ROOT_USER = "root";
55+
private static final String ROOT_PASSWORD = "Milvus";
56+
private static final String DEMO_USER = "java_sdk_example_user";
57+
private static final String DEMO_PASSWORD = "Milvus@2026";
58+
private static final String DEMO_PASSWORD_UPDATED = "Milvus@2027";
59+
private static final String DEMO_ROLE = "java_sdk_example_role";
60+
private static final String DEMO_PRIVILEGE_GROUP = "java_sdk_example_privilege_group";
61+
private static final String DATABASE_NAME = "default";
62+
private static final String COLLECTION_NAME = "*";
63+
64+
public static void main(String[] args) {
65+
ConnectConfig config = ConnectConfig.builder()
66+
.uri(URI)
67+
.username(ROOT_USER)
68+
.password(ROOT_PASSWORD)
69+
.build();
70+
MilvusClientV2 client = new MilvusClientV2(config);
71+
72+
try {
73+
cleanup(client);
74+
runExample(client);
75+
} finally {
76+
cleanup(client);
77+
client.close();
78+
}
79+
}
80+
81+
private static void runExample(MilvusClientV2 client) {
82+
print("Initial users", client.listUsers());
83+
print("Initial roles", client.listRoles());
84+
85+
client.createUser(CreateUserReq.builder()
86+
.userName(DEMO_USER)
87+
.password(DEMO_PASSWORD)
88+
.build());
89+
print("After create user", client.listUsers());
90+
91+
DescribeUserResp userResp = client.describeUser(DescribeUserReq.builder()
92+
.userName(DEMO_USER)
93+
.build());
94+
print("Describe user", userResp);
95+
96+
client.updatePassword(UpdatePasswordReq.builder()
97+
.userName(DEMO_USER)
98+
.password(DEMO_PASSWORD)
99+
.newPassword(DEMO_PASSWORD_UPDATED)
100+
.resetConnection(false)
101+
.build());
102+
print("Password updated", DEMO_USER);
103+
104+
client.createRole(CreateRoleReq.builder()
105+
.roleName(DEMO_ROLE)
106+
.build());
107+
print("After create role", client.listRoles());
108+
109+
client.grantPrivilege(GrantPrivilegeReq.builder()
110+
.roleName(DEMO_ROLE)
111+
.objectType("Global")
112+
.objectName("*")
113+
.privilege("CreateCollection")
114+
.build());
115+
client.grantPrivilege(GrantPrivilegeReq.builder()
116+
.roleName(DEMO_ROLE)
117+
.objectType("Collection")
118+
.objectName("*")
119+
.privilege("Search")
120+
.build());
121+
122+
client.createPrivilegeGroup(CreatePrivilegeGroupReq.builder()
123+
.groupName(DEMO_PRIVILEGE_GROUP)
124+
.build());
125+
client.addPrivilegesToGroup(AddPrivilegesToGroupReq.builder()
126+
.groupName(DEMO_PRIVILEGE_GROUP)
127+
.privileges(Arrays.asList("Query", "Load"))
128+
.build());
129+
client.removePrivilegesFromGroup(RemovePrivilegesFromGroupReq.builder()
130+
.groupName(DEMO_PRIVILEGE_GROUP)
131+
.privileges(Arrays.asList("Load"))
132+
.build());
133+
134+
ListPrivilegeGroupsResp groupsResp = client.listPrivilegeGroups(ListPrivilegeGroupsReq.builder().build());
135+
for (PrivilegeGroup group : groupsResp.getPrivilegeGroups()) {
136+
if (DEMO_PRIVILEGE_GROUP.equals(group.getGroupName())) {
137+
print("Privilege group", group);
138+
}
139+
}
140+
141+
client.grantPrivilegeV2(GrantPrivilegeReqV2.builder()
142+
.roleName(DEMO_ROLE)
143+
.privilege(DEMO_PRIVILEGE_GROUP)
144+
.dbName(DATABASE_NAME)
145+
.collectionName(COLLECTION_NAME)
146+
.build());
147+
148+
DescribeRoleResp roleResp = client.describeRole(DescribeRoleReq.builder()
149+
.roleName(DEMO_ROLE)
150+
.dbName(DATABASE_NAME)
151+
.build());
152+
print("Describe role", roleResp);
153+
154+
client.grantRole(GrantRoleReq.builder()
155+
.userName(DEMO_USER)
156+
.roleName(DEMO_ROLE)
157+
.build());
158+
print("User after grant role", client.describeUser(DescribeUserReq.builder()
159+
.userName(DEMO_USER)
160+
.build()));
161+
162+
MilvusClientV2 userClient = new MilvusClientV2(ConnectConfig.builder()
163+
.uri(URI)
164+
.username(DEMO_USER)
165+
.password(DEMO_PASSWORD_UPDATED)
166+
.build());
167+
try {
168+
print("Updated user client ready", userClient.clientIsReady());
169+
print("Server version", userClient.getServerVersion());
170+
} finally {
171+
userClient.close();
172+
}
173+
174+
client.revokeRole(RevokeRoleReq.builder()
175+
.userName(DEMO_USER)
176+
.roleName(DEMO_ROLE)
177+
.build());
178+
print("User after revoke role", client.describeUser(DescribeUserReq.builder()
179+
.userName(DEMO_USER)
180+
.build()));
181+
182+
cleanup(client);
183+
print("Final users", client.listUsers());
184+
print("Final roles", client.listRoles());
185+
}
186+
187+
private static void cleanup(MilvusClientV2 client) {
188+
if (userHasRole(client, DEMO_USER, DEMO_ROLE)) {
189+
client.revokeRole(RevokeRoleReq.builder()
190+
.userName(DEMO_USER)
191+
.roleName(DEMO_ROLE)
192+
.build());
193+
}
194+
195+
if (client.listRoles().contains(DEMO_ROLE)) {
196+
DescribeRoleResp roleResp = client.describeRole(DescribeRoleReq.builder()
197+
.roleName(DEMO_ROLE)
198+
.dbName(DATABASE_NAME)
199+
.build());
200+
for (DescribeRoleResp.GrantInfo grantInfo : roleResp.getGrantInfos()) {
201+
if (DEMO_PRIVILEGE_GROUP.equals(grantInfo.getPrivilege())) {
202+
client.revokePrivilegeV2(RevokePrivilegeReqV2.builder()
203+
.roleName(DEMO_ROLE)
204+
.privilege(DEMO_PRIVILEGE_GROUP)
205+
.dbName(DATABASE_NAME)
206+
.collectionName(COLLECTION_NAME)
207+
.build());
208+
} else if ("CreateCollection".equals(grantInfo.getPrivilege())) {
209+
client.revokePrivilege(RevokePrivilegeReq.builder()
210+
.roleName(DEMO_ROLE)
211+
.objectType("Global")
212+
.objectName("*")
213+
.privilege("CreateCollection")
214+
.build());
215+
} else if ("Search".equals(grantInfo.getPrivilege())) {
216+
client.revokePrivilege(RevokePrivilegeReq.builder()
217+
.roleName(DEMO_ROLE)
218+
.objectType("Collection")
219+
.objectName("*")
220+
.privilege("Search")
221+
.build());
222+
}
223+
}
224+
225+
client.dropRole(DropRoleReq.builder()
226+
.roleName(DEMO_ROLE)
227+
.forceDrop(true)
228+
.build());
229+
}
230+
231+
if (client.listUsers().contains(DEMO_USER)) {
232+
client.dropUser(DropUserReq.builder()
233+
.userName(DEMO_USER)
234+
.build());
235+
}
236+
237+
if (privilegeGroupExists(client, DEMO_PRIVILEGE_GROUP)) {
238+
client.removePrivilegesFromGroup(RemovePrivilegesFromGroupReq.builder()
239+
.groupName(DEMO_PRIVILEGE_GROUP)
240+
.privileges(Arrays.asList("Query", "Load"))
241+
.build());
242+
client.dropPrivilegeGroup(DropPrivilegeGroupReq.builder()
243+
.groupName(DEMO_PRIVILEGE_GROUP)
244+
.build());
245+
}
246+
}
247+
248+
private static boolean userHasRole(MilvusClientV2 client, String userName, String roleName) {
249+
if (!client.listUsers().contains(userName)) {
250+
return false;
251+
}
252+
253+
DescribeUserResp userResp = client.describeUser(DescribeUserReq.builder()
254+
.userName(userName)
255+
.build());
256+
List<String> roles = userResp.getRoles();
257+
return roles != null && roles.contains(roleName);
258+
}
259+
260+
private static boolean privilegeGroupExists(MilvusClientV2 client, String groupName) {
261+
ListPrivilegeGroupsResp groupsResp = client.listPrivilegeGroups(ListPrivilegeGroupsReq.builder().build());
262+
for (PrivilegeGroup group : groupsResp.getPrivilegeGroups()) {
263+
if (groupName.equals(group.getGroupName())) {
264+
return true;
265+
}
266+
}
267+
return false;
268+
}
269+
270+
private static void print(String label, Object value) {
271+
System.out.printf("%-32s%s%n", label + ":", value);
272+
}
273+
}

sdk-core/src/main/java/io/milvus/v2/client/MilvusClientV2.java

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -938,6 +938,13 @@ public void createUser(CreateUserReq request) {
938938
*/
939939
public void updatePassword(UpdatePasswordReq request) {
940940
rpcUtils.retry(() -> rbacService.updatePassword(this.getRpcStub(), request));
941+
if (Boolean.TRUE.equals(request.getResetConnection()) && connectConfig != null) {
942+
connectConfig.setUsername(request.getUserName());
943+
connectConfig.setPassword(request.getNewPassword());
944+
connectConfig.setToken(null);
945+
connect(connectConfig);
946+
initServices(connectConfig.getDbName());
947+
}
941948
}
942949

943950
/**

sdk-core/src/main/java/io/milvus/v2/service/rbac/request/UpdatePasswordReq.java

Lines changed: 17 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -23,11 +23,13 @@ public class UpdatePasswordReq {
2323
private String userName;
2424
private String password;
2525
private String newPassword;
26+
private Boolean resetConnection;
2627

2728
private UpdatePasswordReq(UpdatePasswordReqBuilder builder) {
2829
this.userName = builder.userName;
2930
this.password = builder.password;
3031
this.newPassword = builder.newPassword;
32+
this.resetConnection = builder.resetConnection;
3133
}
3234

3335
public String getUserName() {
@@ -54,12 +56,21 @@ public void setNewPassword(String newPassword) {
5456
this.newPassword = newPassword;
5557
}
5658

59+
public Boolean getResetConnection() {
60+
return resetConnection;
61+
}
62+
63+
public void setResetConnection(Boolean resetConnection) {
64+
this.resetConnection = resetConnection;
65+
}
66+
5767
@Override
5868
public String toString() {
5969
return "UpdatePasswordReq{" +
6070
"userName='" + userName + '\'' +
6171
", password='" + password + '\'' +
6272
", newPassword='" + newPassword + '\'' +
73+
", resetConnection=" + resetConnection +
6374
'}';
6475
}
6576

@@ -71,6 +82,7 @@ public static class UpdatePasswordReqBuilder {
7182
private String userName;
7283
private String password;
7384
private String newPassword;
85+
private Boolean resetConnection = Boolean.FALSE;
7486

7587
private UpdatePasswordReqBuilder() {
7688
}
@@ -90,6 +102,11 @@ public UpdatePasswordReqBuilder newPassword(String newPassword) {
90102
return this;
91103
}
92104

105+
public UpdatePasswordReqBuilder resetConnection(Boolean resetConnection) {
106+
this.resetConnection = resetConnection;
107+
return this;
108+
}
109+
93110
public UpdatePasswordReq build() {
94111
return new UpdatePasswordReq(this);
95112
}

sdk-core/src/main/java/io/milvus/v2/service/utility/UtilityService.java

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -199,7 +199,8 @@ public CompactResp compact(MilvusServiceGrpc.MilvusServiceBlockingStub blockingS
199199

200200
ManualCompactionRequest.Builder builder = ManualCompactionRequest.newBuilder()
201201
.setCollectionID(descResponse.getCollectionID())
202-
.setMajorCompaction(request.getIsClustering());
202+
.setMajorCompaction(request.getIsClustering())
203+
.setL0Compaction(request.getIsL0());
203204
if (StringUtils.isNotEmpty(dbName)) {
204205
builder.setDbName(dbName);
205206
}

0 commit comments

Comments
 (0)