[b/343069385] Extract foreign keys from HMS (#415)

Extracted foreign keys are added to the existing file `tables-raw.jsonl` as a new field `foreignKeys`.
google · Jun 6, 2024 · 27d24a0 · 27d24a0
1 parent c0e8809
commit 27d24a0
Show file tree

Hide file tree

Showing 4 changed files with 33 additions and 0 deletions.
diff --git a/...m/google/edwmigration/dumper/application/dumper/connector/hive/HiveMetadataConnector.java b/...m/google/edwmigration/dumper/application/dumper/connector/hive/HiveMetadataConnector.java
@@ -336,12 +336,15 @@ private void dumpTable(
               Table table = thriftClient.getTable(databaseName, tableName);
               TBase<?, ?> rawTableThriftObject = table.getRawThriftObject();
               ImmutableList<? extends TBase<?, ?>> primaryKeys = table.getRawPrimaryKeys();
+              ImmutableList<? extends TBase<?, ?>> foreignKeys = table.getRawForeignKeys();
               ThriftJsonSerializer jsonSerializer = new ThriftJsonSerializer();
               synchronized (writer) {
                 writer.write("{\"table\":");
                 writer.write(jsonSerializer.serialize(rawTableThriftObject));
                 writer.write(",\"primaryKeys\":");
                 jsonSerializer.serialize(primaryKeys, writer);
+                writer.write(",\"foreignKeys\":");
+                jsonSerializer.serialize(foreignKeys, writer);
                 writer.write("}");
                 writer.write('\n');
               }

diff --git a/...com/google/edwmigration/dumper/ext/hive/metastore/HiveMetastoreThriftClient_Superset.java b/...com/google/edwmigration/dumper/ext/hive/metastore/HiveMetastoreThriftClient_Superset.java
@@ -27,6 +27,7 @@
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableMap;
 import com.google.edwmigration.dumper.ext.hive.metastore.thrift.api.superset.FieldSchema;
+import com.google.edwmigration.dumper.ext.hive.metastore.thrift.api.superset.ForeignKeysRequest;
 import com.google.edwmigration.dumper.ext.hive.metastore.thrift.api.superset.GetCatalogRequest;
 import com.google.edwmigration.dumper.ext.hive.metastore.thrift.api.superset.GetCatalogsResponse;
 import com.google.edwmigration.dumper.ext.hive.metastore.thrift.api.superset.PrimaryKeysRequest;
@@ -482,6 +483,19 @@ public Boolean isCompressed() {
                 .get_primary_keys(new PrimaryKeysRequest(databaseName, tableName))
                 .getPrimaryKeys());
       }
+
+      @Override
+      public ImmutableList<? extends TBase<?, ?>> getRawForeignKeys() throws Exception {
+        return ImmutableList.copyOf(
+            client
+                .get_foreign_keys(
+                    new ForeignKeysRequest(
+                        /*parent_db_name=*/ null,
+                        /*parent_tbl_name=*/ null,
+                        databaseName,
+                        tableName))
+                .getForeignKeys());
+      }
     };
   }
 

diff --git a/...a/com/google/edwmigration/dumper/ext/hive/metastore/HiveMetastoreThriftClient_v2_3_6.java b/...a/com/google/edwmigration/dumper/ext/hive/metastore/HiveMetastoreThriftClient_v2_3_6.java
@@ -27,6 +27,7 @@
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableMap;
 import com.google.edwmigration.dumper.ext.hive.metastore.thrift.api.v2_3_6.FieldSchema;
+import com.google.edwmigration.dumper.ext.hive.metastore.thrift.api.v2_3_6.ForeignKeysRequest;
 import com.google.edwmigration.dumper.ext.hive.metastore.thrift.api.v2_3_6.PrimaryKeysRequest;
 import java.io.IOException;
 import java.util.ArrayList;
@@ -478,6 +479,19 @@ public Boolean isCompressed() {
                 .get_primary_keys(new PrimaryKeysRequest(databaseName, tableName))
                 .getPrimaryKeys());
       }
+
+      @Override
+      public ImmutableList<? extends TBase<?, ?>> getRawForeignKeys() throws Exception {
+        return ImmutableList.copyOf(
+            client
+                .get_foreign_keys(
+                    new ForeignKeysRequest(
+                        /*parent_db_name=*/ null,
+                        /*parent_tbl_name=*/ null,
+                        databaseName,
+                        tableName))
+                .getForeignKeys());
+      }
     };
   }
 

diff --git a/...hive-metastore/src/main/java/com/google/edwmigration/dumper/ext/hive/metastore/Table.java b/...hive-metastore/src/main/java/com/google/edwmigration/dumper/ext/hive/metastore/Table.java
@@ -97,4 +97,6 @@ public interface Table {
   TBase<?, ?> getRawThriftObject();
 
   ImmutableList<? extends TBase<?, ?>> getRawPrimaryKeys() throws Exception;
+
+  ImmutableList<? extends TBase<?, ?>> getRawForeignKeys() throws Exception;
 }