From 27d24a027605fbab4a47a7e29568b7f88e5560d2 Mon Sep 17 00:00:00 2001 From: Dawid Chodura Date: Thu, 6 Jun 2024 15:50:44 +0200 Subject: [PATCH] [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`. --- .../connector/hive/HiveMetadataConnector.java | 3 +++ .../HiveMetastoreThriftClient_Superset.java | 14 ++++++++++++++ .../HiveMetastoreThriftClient_v2_3_6.java | 14 ++++++++++++++ .../dumper/ext/hive/metastore/Table.java | 2 ++ 4 files changed, 33 insertions(+) diff --git a/dumper/app/src/main/java/com/google/edwmigration/dumper/application/dumper/connector/hive/HiveMetadataConnector.java b/dumper/app/src/main/java/com/google/edwmigration/dumper/application/dumper/connector/hive/HiveMetadataConnector.java index a930ac5b9..ce2c3482b 100644 --- a/dumper/app/src/main/java/com/google/edwmigration/dumper/application/dumper/connector/hive/HiveMetadataConnector.java +++ b/dumper/app/src/main/java/com/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> primaryKeys = table.getRawPrimaryKeys(); + ImmutableList> 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/dumper/lib-ext-hive-metastore/src/main/java/com/google/edwmigration/dumper/ext/hive/metastore/HiveMetastoreThriftClient_Superset.java b/dumper/lib-ext-hive-metastore/src/main/java/com/google/edwmigration/dumper/ext/hive/metastore/HiveMetastoreThriftClient_Superset.java index 1f4c81555..738f230b3 100644 --- a/dumper/lib-ext-hive-metastore/src/main/java/com/google/edwmigration/dumper/ext/hive/metastore/HiveMetastoreThriftClient_Superset.java +++ b/dumper/lib-ext-hive-metastore/src/main/java/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> 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/dumper/lib-ext-hive-metastore/src/main/java/com/google/edwmigration/dumper/ext/hive/metastore/HiveMetastoreThriftClient_v2_3_6.java b/dumper/lib-ext-hive-metastore/src/main/java/com/google/edwmigration/dumper/ext/hive/metastore/HiveMetastoreThriftClient_v2_3_6.java index cc617b879..e8b1774a2 100644 --- a/dumper/lib-ext-hive-metastore/src/main/java/com/google/edwmigration/dumper/ext/hive/metastore/HiveMetastoreThriftClient_v2_3_6.java +++ b/dumper/lib-ext-hive-metastore/src/main/java/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> 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/dumper/lib-ext-hive-metastore/src/main/java/com/google/edwmigration/dumper/ext/hive/metastore/Table.java b/dumper/lib-ext-hive-metastore/src/main/java/com/google/edwmigration/dumper/ext/hive/metastore/Table.java index 95db60570..6485501cc 100644 --- a/dumper/lib-ext-hive-metastore/src/main/java/com/google/edwmigration/dumper/ext/hive/metastore/Table.java +++ b/dumper/lib-ext-hive-metastore/src/main/java/com/google/edwmigration/dumper/ext/hive/metastore/Table.java @@ -97,4 +97,6 @@ public interface Table { TBase getRawThriftObject(); ImmutableList> getRawPrimaryKeys() throws Exception; + + ImmutableList> getRawForeignKeys() throws Exception; }