[#5570] improve(CLI): Add the ability to revoke all roles or groups i…

…n the Gravitino CLI. Add ability to revoke all roles from a user or group in the Gravitino CLI.
apache · Jan 15, 2025 · f6e9fc2 · f6e9fc2
1 parent 3a48aba
commit f6e9fc2
Show file tree

Hide file tree

Showing 7 changed files with 181 additions and 10 deletions.
diff --git a/clients/cli/src/main/java/org/apache/gravitino/cli/GroupCommandHandler.java b/clients/cli/src/main/java/org/apache/gravitino/cli/GroupCommandHandler.java
@@ -130,14 +130,23 @@ private void handleDeleteCommand() {
 
   /** Handles the "REVOKE" command. */
   private void handleRevokeCommand() {
-    String[] revokeRoles = line.getOptionValues(GravitinoOptions.ROLE);
-    for (String role : revokeRoles) {
+    boolean revokeAll = line.hasOption(GravitinoOptions.ALL);
+    if (revokeAll) {
       gravitinoCommandLine
-          .newRemoveRoleFromGroup(url, ignore, metalake, group, role)
+          .newRemoveAllRoles(url, ignore, metalake, group, CommandEntities.GROUP)
           .validate()
           .handle();
+      System.out.printf("Remove all roles from group %s%n", group);
+    } else {
+      String[] revokeRoles = line.getOptionValues(GravitinoOptions.ROLE);
+      for (String role : revokeRoles) {
+        gravitinoCommandLine
+            .newRemoveRoleFromGroup(url, ignore, metalake, group, role)
+            .validate()
+            .handle();
+      }
+      System.out.printf("Remove roles %s from group %s%n", COMMA_JOINER.join(revokeRoles), group);
     }
-    System.out.printf("Remove roles %s from group %s%n", COMMA_JOINER.join(revokeRoles), group);
   }
 
   /** Handles the "GRANT" command. */

diff --git a/clients/cli/src/main/java/org/apache/gravitino/cli/TestableCommandLine.java b/clients/cli/src/main/java/org/apache/gravitino/cli/TestableCommandLine.java
@@ -86,6 +86,7 @@
 import org.apache.gravitino.cli.commands.ModelDetails;
 import org.apache.gravitino.cli.commands.OwnerDetails;
 import org.apache.gravitino.cli.commands.RegisterModel;
+import org.apache.gravitino.cli.commands.RemoveAllRoles;
 import org.apache.gravitino.cli.commands.RemoveAllTags;
 import org.apache.gravitino.cli.commands.RemoveCatalogProperty;
 import org.apache.gravitino.cli.commands.RemoveFilesetProperty;
@@ -457,6 +458,11 @@ protected RemoveRoleFromGroup newRemoveRoleFromGroup(
     return new RemoveRoleFromGroup(url, ignore, metalake, group, role);
   }
 
+  protected RemoveAllRoles newRemoveAllRoles(
+      String url, boolean ignore, String metalake, String entity, String entityType) {
+    return new RemoveAllRoles(url, ignore, metalake, entity, entityType);
+  }
+
   protected AddRoleToGroup newAddRoleToGroup(
       String url, boolean ignore, String metalake, String group, String role) {
     return new AddRoleToGroup(url, ignore, metalake, group, role);

diff --git a/clients/cli/src/main/java/org/apache/gravitino/cli/UserCommandHandler.java b/clients/cli/src/main/java/org/apache/gravitino/cli/UserCommandHandler.java
@@ -150,14 +150,23 @@ private void handleDeleteCommand() {
 
   /** Handles the "REVOKE" command. */
   private void handleRevokeCommand() {
-    String[] revokeRoles = line.getOptionValues(GravitinoOptions.ROLE);
-    for (String role : revokeRoles) {
-      this.gravitinoCommandLine
-          .newRemoveRoleFromUser(this.url, this.ignore, this.metalake, user, role)
+    boolean removeAll = line.hasOption(GravitinoOptions.ALL);
+    if (removeAll) {
+      gravitinoCommandLine
+          .newRemoveAllRoles(url, ignore, metalake, user, CommandEntities.USER)
           .validate()
           .handle();
+      System.out.printf("Remove all roles from user %s%n", user);
+    } else {
+      String[] revokeRoles = line.getOptionValues(GravitinoOptions.ROLE);
+      for (String role : revokeRoles) {
+        this.gravitinoCommandLine
+            .newRemoveRoleFromUser(this.url, this.ignore, this.metalake, user, role)
+            .validate()
+            .handle();
+      }
+      System.out.printf("Remove roles %s from user %s%n", COMMA_JOINER.join(revokeRoles), user);
     }
-    System.out.printf("Remove roles %s from user %s%n", COMMA_JOINER.join(revokeRoles), user);
   }
 
   /** Handles the "GRANT" command. */

diff --git a/clients/cli/src/main/java/org/apache/gravitino/cli/commands/RemoveAllRoles.java b/clients/cli/src/main/java/org/apache/gravitino/cli/commands/RemoveAllRoles.java
@@ -0,0 +1,98 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *  http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.gravitino.cli.commands;
+
+import java.util.List;
+import org.apache.gravitino.authorization.Group;
+import org.apache.gravitino.authorization.User;
+import org.apache.gravitino.cli.CommandEntities;
+import org.apache.gravitino.cli.ErrorMessages;
+import org.apache.gravitino.client.GravitinoClient;
+import org.apache.gravitino.exceptions.NoSuchGroupException;
+import org.apache.gravitino.exceptions.NoSuchMetalakeException;
+import org.apache.gravitino.exceptions.NoSuchUserException;
+
+/** Removes all roles from a group or user. */
+public class RemoveAllRoles extends Command {
+  protected final String metalake;
+  protected final String entity;
+  protected final String entityType;
+
+  /**
+   * Removes all roles from a group or user.
+   *
+   * @param url The URL of the Gravitino server.
+   * @param ignoreVersions If true don't check the client/server versions match.
+   * @param metalake The name of the metalake.
+   * @param entity the name of the group or user.
+   * @param entityType The type of the entity (group or user).
+   */
+  public RemoveAllRoles(
+      String url, boolean ignoreVersions, String metalake, String entity, String entityType) {
+    super(url, ignoreVersions);
+    this.metalake = metalake;
+    this.entity = entity;
+    this.entityType = entityType;
+  }
+
+  /** Removes all roles from a group or user. */
+  @Override
+  public void handle() {
+    if (CommandEntities.GROUP.equals(entityType)) {
+      revokeAllRolesFromGroup();
+    } else {
+      revokeAllRolesFromUser();
+    }
+  }
+
+  /** Removes all roles from a group. */
+  private void revokeAllRolesFromGroup() {
+    List<String> roles;
+    try {
+      GravitinoClient client = buildClient(metalake);
+      Group group = client.getGroup(entity);
+      roles = group.roles();
+      client.revokeRolesFromGroup(roles, entity);
+    } catch (NoSuchMetalakeException e) {
+      exitWithError(ErrorMessages.UNKNOWN_METALAKE);
+    } catch (NoSuchGroupException e) {
+      exitWithError(ErrorMessages.UNKNOWN_GROUP);
+    } catch (Exception e) {
+      exitWithError(e.getMessage());
+    }
+  }
+
+  /** Removes all roles from a user. */
+  private void revokeAllRolesFromUser() {
+    List<String> roles;
+    try {
+      GravitinoClient client = buildClient(metalake);
+      User user = client.getUser(entity);
+      roles = user.roles();
+      client.revokeRolesFromUser(roles, entity);
+    } catch (NoSuchMetalakeException e) {
+      exitWithError(ErrorMessages.UNKNOWN_METALAKE);
+    } catch (NoSuchUserException e) {
+      exitWithError(ErrorMessages.UNKNOWN_USER);
+    } catch (Exception e) {
+      exitWithError(e.getMessage());
+    }
+  }
+}
diff --git a/clients/cli/src/main/java/org/apache/gravitino/cli/commands/RemoveRoleFromGroup.java b/clients/cli/src/main/java/org/apache/gravitino/cli/commands/RemoveRoleFromGroup.java
@@ -50,7 +50,7 @@ public RemoveRoleFromGroup(
     this.role = role;
   }
 
-  /** Adds a role to a group. */
+  /** Removes a from a group. */
   @Override
   public void handle() {
     try {

diff --git a/clients/cli/src/test/java/org/apache/gravitino/cli/TestGroupCommands.java b/clients/cli/src/test/java/org/apache/gravitino/cli/TestGroupCommands.java
@@ -39,6 +39,7 @@
 import org.apache.gravitino.cli.commands.GroupAudit;
 import org.apache.gravitino.cli.commands.GroupDetails;
 import org.apache.gravitino.cli.commands.ListGroups;
+import org.apache.gravitino.cli.commands.RemoveAllRoles;
 import org.apache.gravitino.cli.commands.RemoveRoleFromGroup;
 import org.junit.jupiter.api.AfterEach;
 import org.junit.jupiter.api.BeforeEach;
@@ -260,6 +261,32 @@ void testRemoveRolesFromGroupCommand() {
     verify(mockRemoveSecondRole).handle();
   }
 
+  @Test
+  void testRemoveAllRolesFromGroupCommand() {
+    RemoveAllRoles mock = mock(RemoveAllRoles.class);
+    when(mockCommandLine.hasOption(GravitinoOptions.METALAKE)).thenReturn(true);
+    when(mockCommandLine.getOptionValue(GravitinoOptions.METALAKE)).thenReturn("metalake_demo");
+    when(mockCommandLine.hasOption(GravitinoOptions.GROUP)).thenReturn(true);
+    when(mockCommandLine.getOptionValue(GravitinoOptions.GROUP)).thenReturn("groupA");
+    when(mockCommandLine.hasOption(GravitinoOptions.ALL)).thenReturn(true);
+    GravitinoCommandLine commandLine =
+        spy(
+            new GravitinoCommandLine(
+                mockCommandLine, mockOptions, CommandEntities.GROUP, CommandActions.REVOKE));
+
+    doReturn(mock)
+        .when(commandLine)
+        .newRemoveAllRoles(
+            GravitinoCommandLine.DEFAULT_URL,
+            false,
+            "metalake_demo",
+            "groupA",
+            CommandEntities.GROUP);
+    doReturn(mock).when(mock).validate();
+    commandLine.handleCommandLine();
+    verify(mock).handle();
+  }
+
   @Test
   void testAddRolesToGroupCommand() {
     AddRoleToGroup mockAddFirstRole = mock(AddRoleToGroup.class);

diff --git a/clients/cli/src/test/java/org/apache/gravitino/cli/TestUserCommands.java b/clients/cli/src/test/java/org/apache/gravitino/cli/TestUserCommands.java
@@ -37,6 +37,7 @@
 import org.apache.gravitino.cli.commands.CreateUser;
 import org.apache.gravitino.cli.commands.DeleteUser;
 import org.apache.gravitino.cli.commands.ListUsers;
+import org.apache.gravitino.cli.commands.RemoveAllRoles;
 import org.apache.gravitino.cli.commands.RemoveRoleFromUser;
 import org.apache.gravitino.cli.commands.UserAudit;
 import org.apache.gravitino.cli.commands.UserDetails;
@@ -261,6 +262,27 @@ void testRemoveRolesFromUserCommand() {
     verify(mockRemoveFirstRole).handle();
   }
 
+  @Test
+  void removeAllRolesFromUserCommand() {
+    RemoveAllRoles mockRemove = mock(RemoveAllRoles.class);
+    when(mockCommandLine.hasOption(GravitinoOptions.METALAKE)).thenReturn(true);
+    when(mockCommandLine.getOptionValue(GravitinoOptions.METALAKE)).thenReturn("metalake_demo");
+    when(mockCommandLine.hasOption(GravitinoOptions.USER)).thenReturn(true);
+    when(mockCommandLine.getOptionValue(GravitinoOptions.USER)).thenReturn("user");
+    when(mockCommandLine.hasOption(GravitinoOptions.ALL)).thenReturn(true);
+    GravitinoCommandLine commandLine =
+        spy(
+            new GravitinoCommandLine(
+                mockCommandLine, mockOptions, CommandEntities.USER, CommandActions.REVOKE));
+    doReturn(mockRemove)
+        .when(commandLine)
+        .newRemoveAllRoles(
+            GravitinoCommandLine.DEFAULT_URL, false, "metalake_demo", "user", CommandEntities.USER);
+    doReturn(mockRemove).when(mockRemove).validate();
+    commandLine.handleCommandLine();
+    verify(mockRemove).handle();
+  }
+
   @Test
   void testAddRolesToUserCommand() {
     AddRoleToUser mockAddFirstRole = mock(AddRoleToUser.class);