Add grantPrivilegeV2 and revokePrivilegeV2 (#395)

* Add grantPrivilegeV2 and revokePrivilegeV2 Signed-off-by: ryjiang <[email protected]> * disable debug Signed-off-by: ryjiang <[email protected]> --------- Signed-off-by: ryjiang <[email protected]>
milvus-io · Dec 24, 2024 · 1a962f9 · 1a962f9
1 parent 1ffa69d
commit 1a962f9
Show file tree

Hide file tree

Showing 3 changed files with 125 additions and 44 deletions.
diff --git a/milvus/grpc/User.ts b/milvus/grpc/User.ts
@@ -23,6 +23,9 @@ import {
   OperatePrivilegeGroupReq,
   AddPrivilegesToGroupReq,
   RemovePrivilegesFromGroupReq,
+  OperatePrivilegeV2Request,
+  GrantPrivilegeV2Request,
+  RevokePrivilegeV2Request,
   OperatePrivilegeGroupType,
   GrpcTimeOut,
   ListCredUsersResponse,
@@ -774,37 +777,6 @@ export class User extends Resource {
     return promise;
   }
 
-  /**
-   * Operate a privilege group in Milvus.
-   * @param {OperatePrivilegeGroupReq} data - The privilege group data.
-   * @param {string} data.group_name - The name of the privilege group to be operated.
-   * @param {PrivilegeEntity[]} data.privileges - The privileges to be operated.
-   * @param {OperatePrivilegeGroupType} data.type - The operation type.
-   * @param {number} [data.timeout] - An optional duration of time in milliseconds to allow for the RPC. If it is set to undefined, the client keeps waiting until the server responds or error occurs. Default is undefined.
-   * @returns {Promise<any>} The response object.
-   *
-   * @example
-   * ```javascript
-   * await milvusClient.operatePrivilegeGroup({
-   * group_name: 'exampleGroup',
-   * privileges: [{name: 'CreateCollection'}],
-   * type: OperatePrivilegeGroupType.AddPrivilegesToGroup,
-   * });
-   * ```
-   */
-  async operatePrivilegeGroup(
-    data: OperatePrivilegeGroupReq
-  ): Promise<ResStatus> {
-    const promise = await promisify(
-      this.channelPool,
-      'OperatePrivilegeGroup',
-      data,
-      data.timeout || this.timeout
-    );
-
-    return promise;
-  }
-
   /**
    * add privileges to a privilege group in Milvus.
    * @param {AddPrivilegesToGroupReq} data - The privilege group data.
@@ -878,4 +850,76 @@ export class User extends Resource {
 
     return promise;
   }
+
+  /**
+   * Grant a privilege to a role in Milvus.
+   * @param {GrantPrivilegeV2Request} data - The privilege data.
+   * @param {string} data.role - The name of the role.
+   * @param {string} data.privilege - The name of the privilege.
+   * @param {string} data.db_name - The name of the database.
+   * @param {string} data.collection_name - The name of the collection.
+   * @returns {Promise<ResStatus>} The response object.
+   *
+   * @example
+   * ```javascript
+   * await milvusClient.grantPrivilegeV2({
+   *  role: 'exampleRole',
+   *  privilege: 'CreateCollection',
+   *  db_name: 'exampleDB',
+   *  collection_name: 'exampleCollection',
+   * });
+   * ```
+   */
+  async grantPrivilegeV2(data: GrantPrivilegeV2Request): Promise<ResStatus> {
+    const promise = await promisify(
+      this.channelPool,
+      'OperatePrivilegeV2',
+      {
+        role: { name: data.role },
+        grantor: { privilege: { name: data.privilege } },
+        type: OperatePrivilegeType.Grant,
+        db_name: data.db_name,
+        collection_name: data.collection_name,
+      },
+      data.timeout || this.timeout
+    );
+
+    return promise;
+  }
+
+  /**
+   * Revoke a privilege from a role in Milvus.
+   * @param {RevokePrivilegeV2Request} data - The privilege data.
+   * @param {string} data.role - The name of the role.
+   * @param {string} data.privilege - The name of the privilege.
+   * @param {string} data.db_name - The name of the database.
+   * @param {string} data.collection_name - The name of the collection.
+   * @returns {Promise<ResStatus>} The response object.
+   *
+   * @example
+   * ```javascript
+   * await milvusClient.revokePrivilegeV2({
+   *  role: 'exampleRole',
+   *  privilege: 'CreateCollection',
+   *  db_name: 'exampleDB',
+   *  collection_name: 'exampleCollection',
+   * });
+   * ```
+   * */
+  async revokePrivilegeV2(data: RevokePrivilegeV2Request): Promise<ResStatus> {
+    const promise = await promisify(
+      this.channelPool,
+      'OperatePrivilegeV2',
+      {
+        role: { name: data.role },
+        grantor: { privilege: { name: data.privilege } },
+        type: OperatePrivilegeType.Revoke,
+        db_name: data.db_name,
+        collection_name: data.collection_name,
+      },
+      data.timeout || this.timeout
+    );
+
+    return promise;
+  }
 }
diff --git a/milvus/types/User.ts b/milvus/types/User.ts
@@ -1,5 +1,9 @@
 import { GrpcTimeOut, PrivilegesTypes, resStatusResponse } from './Common';
-import { RbacObjects, OperatePrivilegeGroupType } from '../';
+import {
+  RbacObjects,
+  OperatePrivilegeGroupType,
+  OperatePrivilegeType,
+} from '../';
 
 // base
 export interface usernameReq extends GrpcTimeOut {
@@ -85,6 +89,24 @@ export interface HasRoleResponse extends resStatusResponse {
   hasRole: boolean;
 }
 
+// V2
+export interface OperatePrivilegeV2Request extends GrpcTimeOut {
+  role: RoleEntity; // required, role
+  grantor: Grantor; // required, grantor
+  type: OperatePrivilegeType; // required, operation type
+  db_name: string; // required, db name
+  collection_name: string; // required, collection name
+}
+
+export interface GrantPrivilegeV2Request extends GrpcTimeOut {
+  role: string; // required, string
+  privilege: PrivilegesTypes; // required, privilege
+  db_name: string; // required, db name or '*'
+  collection_name: string; // required, collection name or '*'
+}
+
+export interface RevokePrivilegeV2Request extends GrantPrivilegeV2Request {}
+
 export interface CreatePrivilegeGroupReq extends GrpcTimeOut {
   group_name: string; // required, name
 }

diff --git a/test/grpc/User.spec.ts b/test/grpc/User.spec.ts
@@ -1,4 +1,3 @@
-import exp from 'constants';
 import {
   MilvusClient,
   ERROR_REASONS,
@@ -7,7 +6,6 @@ import {
   Roles,
   Privileges,
   RbacObjects,
-  OperatePrivilegeGroupType,
 } from '../../milvus';
 import { timeoutTest } from '../tools';
 import { IP, genCollectionParams, GENERATE_NAME } from '../tools';
@@ -215,14 +213,27 @@ describe(`User Api`, () => {
     expect(grant.error_code).toEqual(ErrorCode.SUCCESS);
   });
 
+  it(`it should grant privilege use grantPrivilegeV2 successfully`, async () => {
+    const res = await authClient.grantPrivilegeV2({
+      role: ROLE_NAME,
+      collection_name: COLLECTION_NAME,
+      db_name: 'default',
+      privilege: Privileges.Query,
+    });
+    expect(res.error_code).toEqual(ErrorCode.SUCCESS);
+  });
+
   it(`It should list grants successfully`, async () => {
     const res = await authClient.listGrants({
       roleName: ROLE_NAME,
     });
-    expect(res.entities.length).toEqual(1);
+    expect(res.entities.length).toEqual(2);
     expect(res.entities[0].object_name).toEqual(COLLECTION_NAME);
     expect(res.entities[0].object.name).toEqual(RbacObjects.Collection);
-    expect(res.entities[0].grantor.privilege.name).toEqual(Privileges.Search);
+    expect(res.entities[0].grantor.privilege.name).toEqual(Privileges.Query);
+    expect(res.entities[1].object_name).toEqual(COLLECTION_NAME);
+    expect(res.entities[1].object.name).toEqual(RbacObjects.Collection);
+    expect(res.entities[1].grantor.privilege.name).toEqual(Privileges.Search);
     expect(res.status.error_code).toEqual(ErrorCode.SUCCESS);
   });
 
@@ -254,6 +265,16 @@ describe(`User Api`, () => {
     expect(res.error_code).toEqual(ErrorCode.SUCCESS);
   });
 
+  it(`It should revoke privilege use revokePrivilegeV2 successfully`, async () => {
+    const res = await authClient.revokePrivilegeV2({
+      role: ROLE_NAME,
+      collection_name: COLLECTION_NAME,
+      db_name: 'default',
+      privilege: Privileges.Query,
+    });
+    expect(res.error_code).toEqual(ErrorCode.SUCCESS);
+  });
+
   it(`It should remove user from role successfully`, async () => {
     const res = await authClient.removeUserFromRole({
       username: USERNAME,
@@ -293,17 +314,10 @@ describe(`User Api`, () => {
   it(`add privileges to a privilege group`, async () => {
     const res = await authClient.addPrivilegesToGroup({
       group_name: PRIVILEGE_GRP_NAME,
-      privileges: [Privileges.Query],
+      privileges: [Privileges.Query, Privileges.Search],
     });
 
     expect(res.error_code).toEqual(ErrorCode.SUCCESS);
-
-    const res2 = await authClient.operatePrivilegeGroup({
-      group_name: PRIVILEGE_GRP_NAME,
-      privileges: [{ name: Privileges.Search }],
-      type: OperatePrivilegeGroupType.AddPrivilegesToGroup,
-    });
-    expect(res2.error_code).toEqual(ErrorCode.SUCCESS);
   });
 
   it(`remove privileges from a privilege group`, async () => {
@@ -320,6 +334,7 @@ describe(`User Api`, () => {
     const grp = res.privilege_groups.find(
       g => g.group_name === PRIVILEGE_GRP_NAME
     )!;
+
     expect(grp.group_name).toEqual(PRIVILEGE_GRP_NAME);
     expect(grp.privileges.map(p => p.name)).toContain(Privileges.Search);
   });