[b/343069385] Extract constraints from HMS

Extracted unique, non-null, default and check constraints are added to the existing file `tables-raw.jsonl` as new fields: `uniqueConstraints`, `nonNullConstraints`, `defaultConstraints` and `checkConstraints`.
google · Jun 6, 2024 · 2f38d2a · 2f38d2a
1 parent 27d24a0
commit 2f38d2a
Show file tree

Hide file tree

Showing 4 changed files with 84 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
@@ -337,6 +337,14 @@ private void dumpTable(
               TBase<?, ?> rawTableThriftObject = table.getRawThriftObject();
               ImmutableList<? extends TBase<?, ?>> primaryKeys = table.getRawPrimaryKeys();
               ImmutableList<? extends TBase<?, ?>> foreignKeys = table.getRawForeignKeys();
+              ImmutableList<? extends TBase<?, ?>> uniqueConstraints =
+                  table.getRawUniqueConstraints();
+              ImmutableList<? extends TBase<?, ?>> nonNullConstraints =
+                  table.getRawNonNullConstraints();
+              ImmutableList<? extends TBase<?, ?>> defaultConstraints =
+                  table.getRawDefaultConstraints();
+              ImmutableList<? extends TBase<?, ?>> checkConstraints =
+                  table.getRawCheckConstraints();
               ThriftJsonSerializer jsonSerializer = new ThriftJsonSerializer();
               synchronized (writer) {
                 writer.write("{\"table\":");
@@ -345,6 +353,14 @@ private void dumpTable(
                 jsonSerializer.serialize(primaryKeys, writer);
                 writer.write(",\"foreignKeys\":");
                 jsonSerializer.serialize(foreignKeys, writer);
+                writer.write(",\"uniqueConstraints\":");
+                jsonSerializer.serialize(uniqueConstraints, writer);
+                writer.write(",\"nonNullConstraints\":");
+                jsonSerializer.serialize(nonNullConstraints, writer);
+                writer.write(",\"defaultConstraints\":");
+                jsonSerializer.serialize(defaultConstraints, writer);
+                writer.write(",\"checkConstraints\":");
+                jsonSerializer.serialize(checkConstraints, 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
@@ -26,11 +26,15 @@
 
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableMap;
+import com.google.edwmigration.dumper.ext.hive.metastore.thrift.api.superset.CheckConstraintsRequest;
+import com.google.edwmigration.dumper.ext.hive.metastore.thrift.api.superset.DefaultConstraintsRequest;
 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.NotNullConstraintsRequest;
 import com.google.edwmigration.dumper.ext.hive.metastore.thrift.api.superset.PrimaryKeysRequest;
+import com.google.edwmigration.dumper.ext.hive.metastore.thrift.api.superset.UniqueConstraintsRequest;
 import com.google.edwmigration.dumper.ext.hive.metastore.thrift.api.superset.WMGetAllResourcePlanRequest;
 import java.io.IOException;
 import java.util.ArrayList;
@@ -496,6 +500,42 @@ public Boolean isCompressed() {
                         tableName))
                 .getForeignKeys());
       }
+
+      @Override
+      public ImmutableList<? extends TBase<?, ?>> getRawUniqueConstraints() throws Exception {
+        return ImmutableList.copyOf(
+            client
+                .get_unique_constraints(
+                    new UniqueConstraintsRequest(table.catName, databaseName, tableName))
+                .getUniqueConstraints());
+      }
+
+      @Override
+      public ImmutableList<? extends TBase<?, ?>> getRawNonNullConstraints() throws Exception {
+        return ImmutableList.copyOf(
+            client
+                .get_not_null_constraints(
+                    new NotNullConstraintsRequest(table.catName, databaseName, tableName))
+                .getNotNullConstraints());
+      }
+
+      @Override
+      public ImmutableList<? extends TBase<?, ?>> getRawDefaultConstraints() throws Exception {
+        return ImmutableList.copyOf(
+            client
+                .get_default_constraints(
+                    new DefaultConstraintsRequest(table.catName, databaseName, tableName))
+                .getDefaultConstraints());
+      }
+
+      @Override
+      public ImmutableList<? extends TBase<?, ?>> getRawCheckConstraints() throws Exception {
+        return ImmutableList.copyOf(
+            client
+                .get_check_constraints(
+                    new CheckConstraintsRequest(table.catName, databaseName, tableName))
+                .getCheckConstraints());
+      }
     };
   }
 

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
@@ -492,6 +492,26 @@ public Boolean isCompressed() {
                         tableName))
                 .getForeignKeys());
       }
+
+      @Override
+      public ImmutableList<? extends TBase<?, ?>> getRawUniqueConstraints() {
+        return ImmutableList.of();
+      }
+
+      @Override
+      public ImmutableList<? extends TBase<?, ?>> getRawNonNullConstraints() {
+        return ImmutableList.of();
+      }
+
+      @Override
+      public ImmutableList<? extends TBase<?, ?>> getRawDefaultConstraints() {
+        return ImmutableList.of();
+      }
+
+      @Override
+      public ImmutableList<? extends TBase<?, ?>> getRawCheckConstraints() {
+        return ImmutableList.of();
+      }
     };
   }
 

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
@@ -99,4 +99,12 @@ public interface Table {
   ImmutableList<? extends TBase<?, ?>> getRawPrimaryKeys() throws Exception;
 
   ImmutableList<? extends TBase<?, ?>> getRawForeignKeys() throws Exception;
+
+  ImmutableList<? extends TBase<?, ?>> getRawUniqueConstraints() throws Exception;
+
+  ImmutableList<? extends TBase<?, ?>> getRawNonNullConstraints() throws Exception;
+
+  ImmutableList<? extends TBase<?, ?>> getRawDefaultConstraints() throws Exception;
+
+  ImmutableList<? extends TBase<?, ?>> getRawCheckConstraints() throws Exception;
 }