feat(ingest): add import descriptions of two or more nested messages (d…

…atahub-project#6959)

Co-authored-by: 서주현[G플레이스데이터개발] <juhyun.seo@navercorp.com>

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

@@ -24,8 +24,7 @@ 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
-                        && messageProto() == fileProto().getMessageType(loc.getPath(1)));
+                        && loc.getPath(0) == FileDescriptorProto.MESSAGE_TYPE_FIELD_NUMBER);
     }
 
     <T> Stream<T> accept(ProtobufModelVisitor<T> v, VisitContext context);

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

@@ -21,6 +21,8 @@
 import lombok.Builder;
 import lombok.Getter;
 
+import java.util.ArrayList;
+import java.util.List;
 import java.util.Optional;
 import java.util.stream.Collectors;
 import java.util.stream.Stream;
@@ -37,6 +39,7 @@ public class ProtobufField implements ProtobufElement {
     private final String fieldPathType;
     private final Boolean isMessageType;
     private final SchemaFieldDataType schemaFieldDataType;
+    private final Boolean isNestedType;
 
     public OneofDescriptorProto oneOfProto() {
         if (fieldProto.hasOneofIndex()) {
@@ -208,14 +211,59 @@ public SchemaFieldDataType schemaFieldDataType() throws IllegalStateException {
     @Override
     public String comment() {
         return messageLocations()
-                .filter(loc -> loc.getPathCount() > 3
-                        && loc.getPath(2) == DescriptorProto.FIELD_FIELD_NUMBER
-                        && fieldProto == messageProto().getField(loc.getPath(3)))
+                .filter(location -> location.getPathCount() > 3)
+                .filter(location -> !ProtobufUtils.collapseLocationComments(location).isEmpty()
+                        && !isEnumType(location.getPathList()))
+                .filter(location -> {
+                    List<Integer> pathList = location.getPathList();
+                    DescriptorProto messageType = fileProto().getMessageType(pathList.get(1));
+
+                    if (!isNestedType
+                            && location.getPath(2) == DescriptorProto.FIELD_FIELD_NUMBER
+                            && fieldProto == messageType.getField(location.getPath(3))) {
+                        return true;
+                    } else if (isNestedType
+                            && location.getPath(2) == DescriptorProto.NESTED_TYPE_FIELD_NUMBER
+                            && fieldProto == getNestedTypeFields(pathList, messageType)) {
+                        return true;
+                    }
+                    return false;
+                })
                 .map(ProtobufUtils::collapseLocationComments)
                 .collect(Collectors.joining("\n"))
                 .trim();
     }
 
+    private FieldDescriptorProto getNestedTypeFields(List<Integer> pathList, DescriptorProto messageType) {
+        int pathSize = pathList.size();
+        List<Integer> nestedValues = new ArrayList<>(pathSize);
+
+        for (int index = 0; index < pathSize; index++) {
+            if (index > 1
+                    && index % 2 == 0
+                    && pathList.get(index) == DescriptorProto.NESTED_TYPE_FIELD_NUMBER) {
+                nestedValues.add(pathList.get(index + 1));
+            }
+        }
+
+        for (Integer value : nestedValues) {
+            messageType = messageType.getNestedType(value);
+        }
+
+        return messageType.getField(pathList.get(pathList.size() - 1));
+    }
+
+    private boolean isEnumType(List<Integer> pathList) {
+        for (int index = 0; index < pathList.size(); index++) {
+            if (index > 1
+                    && index % 2 == 0
+                    && pathList.get(index) == DescriptorProto.ENUM_TYPE_FIELD_NUMBER) {
+                return true;
+            }
+        }
+        return false;
+    }
+
     @Override
     public <T> Stream<T> accept(ProtobufModelVisitor<T> visitor, VisitContext context) {
         return visitor.visitField(this, context);

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

@@ -170,6 +170,7 @@ private void buildProtobufGraph(DescriptorProtos.FileDescriptorSet fileSet) {
                         ProtobufField fieldVertex = ProtobufField.builder()
                                 .protobufMessage(messageVertex)
                                 .fieldProto(fieldProto)
+                                .isNestedType(false)
                                 .build();
 
                         // Add field vertex
@@ -210,6 +211,10 @@ private void addEnum(DescriptorProtos.FileDescriptorProto fileProto, DescriptorP
     }
 
     private void addNestedMessage(DescriptorProtos.FileDescriptorProto fileProto, DescriptorProtos.DescriptorProto messageProto) {
+        if (messageProto.getNestedTypeCount() < 1) {
+            return;
+        }
+
         messageProto.getNestedTypeList().forEach(nestedMessageProto -> {
             ProtobufMessage nestedMessageVertex = ProtobufMessage.builder()
                     .fileProto(fileProto)
@@ -222,6 +227,7 @@ private void addNestedMessage(DescriptorProtos.FileDescriptorProto fileProto, De
                 ProtobufField field = ProtobufField.builder()
                         .protobufMessage(nestedMessageVertex)
                         .fieldProto(nestedFieldProto)
+                        .isNestedType(true)
                         .build();
 
                 // Add field vertex
@@ -236,6 +242,8 @@ private void addNestedMessage(DescriptorProtos.FileDescriptorProto fileProto, De
                             .build().inGraph(this);
                 }
             });
+
+            addNestedMessage(fileProto, nestedMessageProto);
         });
     }
 

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

@@ -19,7 +19,7 @@ public class ProtobufOneOfField extends ProtobufField {
     @Builder(builderMethodName = "oneOfBuilder")
     public ProtobufOneOfField(ProtobufMessage protobufMessage,
                               FieldDescriptorProto fieldProto) {
-        super(protobufMessage, fieldProto, null, null, null, null);
+        super(protobufMessage, fieldProto, null, null, null, null, null);
     }
 
     @Override

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

@@ -4,6 +4,7 @@
 import com.google.protobuf.DescriptorProtos.FieldDescriptorProto;
 import com.google.protobuf.DescriptorProtos.FileDescriptorProto;
 import com.linkedin.data.template.StringArray;
+
 import com.linkedin.schema.ArrayType;
 import com.linkedin.schema.BooleanType;
 import com.linkedin.schema.BytesType;
@@ -12,12 +13,17 @@
 import com.linkedin.schema.NumberType;
 import com.linkedin.schema.RecordType;
 import com.linkedin.schema.SchemaFieldDataType;
+import com.linkedin.schema.SchemaField;
+import com.linkedin.schema.SchemaMetadata;
 import com.linkedin.schema.StringType;
+import datahub.protobuf.ProtobufDataset;
 import org.junit.Test;
 
+import java.io.IOException;
 import java.util.Arrays;
 import java.util.Set;
 
+import static datahub.protobuf.TestFixtures.*;
 import static org.junit.jupiter.api.Assertions.*;
 
 
@@ -227,4 +233,27 @@ public void schemaFieldTypeArrayTest() {
                     .setNestedType(new StringArray()))), test.schemaFieldDataType());
         });
     }
+
+    @Test
+    public void nestedTypeFieldTest() throws IOException {
+        ProtobufDataset test = getTestProtobufDataset("extended_protobuf", "messageC");
+        SchemaMetadata testMetadata = test.getSchemaMetadata();
+
+        SchemaField nicknameField = testMetadata.getFields()
+                .stream()
+                .filter(f -> f.getFieldPath()
+                        .equals("[version=2.0].[type=extended_protobuf_UserMsg].[type=extended_protobuf_UserMsg_UserInfo].user_info.[type=string].nickname"))
+                .findFirst()
+                .orElseThrow();
+
+        assertEquals("nickname info", nicknameField.getDescription());
+
+        SchemaField profileUrlField = testMetadata.getFields()
+                .stream().filter(f -> f.getFieldPath()
+                        .equals("[version=2.0].[type=extended_protobuf_UserMsg].[type=extended_protobuf_UserMsg_UserInfo].user_info.[type=string].profile_url"))
+                .findFirst()
+                .orElseThrow();
+
+        assertEquals("profile url info", profileUrlField.getDescription());
+    }
 }

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

@@ -0,0 +1,14 @@
+syntax = "proto3";
+package extended_protobuf;
+
+
+message UserMsg {
+  message UserInfo {
+    string nickname = 1; // nickname info
+    string profile_url = 2; // profile url info
+  }
+
+  string id = 1; // user id
+  string name = 2; // user name
+  UserInfo user_info = 3; // user info
+}