fix(ingest): add message proto check (datahub-project#7130)

Co-authored-by: Tino Lyu <[email protected]>
Co-authored-by: david-leifker <[email protected]>

metadata-integration/java/datahub-protobuf/src/main/java/datahub/protobuf/model/ProtobufElement.java

@@ -24,7 +24,8 @@ default Stream<SourceCodeInfo.Location> messageLocations() {
         List<SourceCodeInfo.Location> fileLocations = fileProto().getSourceCodeInfo().getLocationList();
         return fileLocations.stream()
                 .filter(loc -> loc.getPathCount() > 1
-                        && loc.getPath(0) == FileDescriptorProto.MESSAGE_TYPE_FIELD_NUMBER);
+                        && loc.getPath(0) == FileDescriptorProto.MESSAGE_TYPE_FIELD_NUMBER
+                        && messageProto() == fileProto().getMessageType(loc.getPath(1)));
     }
 
     <T> Stream<T> accept(ProtobufModelVisitor<T> v, VisitContext context);

metadata-integration/java/datahub-protobuf/src/main/java/datahub/protobuf/model/ProtobufField.java

@@ -4,6 +4,7 @@
 import com.google.protobuf.DescriptorProtos.FieldDescriptorProto;
 import com.google.protobuf.DescriptorProtos.FileDescriptorProto;
 import com.google.protobuf.DescriptorProtos.OneofDescriptorProto;
+import com.google.protobuf.DescriptorProtos.SourceCodeInfo;
 import com.linkedin.data.template.StringArray;
 import com.linkedin.schema.ArrayType;
 import com.linkedin.schema.BooleanType;
@@ -208,6 +209,14 @@ public SchemaFieldDataType schemaFieldDataType() throws IllegalStateException {
         });
     }
 
+    @Override
+    public Stream<SourceCodeInfo.Location> messageLocations() {
+        List<SourceCodeInfo.Location> fileLocations = fileProto().getSourceCodeInfo().getLocationList();
+        return fileLocations.stream()
+                .filter(loc -> loc.getPathCount() > 1
+                        && loc.getPath(0) == FileDescriptorProto.MESSAGE_TYPE_FIELD_NUMBER);
+    }
+
     @Override
     public String comment() {
         return messageLocations()

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

@@ -63,6 +63,10 @@ public static ProtobufGraph getTestProtobufGraph(String protoPackage, String fil
         return new ProtobufGraph(getTestProtobufFileSet(protoPackage, filename));
     }
 
+    public static ProtobufGraph getTestProtobufGraph(String protoPackage, String filename, String messageName) throws IOException {
+        return new ProtobufGraph(getTestProtobufFileSet(protoPackage, filename), messageName);
+    }
+
     public static Object extractAspect(MetadataChangeProposalWrapper<? extends RecordTemplate> mcp, String aspect) {
         return mcp.getAspect().data().get(aspect);
     }

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

@@ -16,11 +16,11 @@ public class DescriptionVisitorTest {
 
     @Test
     public void visitorTest() throws IOException {
-        ProtobufGraph graph = getTestProtobufGraph("protobuf", "messageB");
+        ProtobufGraph graph = getTestProtobufGraph("protobuf", "messageC2", "protobuf.MessageC2");
 
         DescriptionVisitor test = new DescriptionVisitor();
 
-        assertEquals(Set.of("This contains nested types.\n\nOwned by TeamB"),
-                graph.accept(getVisitContextBuilder("protobuf.MessageB"), List.of(test)).collect(Collectors.toSet()));
+        assertEquals(Set.of("This contains nested type\n\nDescription for MessageC2"),
+                graph.accept(getVisitContextBuilder("protobuf.MessageC2"), List.of(test)).collect(Collectors.toSet()));
     }
 }

metadata-integration/java/datahub-protobuf/src/test/resources/protobuf/messageC2.proto

@@ -8,6 +8,11 @@ message MessageC1 {
   MessageC2 messageList = 1;
 }
 
+/*
+   This contains nested type
+
+   Description for MessageC2
+*/
 message MessageC2 {
   repeated MessageC3 list = 1;
 }