[HUDI-4609] Improve usability of upgrade/downgrade commands in Hudi CLI

apache · Aug 12, 2022 · ee73326 · ee73326
1 parent a41bc41
commit ee73326
Show file tree

Hide file tree

Showing 2 changed files with 38 additions and 6 deletions.
diff --git a/hudi-cli/src/main/java/org/apache/hudi/cli/commands/UpgradeOrDowngradeCommand.java b/hudi-cli/src/main/java/org/apache/hudi/cli/commands/UpgradeOrDowngradeCommand.java
@@ -23,6 +23,8 @@
 import org.apache.hudi.cli.utils.InputStreamConsumer;
 import org.apache.hudi.cli.utils.SparkUtil;
 import org.apache.hudi.common.table.HoodieTableMetaClient;
+import org.apache.hudi.common.table.HoodieTableVersion;
+import org.apache.hudi.common.util.StringUtils;
 
 import org.apache.spark.launcher.SparkLauncher;
 import org.springframework.shell.core.CommandMarker;
@@ -48,15 +50,16 @@ public String upgradeHoodieTable(
     HoodieTableMetaClient metaClient = HoodieCLI.getTableMetaClient();
 
     SparkLauncher sparkLauncher = SparkUtil.initLauncher(sparkPropertiesPath);
-    sparkLauncher.addAppArgs(SparkCommand.UPGRADE.toString(), master, sparkMemory, metaClient.getBasePath(), toVersion);
+    String toVersionName = getHoodieTableVersionName(toVersion, true);
+    sparkLauncher.addAppArgs(SparkCommand.UPGRADE.toString(), master, sparkMemory, metaClient.getBasePath(), toVersionName);
     Process process = sparkLauncher.launch();
     InputStreamConsumer.captureOutput(process);
     int exitCode = process.waitFor();
     HoodieCLI.refreshTableMetadata();
     if (exitCode != 0) {
-      return String.format("Failed: Could not Upgrade/Downgrade Hoodie table to \"%s\".", toVersion);
+      return String.format("Failed: Could not Upgrade/Downgrade Hoodie table to \"%s\".", toVersionName);
     }
-    return String.format("Hoodie table upgraded/downgraded to %s", toVersion);
+    return String.format("Hoodie table upgraded/downgraded to %s", toVersionName);
   }
 
   @CliCommand(value = "downgrade table", help = "Downgrades a table")
@@ -70,14 +73,29 @@ public String downgradeHoodieTable(
 
     HoodieTableMetaClient metaClient = HoodieCLI.getTableMetaClient();
     SparkLauncher sparkLauncher = SparkUtil.initLauncher(sparkPropertiesPath);
-    sparkLauncher.addAppArgs(SparkCommand.DOWNGRADE.toString(), master, sparkMemory, metaClient.getBasePath(), toVersion);
+    String toVersionName = getHoodieTableVersionName(toVersion, false);
+    sparkLauncher.addAppArgs(SparkCommand.DOWNGRADE.toString(), master, sparkMemory, metaClient.getBasePath(), toVersionName);
     Process process = sparkLauncher.launch();
     InputStreamConsumer.captureOutput(process);
     int exitCode = process.waitFor();
     HoodieCLI.refreshTableMetadata();
     if (exitCode != 0) {
-      return String.format("Failed: Could not Upgrade/Downgrade Hoodie table to \"%s\".", toVersion);
+      return String.format("Failed: Could not Upgrade/Downgrade Hoodie table to \"%s\".", toVersionName);
+    }
+    return String.format("Hoodie table upgraded/downgraded to %s", toVersionName);
+  }
+
+  static String getHoodieTableVersionName(String versionOption, boolean overrideWithDefault) {
+    if (StringUtils.isNullOrEmpty(versionOption) && overrideWithDefault) {
+      return HoodieTableVersion.current().name();
+    }
+
+    try {
+      int versionCode = Integer.parseInt(versionOption);
+      return HoodieTableVersion.versionFromCode(versionCode).name();
+    } catch (NumberFormatException e) {
+      // The version option from the CLI is not a number, returns the original String
+      return versionOption;
     }
-    return String.format("Hoodie table upgraded/downgraded to %s", toVersion);
   }
 }
diff --git a/hudi-cli/src/test/java/org/apache/hudi/cli/commands/TestUpgradeDowngradeCommand.java b/hudi-cli/src/test/java/org/apache/hudi/cli/commands/TestUpgradeDowngradeCommand.java
@@ -42,6 +42,7 @@
 import org.junit.jupiter.params.ParameterizedTest;
 import org.junit.jupiter.params.provider.Arguments;
 import org.junit.jupiter.params.provider.MethodSource;
+import org.junit.jupiter.params.provider.ValueSource;
 
 import java.io.IOException;
 import java.util.Arrays;
@@ -140,6 +141,19 @@ public void testUpgradeDowngradeCommand(HoodieTableVersion fromVersion, HoodieTa
     }
   }
 
+  @ParameterizedTest
+  @ValueSource(booleans = {true, false})
+  public void testGetHoodieTableVersionName(boolean overrideWithDefault) {
+    assertEquals(overrideWithDefault ? HoodieTableVersion.current().name() : null,
+        UpgradeOrDowngradeCommand.getHoodieTableVersionName(null, overrideWithDefault));
+    assertEquals(overrideWithDefault ? HoodieTableVersion.current().name() : "",
+        UpgradeOrDowngradeCommand.getHoodieTableVersionName("", overrideWithDefault));
+    assertEquals("FIVE",
+        UpgradeOrDowngradeCommand.getHoodieTableVersionName("FIVE", overrideWithDefault));
+    assertEquals("FIVE",
+        UpgradeOrDowngradeCommand.getHoodieTableVersionName("5", overrideWithDefault));
+  }
+
   private void verifyTableVersion(HoodieTableVersion expectedVersion) throws IOException {
     metaClient = HoodieTableMetaClient.reload(HoodieCLI.getTableMetaClient());
     assertEquals(expectedVersion.versionCode(), metaClient.getTableConfig().getTableVersion().versionCode());