[apache#6449] improve(CLI): Refactor getAuth method in CLI (apache#6451)

### What changes were proposed in this pull request? Refactor getAuth method in CLI. ### Why are the changes needed? Fix: apache#6449 ### Does this PR introduce _any_ user-facing change? No ### How was this patch tested? local test.
diqiu50 · Feb 19, 2025 · 703b3ed · 703b3ed
1 parent db7f06f
commit 703b3ed
Show file tree

Hide file tree

Showing 15 changed files with 56 additions and 54 deletions.
diff --git a/clients/cli/src/main/java/org/apache/gravitino/cli/CatalogCommandHandler.java b/clients/cli/src/main/java/org/apache/gravitino/cli/CatalogCommandHandler.java
@@ -64,7 +64,7 @@ public CatalogCommandHandler(
   @Override
   protected void handle() {
     String userName = line.getOptionValue(GravitinoOptions.LOGIN);
-    Command.setAuthenticationMode(getAuth(line), userName);
+    Command.setAuthenticationMode(context.auth(), userName);
     List<String> missingEntities = Lists.newArrayList();
 
     if (CommandActions.LIST.equals(command)) {

diff --git a/clients/cli/src/main/java/org/apache/gravitino/cli/ColumnCommandHandler.java b/clients/cli/src/main/java/org/apache/gravitino/cli/ColumnCommandHandler.java
@@ -66,7 +66,7 @@ public ColumnCommandHandler(
   @Override
   protected void handle() {
     String userName = line.getOptionValue(GravitinoOptions.LOGIN);
-    Command.setAuthenticationMode(getAuth(line), userName);
+    Command.setAuthenticationMode(context.auth(), userName);
 
     List<String> missingEntities = Lists.newArrayList();
     if (catalog == null) missingEntities.add(CommandEntities.CATALOG);

diff --git a/clients/cli/src/main/java/org/apache/gravitino/cli/CommandContext.java b/clients/cli/src/main/java/org/apache/gravitino/cli/CommandContext.java
@@ -31,11 +31,14 @@ public class CommandContext {
   private final String url;
   private final boolean quiet;
   private final CommandLine line;
+  private final String auth;
 
   private String ignoreEnv;
   private boolean ignoreSet = false;
   private String urlEnv;
   private boolean urlSet = false;
+  private String authEnv;
+  private boolean authSet = false;
   // Can add more "global" command flags here without any major changes e.g. a guiet flag
 
   /**
@@ -55,6 +58,7 @@ public CommandContext(CommandLine line) {
 
     this.url = getUrl();
     this.ignoreVersions = getIgnore();
+    this.auth = this.getAuth();
   }
 
   /**
@@ -102,6 +106,15 @@ public boolean quiet() {
     return quiet;
   }
 
+  /**
+   * Returns the authentication type.
+   *
+   * @return The authentication type.
+   */
+  public String auth() {
+    return auth;
+  }
+
   /**
    * Retrieves the Gravitino URL from the command line options or the GRAVITINO_URL environment
    * variable or the Gravitino config file.
@@ -166,4 +179,34 @@ private boolean getIgnore() {
 
     return ignore;
   }
+
+  private String getAuth() {
+    // If specified on the command line use that
+    if (line.hasOption(GravitinoOptions.SIMPLE)) {
+      return GravitinoOptions.SIMPLE;
+    }
+
+    // Cache the Gravitino authentication type environment variable
+    if (authEnv == null && !authSet) {
+      authEnv = System.getenv("GRAVITINO_AUTH");
+      authSet = true;
+    }
+
+    // If set return the Gravitino authentication type environment variable
+    if (authEnv != null) {
+      return authEnv;
+    }
+
+    // Check if the authentication type is specified in the configuration file
+    GravitinoConfig config = new GravitinoConfig(null);
+    if (config.fileExists()) {
+      config.read();
+      String configAuthType = config.getGravitinoAuthType();
+      if (configAuthType != null) {
+        return configAuthType;
+      }
+    }
+
+    return null;
+  }
 }
diff --git a/clients/cli/src/main/java/org/apache/gravitino/cli/CommandHandler.java b/clients/cli/src/main/java/org/apache/gravitino/cli/CommandHandler.java
@@ -21,49 +21,9 @@
 
 import com.google.common.base.Joiner;
 import java.util.List;
-import org.apache.commons.cli.CommandLine;
 
 public abstract class CommandHandler {
   public static final Joiner COMMA_JOINER = Joiner.on(", ").skipNulls();
-  private String authEnv;
-  private boolean authSet = false;
-
-  /**
-   * Retrieves the Gravitino authentication from the command line options or the GRAVITINO_AUTH
-   * environment variable or the Gravitino config file.
-   *
-   * @param line The command line instance.
-   * @return The Gravitino authentication, or null if not found.
-   */
-  public String getAuth(CommandLine line) {
-    // If specified on the command line use that
-    if (line.hasOption(GravitinoOptions.SIMPLE)) {
-      return GravitinoOptions.SIMPLE;
-    }
-
-    // Cache the Gravitino authentication type environment variable
-    if (authEnv == null && !authSet) {
-      authEnv = System.getenv("GRAVITINO_AUTH");
-      authSet = true;
-    }
-
-    // If set return the Gravitino authentication type environment variable
-    if (authEnv != null) {
-      return authEnv;
-    }
-
-    // Check if the authentication type is specified in the configuration file
-    GravitinoConfig config = new GravitinoConfig(null);
-    if (config.fileExists()) {
-      config.read();
-      String configAuthType = config.getGravitinoAuthType();
-      if (configAuthType != null) {
-        return configAuthType;
-      }
-    }
-
-    return null;
-  }
 
   protected abstract void handle();
 

diff --git a/clients/cli/src/main/java/org/apache/gravitino/cli/FilesetCommandHandler.java b/clients/cli/src/main/java/org/apache/gravitino/cli/FilesetCommandHandler.java
@@ -66,7 +66,7 @@ public FilesetCommandHandler(
   @Override
   protected void handle() {
     String userName = line.getOptionValue(GravitinoOptions.LOGIN);
-    Command.setAuthenticationMode(getAuth(line), userName);
+    Command.setAuthenticationMode(context.auth(), userName);
 
     List<String> missingEntities = Lists.newArrayList();
     if (catalog == null) missingEntities.add(CommandEntities.CATALOG);

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
@@ -58,7 +58,7 @@ public GroupCommandHandler(
   @Override
   protected void handle() {
     String userName = line.getOptionValue(GravitinoOptions.LOGIN);
-    Command.setAuthenticationMode(getAuth(line), userName);
+    Command.setAuthenticationMode(context.auth(), userName);
 
     if (CommandActions.LIST.equals(command)) {
       handleListCommand();

diff --git a/clients/cli/src/main/java/org/apache/gravitino/cli/MetalakeCommandHandler.java b/clients/cli/src/main/java/org/apache/gravitino/cli/MetalakeCommandHandler.java
@@ -56,7 +56,7 @@ public MetalakeCommandHandler(
   public void handle() {
     String userName = line.getOptionValue(GravitinoOptions.LOGIN);
     FullName name = new FullName(line);
-    Command.setAuthenticationMode(getAuth(line), userName);
+    Command.setAuthenticationMode(context.auth(), userName);
 
     if (CommandActions.LIST.equals(command)) {
       handleListCommand();

diff --git a/clients/cli/src/main/java/org/apache/gravitino/cli/ModelCommandHandler.java b/clients/cli/src/main/java/org/apache/gravitino/cli/ModelCommandHandler.java
@@ -65,7 +65,7 @@ public ModelCommandHandler(
   @Override
   protected void handle() {
     String userName = line.getOptionValue(GravitinoOptions.LOGIN);
-    Command.setAuthenticationMode(getAuth(line), userName);
+    Command.setAuthenticationMode(context.auth(), userName);
 
     List<String> missingEntities = Lists.newArrayList();
     if (catalog == null) missingEntities.add(CommandEntities.CATALOG);

diff --git a/clients/cli/src/main/java/org/apache/gravitino/cli/OwnerCommandHandler.java b/clients/cli/src/main/java/org/apache/gravitino/cli/OwnerCommandHandler.java
@@ -66,7 +66,7 @@ public OwnerCommandHandler(
   @Override
   protected void handle() {
     String userName = line.getOptionValue(GravitinoOptions.LOGIN);
-    Command.setAuthenticationMode(getAuth(line), userName);
+    Command.setAuthenticationMode(context.auth(), userName);
 
     if (entityName == null && !CommandEntities.METALAKE.equals(entity)) {
       System.err.println(ErrorMessages.MISSING_NAME);

diff --git a/clients/cli/src/main/java/org/apache/gravitino/cli/RoleCommandHandler.java b/clients/cli/src/main/java/org/apache/gravitino/cli/RoleCommandHandler.java
@@ -46,9 +46,8 @@ public RoleCommandHandler(
 
   /** Handles the command execution logic based on the provided command. */
   public void handle() {
-    String auth = getAuth(line);
     String userName = line.getOptionValue(GravitinoOptions.LOGIN);
-    Command.setAuthenticationMode(auth, userName);
+    Command.setAuthenticationMode(context.auth(), userName);
 
     metalake = new FullName(line).getMetalakeName();
 

diff --git a/clients/cli/src/main/java/org/apache/gravitino/cli/SchemaCommandHandler.java b/clients/cli/src/main/java/org/apache/gravitino/cli/SchemaCommandHandler.java
@@ -61,7 +61,7 @@ public SchemaCommandHandler(
   @Override
   protected void handle() {
     String userName = line.getOptionValue(GravitinoOptions.LOGIN);
-    Command.setAuthenticationMode(getAuth(line), userName);
+    Command.setAuthenticationMode(context.auth(), userName);
 
     List<String> missingEntities = Lists.newArrayList();
     if (metalake == null) missingEntities.add(CommandEntities.METALAKE);

diff --git a/clients/cli/src/main/java/org/apache/gravitino/cli/TableCommandHandler.java b/clients/cli/src/main/java/org/apache/gravitino/cli/TableCommandHandler.java
@@ -64,7 +64,7 @@ public TableCommandHandler(
   @Override
   protected void handle() {
     String userName = line.getOptionValue(GravitinoOptions.LOGIN);
-    Command.setAuthenticationMode(getAuth(line), userName);
+    Command.setAuthenticationMode(context.auth(), userName);
 
     List<String> missingEntities = Lists.newArrayList();
     if (catalog == null) missingEntities.add(CommandEntities.CATALOG);

diff --git a/clients/cli/src/main/java/org/apache/gravitino/cli/TagCommandHandler.java b/clients/cli/src/main/java/org/apache/gravitino/cli/TagCommandHandler.java
@@ -50,7 +50,7 @@ public TagCommandHandler(
   public void handle() {
     String userName = line.getOptionValue(GravitinoOptions.LOGIN);
     FullName name = new FullName(line);
-    Command.setAuthenticationMode(getAuth(line), userName);
+    Command.setAuthenticationMode(context.auth(), userName);
 
     metalake = name.getMetalakeName();
 

diff --git a/clients/cli/src/main/java/org/apache/gravitino/cli/TopicCommandHandler.java b/clients/cli/src/main/java/org/apache/gravitino/cli/TopicCommandHandler.java
@@ -64,7 +64,7 @@ public TopicCommandHandler(
   @Override
   protected void handle() {
     String userName = line.getOptionValue(GravitinoOptions.LOGIN);
-    Command.setAuthenticationMode(getAuth(line), userName);
+    Command.setAuthenticationMode(context.auth(), userName);
 
     List<String> missingEntities = Lists.newArrayList();
     if (catalog == null) missingEntities.add(CommandEntities.CATALOG);

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
@@ -58,7 +58,7 @@ public UserCommandHandler(
   @Override
   protected void handle() {
     String userName = line.getOptionValue(GravitinoOptions.LOGIN);
-    Command.setAuthenticationMode(getAuth(line), userName);
+    Command.setAuthenticationMode(context.auth(), userName);
 
     user = line.getOptionValue(GravitinoOptions.USER);