[b/343058608] Extract master keys from HMS (GoogleCloudPlatform#401)

Extract master keys and save them in the `master-keys.csv` file.

Author: dawidxc

dumper/app/src/main/java/com/google/edwmigration/dumper/application/dumper/connector/hive/HiveMetadataConnector.java

@@ -45,6 +45,7 @@
 import java.util.List;
 import java.util.function.Predicate;
 import javax.annotation.Nonnull;
+import org.apache.commons.csv.CSVFormat;
 import org.apache.commons.csv.CSVPrinter;
 import org.apache.commons.lang3.BooleanUtils;
 import org.slf4j.Logger;
@@ -370,6 +371,42 @@ protected String toCallDescription() {
     }
   }
 
+  private static class MasterKeysTask extends AbstractHiveMetadataTask {
+    private static final CSVFormat MASTER_KEYS_FORMAT =
+        FORMAT.builder().setHeader("MasterKey").build();
+
+    private MasterKeysTask() {
+      super("master-keys.csv", unused -> true);
+    }
+
+    @Override
+    protected void run(@Nonnull Writer writer, @Nonnull ThriftClientHandle thriftClientHandle)
+        throws Exception {
+      try (HiveMetastoreThriftClient client =
+              thriftClientHandle.newClient("master-keys-task-client");
+          RecordProgressMonitor monitor =
+              new RecordProgressMonitor("Writing Master Keys to " + getTargetPath())) {
+        ImmutableList<String> masterKeys = client.getMasterKeys();
+        try (CSVPrinter printer = MASTER_KEYS_FORMAT.print(writer)) {
+          for (String masterKey : masterKeys) {
+            monitor.count();
+            printer.printRecord(masterKey);
+          }
+        }
+      }
+    }
+
+    @Override
+    public TaskCategory getCategory() {
+      return TaskCategory.OPTIONAL;
+    }
+
+    @Override
+    protected String toCallDescription() {
+      return "get_master_keys()";
+    }
+  }
+
   public HiveMetadataConnector() {
     super("hiveql");
   }
@@ -389,6 +426,7 @@ public void addTasksTo(List<? super Task<?>> out, @Nonnull ConnectorArguments ar
         arguments.getDefinitionOrDefault(HiveConnectorProperty.MIGRATION_METADATA))) {
       out.add(new CatalogsTask());
       out.add(new DatabasesJsonlTask());
+      out.add(new MasterKeysTask());
     }
 
     if (arguments.isAssessment()) {

dumper/app/src/test/java/com/google/edwmigration/dumper/application/dumper/connector/hive/HiveMetadataConnectorTest.java

@@ -102,7 +102,7 @@ public void addTasksTo_taskExists_success(
 
   @DataPoints("migrationMetadataTaskNames")
   public static final ImmutableList<String> EXPECTED_TASK_NAMES_WITH_MIGRATION_METADATA_ENABLED =
-      ImmutableList.of("catalogs.jsonl", "databases.jsonl");
+      ImmutableList.of("catalogs.jsonl", "databases.jsonl", "master-keys.csv");
 
   @Theory
   public void addTasksTo_migrationMetadataTaskExists_success(

dumper/lib-ext-hive-metastore/src/main/java/com/google/edwmigration/dumper/ext/hive/metastore/HiveMetastoreThriftClient.java

@@ -240,6 +240,9 @@ public String getName() {
   @Nonnull
   public abstract ImmutableList<String> getDatabasesAsJsonl() throws Exception;
 
+  @Nonnull
+  public abstract ImmutableList<String> getMasterKeys() throws Exception;
+
   @Nonnull
   public abstract ImmutableList<String> getAllTableNamesInDatabase(@Nonnull String databaseName)
       throws Exception;

dumper/lib-ext-hive-metastore/src/main/java/com/google/edwmigration/dumper/ext/hive/metastore/HiveMetastoreThriftClient_Superset.java

@@ -130,6 +130,11 @@ public ImmutableList<String> getDatabasesAsJsonl() throws Exception {
         .collect(toImmutableList());
   }
 
+  @Override
+  public ImmutableList<String> getMasterKeys() throws Exception {
+    return ImmutableList.copyOf(client.get_master_keys());
+  }
+
   @Nonnull
   @Override
   public ImmutableList<String> getAllTableNamesInDatabase(@Nonnull String databaseName)

dumper/lib-ext-hive-metastore/src/main/java/com/google/edwmigration/dumper/ext/hive/metastore/HiveMetastoreThriftClient_v2_3_6.java

@@ -127,6 +127,11 @@ public ImmutableList<String> getDatabasesAsJsonl() throws Exception {
         .collect(toImmutableList());
   }
 
+  @Override
+  public ImmutableList<String> getMasterKeys() throws Exception {
+    return ImmutableList.copyOf(client.get_master_keys());
+  }
+
   @Nonnull
   @Override
   public ImmutableList<String> getAllTableNamesInDatabase(@Nonnull String databaseName)