Allow message_name field for protobuf ingestion (#9480)

metadata-integration/java/datahub-protobuf/build.gradle

@@ -31,10 +31,10 @@ dependencies {
     implementation externalDependency.commonsCli
     implementation externalDependency.httpAsyncClient
     implementation externalDependency.slf4jApi
+    implementation externalDependency.jacksonCore
     compileOnly externalDependency.lombok
     annotationProcessor externalDependency.lombok
-    testImplementation externalDependency.junitJupiterApi
+    testImplementation externalDependency.testng
-    testRuntimeOnly externalDependency.junitJupiterEngine
 }
 
 import java.nio.file.Paths
@@ -61,10 +61,7 @@ jacocoTestReport {
     dependsOn test // tests are required to run before generating the report
 }
 
-test {
+test.finalizedBy jacocoTestReport
-    useJUnit()
-    finalizedBy jacocoTestReport
-}
 
 
 task checkShadowJar(type: Exec) {

metadata-integration/java/datahub-protobuf/src/main/java/datahub/protobuf/Proto2DataHub.java

@@ -67,6 +67,13 @@ public class Proto2DataHub {
               "[Optional if using --directory] The protobuf source file. Typically a .proto file.")
           .build();
 
+  private static final Option OPTION_MESSAGE_NAME =
+      Option.builder()
+          .longOpt("message_name")
+          .hasArg()
+          .desc("[Optional] The protobuf message name to read from.")
+          .build();
+
   private static final Option OPTION_DIR =
       Option.builder()
           .longOpt("directory")
@@ -166,6 +173,7 @@ public class Proto2DataHub {
     private final String dataPlatform;
     private final String protoc;
     private final String inputFile;
+    private final String messageName;
     private final String inputDir;
     private final TransportOptions transport;
     private final String filename;
@@ -191,6 +199,7 @@ public class Proto2DataHub {
       dataPlatform = cli.getOptionValue(OPTION_DATAHUB_PLATFORM, "kafka").toLowerCase(Locale.ROOT);
       protoc = cli.getOptionValue(OPTION_DESCRIPTOR);
       inputFile = cli.getOptionValue(OPTION_FILE, null);
+      messageName = cli.getOptionValue(OPTION_MESSAGE_NAME, null);
       transport =
           TransportOptions.valueOf(
               cli.getOptionValue(OPTION_TRANSPORT, "rest").toUpperCase(Locale.ROOT));
@@ -250,6 +259,7 @@ public class Proto2DataHub {
         .addOption(OPTION_DATAHUB_TOKEN)
         .addOption(OPTION_DESCRIPTOR)
         .addOption(OPTION_FILE)
+        .addOption(OPTION_MESSAGE_NAME)
         .addOption(OPTION_DIR)
         .addOption(OPTION_EXCLUDE_PATTERN)
         .addOption(OPTION_DATAHUB_USER)
@@ -354,6 +364,7 @@ public class Proto2DataHub {
                       .setGithubOrganization(config.githubOrg)
                       .setSlackTeamId(config.slackId)
                       .setSubType(config.subType)
+                      .setMessageName(config.messageName)
                       .build();
 
               dataset

metadata-integration/java/datahub-protobuf/src/test/java/datahub/protobuf/ProtobufDatasetTest.java

@@ -1,8 +1,8 @@
 package datahub.protobuf;
 
 import static datahub.protobuf.TestFixtures.*;
-import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.testng.Assert.assertEquals;
-import static org.junit.jupiter.api.Assertions.assertNotNull;
+import static org.testng.Assert.assertNotNull;
 
 import com.linkedin.common.FabricType;
 import com.linkedin.common.GlobalTags;
@@ -34,7 +34,7 @@ import java.io.IOException;
 import java.util.Set;
 import java.util.stream.Collectors;
 import java.util.stream.Stream;
-import org.junit.jupiter.api.Test;
+import org.testng.annotations.Test;
 
 public class ProtobufDatasetTest {
 

metadata-integration/java/datahub-protobuf/src/test/java/datahub/protobuf/ProtobufUtilsTest.java

@@ -2,13 +2,13 @@ package datahub.protobuf;
 
 import static datahub.protobuf.TestFixtures.getTestProtobufFileSet;
 import static datahub.protobuf.TestFixtures.getTestProtoc;
-import static org.junit.jupiter.api.Assertions.*;
+import static org.testng.Assert.*;
 
 import com.google.protobuf.DescriptorProtos;
 import com.google.protobuf.ExtensionRegistry;
 import datahub.protobuf.model.ProtobufGraph;
 import java.io.IOException;
-import org.junit.jupiter.api.Test;
+import org.testng.annotations.Test;
 
 public class ProtobufUtilsTest {
 

metadata-integration/java/datahub-protobuf/src/test/java/datahub/protobuf/model/ProtobufEnumTest.java

@@ -1,6 +1,6 @@
 package datahub.protobuf.model;
 
-import static org.junit.jupiter.api.Assertions.*;
+import static org.testng.Assert.*;
 
 import com.google.protobuf.DescriptorProtos.DescriptorProto;
 import com.google.protobuf.DescriptorProtos.EnumDescriptorProto;
@@ -11,7 +11,7 @@ import java.util.List;
 import java.util.Set;
 import java.util.stream.Collectors;
 import java.util.stream.Stream;
-import org.junit.jupiter.api.Test;
+import org.testng.annotations.Test;
 
 public class ProtobufEnumTest {
 

metadata-integration/java/datahub-protobuf/src/test/java/datahub/protobuf/model/ProtobufFieldTest.java

@@ -1,7 +1,7 @@
 package datahub.protobuf.model;
 
 import static datahub.protobuf.TestFixtures.*;
-import static org.junit.jupiter.api.Assertions.*;
+import static org.testng.Assert.*;
 
 import com.google.protobuf.DescriptorProtos.DescriptorProto;
 import com.google.protobuf.DescriptorProtos.FieldDescriptorProto;
@@ -22,7 +22,7 @@ import datahub.protobuf.ProtobufDataset;
 import java.io.IOException;
 import java.util.Arrays;
 import java.util.Set;
-import org.junit.jupiter.api.Test;
+import org.testng.annotations.Test;
 
 public class ProtobufFieldTest {
   private static final DescriptorProto EXPECTED_MESSAGE_PROTO =

metadata-integration/java/datahub-protobuf/src/test/java/datahub/protobuf/model/ProtobufGraphTest.java

@@ -2,14 +2,14 @@ package datahub.protobuf.model;
 
 import static datahub.protobuf.TestFixtures.getTestProtobufFileSet;
 import static datahub.protobuf.TestFixtures.getTestProtobufGraph;
-import static org.junit.jupiter.api.Assertions.*;
+import static org.testng.Assert.*;
 
 import com.google.protobuf.DescriptorProtos.FileDescriptorSet;
 import java.io.IOException;
 import java.util.HashSet;
 import java.util.List;
 import java.util.stream.Collectors;
-import org.junit.jupiter.api.Test;
+import org.testng.annotations.Test;
 
 public class ProtobufGraphTest {
 

metadata-integration/java/datahub-protobuf/src/test/java/datahub/protobuf/model/ProtobufMessageTest.java

@@ -1,6 +1,6 @@
 package datahub.protobuf.model;
 
-import static org.junit.jupiter.api.Assertions.*;
+import static org.testng.Assert.*;
 
 import com.google.protobuf.DescriptorProtos.DescriptorProto;
 import com.google.protobuf.DescriptorProtos.FileDescriptorProto;
@@ -11,7 +11,7 @@ import java.util.List;
 import java.util.Set;
 import java.util.stream.Collectors;
 import java.util.stream.Stream;
-import org.junit.jupiter.api.Test;
+import org.testng.annotations.Test;
 
 public class ProtobufMessageTest {
 

metadata-integration/java/datahub-protobuf/src/test/java/datahub/protobuf/model/ProtobufOneOfFieldTest.java

@@ -1,6 +1,6 @@
 package datahub.protobuf.model;
 
-import static org.junit.jupiter.api.Assertions.*;
+import static org.testng.Assert.*;
 
 import com.google.protobuf.DescriptorProtos.DescriptorProto;
 import com.google.protobuf.DescriptorProtos.FieldDescriptorProto;
@@ -12,7 +12,7 @@ import java.util.List;
 import java.util.Set;
 import java.util.stream.Collectors;
 import java.util.stream.Stream;
-import org.junit.jupiter.api.Test;
+import org.testng.annotations.Test;
 
 public class ProtobufOneOfFieldTest {
 

metadata-integration/java/datahub-protobuf/src/test/java/datahub/protobuf/visitors/VisitContextTest.java

@@ -2,7 +2,7 @@ package datahub.protobuf.visitors;
 
 import static datahub.protobuf.TestFixtures.getTestProtobufFileSet;
 import static datahub.protobuf.TestFixtures.getTestProtobufGraph;
-import static org.junit.jupiter.api.Assertions.assertNotEquals;
+import static org.testng.Assert.assertNotEquals;
 
 import com.google.protobuf.DescriptorProtos.FileDescriptorSet;
 import datahub.protobuf.model.FieldTypeEdge;
@@ -13,7 +13,7 @@ import java.util.List;
 import java.util.Set;
 import java.util.stream.Collectors;
 import org.jgrapht.GraphPath;
-import org.junit.jupiter.api.Test;
+import org.testng.annotations.Test;
 
 public class VisitContextTest {
 

metadata-integration/java/datahub-protobuf/src/test/java/datahub/protobuf/visitors/dataset/DatasetVisitorTest.java

@@ -1,7 +1,7 @@
 package datahub.protobuf.visitors.dataset;
 
 import static datahub.protobuf.TestFixtures.*;
-import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.testng.Assert.assertEquals;
 
 import com.linkedin.common.urn.DatasetUrn;
 import com.linkedin.data.template.RecordTemplate;
@@ -14,7 +14,7 @@ import java.net.URISyntaxException;
 import java.util.List;
 import java.util.stream.Collectors;
 import java.util.stream.Stream;
-import org.junit.jupiter.api.Test;
+import org.testng.annotations.Test;
 
 public class DatasetVisitorTest {
 

metadata-integration/java/datahub-protobuf/src/test/java/datahub/protobuf/visitors/dataset/DescriptionVisitorTest.java

@@ -1,14 +1,14 @@
 package datahub.protobuf.visitors.dataset;
 
 import static datahub.protobuf.TestFixtures.*;
-import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.testng.Assert.assertEquals;
 
 import datahub.protobuf.model.ProtobufGraph;
 import java.io.IOException;
 import java.util.List;
 import java.util.Set;
 import java.util.stream.Collectors;
-import org.junit.jupiter.api.Test;
+import org.testng.annotations.Test;
 
 public class DescriptionVisitorTest {
 

metadata-integration/java/datahub-protobuf/src/test/java/datahub/protobuf/visitors/dataset/DomainVisitorTest.java

@@ -2,7 +2,7 @@ package datahub.protobuf.visitors.dataset;
 
 import static datahub.protobuf.TestFixtures.getTestProtobufGraph;
 import static datahub.protobuf.TestFixtures.getVisitContextBuilder;
-import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.testng.Assert.assertEquals;
 
 import com.linkedin.common.urn.Urn;
 import datahub.protobuf.model.ProtobufGraph;
@@ -10,7 +10,7 @@ import java.io.IOException;
 import java.util.List;
 import java.util.Set;
 import java.util.stream.Collectors;
-import org.junit.jupiter.api.Test;
+import org.testng.annotations.Test;
 
 public class DomainVisitorTest {
 

metadata-integration/java/datahub-protobuf/src/test/java/datahub/protobuf/visitors/dataset/InstitutionalMemoryVisitorTest.java

@@ -1,7 +1,7 @@
 package datahub.protobuf.visitors.dataset;
 
 import static datahub.protobuf.TestFixtures.*;
-import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.testng.Assert.assertEquals;
 
 import com.linkedin.common.InstitutionalMemoryMetadata;
 import com.linkedin.common.url.Url;
@@ -9,7 +9,7 @@ import java.io.IOException;
 import java.util.List;
 import java.util.Set;
 import java.util.stream.Collectors;
-import org.junit.jupiter.api.Test;
+import org.testng.annotations.Test;
 
 public class InstitutionalMemoryVisitorTest {
 

metadata-integration/java/datahub-protobuf/src/test/java/datahub/protobuf/visitors/dataset/KafkaTopicPropertyVisitorTest.java

@@ -2,7 +2,7 @@ package datahub.protobuf.visitors.dataset;
 
 import static datahub.protobuf.TestFixtures.getTestProtobufGraph;
 import static datahub.protobuf.TestFixtures.getVisitContextBuilder;
-import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.testng.Assert.assertEquals;
 
 import com.linkedin.data.template.StringMap;
 import com.linkedin.dataset.DatasetProperties;
@@ -11,7 +11,7 @@ import java.util.List;
 import java.util.Map;
 import java.util.Set;
 import java.util.stream.Collectors;
-import org.junit.jupiter.api.Test;
+import org.testng.annotations.Test;
 
 public class KafkaTopicPropertyVisitorTest {
 

metadata-integration/java/datahub-protobuf/src/test/java/datahub/protobuf/visitors/dataset/OwnershipVisitorTest.java

@@ -2,7 +2,7 @@ package datahub.protobuf.visitors.dataset;
 
 import static datahub.protobuf.TestFixtures.getTestProtobufGraph;
 import static datahub.protobuf.TestFixtures.getVisitContextBuilder;
-import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.testng.Assert.assertEquals;
 
 import com.linkedin.common.Owner;
 import com.linkedin.common.OwnershipSource;
@@ -14,7 +14,7 @@ import java.io.IOException;
 import java.util.List;
 import java.util.Set;
 import java.util.stream.Collectors;
-import org.junit.jupiter.api.Test;
+import org.testng.annotations.Test;
 
 public class OwnershipVisitorTest {
 

metadata-integration/java/datahub-protobuf/src/test/java/datahub/protobuf/visitors/dataset/PropertyVisitorTest.java

@@ -3,7 +3,7 @@ package datahub.protobuf.visitors.dataset;
 import static datahub.protobuf.TestFixtures.getTestProtobufGraph;
 import static datahub.protobuf.TestFixtures.getVisitContextBuilder;
 import static java.util.Map.entry;
-import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.testng.Assert.assertEquals;
 
 import com.linkedin.data.template.StringMap;
 import com.linkedin.dataset.DatasetProperties;
@@ -11,7 +11,7 @@ import java.io.IOException;
 import java.util.List;
 import java.util.Map;
 import java.util.stream.Collectors;
-import org.junit.jupiter.api.Test;
+import org.testng.annotations.Test;
 
 public class PropertyVisitorTest {
 

metadata-integration/java/datahub-protobuf/src/test/java/datahub/protobuf/visitors/dataset/TermAssociationVisitorTest.java

@@ -2,7 +2,7 @@ package datahub.protobuf.visitors.dataset;
 
 import static datahub.protobuf.TestFixtures.getTestProtobufGraph;
 import static datahub.protobuf.TestFixtures.getVisitContextBuilder;
-import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.testng.Assert.assertEquals;
 
 import com.linkedin.common.GlossaryTermAssociation;
 import com.linkedin.common.urn.GlossaryTermUrn;
@@ -10,7 +10,7 @@ import java.io.IOException;
 import java.util.List;
 import java.util.Set;
 import java.util.stream.Collectors;
-import org.junit.jupiter.api.Test;
+import org.testng.annotations.Test;
 
 public class TermAssociationVisitorTest {
 

metadata-integration/java/datahub-protobuf/src/test/java/datahub/protobuf/visitors/field/ProtobufExtensionFieldVisitorTest.java

@@ -1,7 +1,7 @@
 package datahub.protobuf.visitors.field;
 
 import static datahub.protobuf.TestFixtures.*;
-import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.testng.Assert.assertEquals;
 
 import com.linkedin.common.GlobalTags;
 import com.linkedin.common.GlossaryTermAssociation;
@@ -23,7 +23,7 @@ import java.net.URISyntaxException;
 import java.util.List;
 import java.util.stream.Collectors;
 import java.util.stream.Stream;
-import org.junit.jupiter.api.Test;
+import org.testng.annotations.Test;
 
 public class ProtobufExtensionFieldVisitorTest {
 

metadata-integration/java/datahub-protobuf/src/test/java/datahub/protobuf/visitors/field/SchemaFieldVisitorTest.java

@@ -2,7 +2,7 @@ package datahub.protobuf.visitors.field;
 
 import static datahub.protobuf.TestFixtures.getTestProtobufGraph;
 import static datahub.protobuf.TestFixtures.getVisitContextBuilder;
-import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.testng.Assert.assertEquals;
 
 import com.linkedin.schema.NumberType;
 import com.linkedin.schema.SchemaField;
@@ -15,7 +15,7 @@ import java.io.IOException;
 import java.util.List;
 import java.util.stream.Collectors;
 import java.util.stream.Stream;
-import org.junit.jupiter.api.Test;
+import org.testng.annotations.Test;
 
 public class SchemaFieldVisitorTest {
 

metadata-integration/java/datahub-protobuf/src/test/java/datahub/protobuf/visitors/tag/TagVisitorTest.java

@@ -2,7 +2,7 @@ package datahub.protobuf.visitors.tag;
 
 import static datahub.protobuf.TestFixtures.getTestProtobufGraph;
 import static datahub.protobuf.TestFixtures.getVisitContextBuilder;
-import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.testng.Assert.assertEquals;
 
 import com.linkedin.tag.TagProperties;
 import datahub.event.MetadataChangeProposalWrapper;
@@ -11,7 +11,7 @@ import java.io.IOException;
 import java.util.List;
 import java.util.Set;
 import java.util.stream.Collectors;
-import org.junit.jupiter.api.Test;
+import org.testng.annotations.Test;
 
 public class TagVisitorTest {