feat(model 2.0): implement extractRelationshipsFromAspect util method

dao-impl/ebean-dao/build.gradle

@@ -9,6 +9,7 @@ configurations {
 dependencies {
   compile project(':core-models-utils')
   compile project(':dao-api')
+  compile project(':gradle-plugins:metadata-annotations-lib')
   compile externalDependency.ebean
   compile externalDependency.flywayCore
   compile externalDependency.guava
@@ -20,6 +21,7 @@ dependencies {
 
   annotationProcessor externalDependency.lombok
 
+  testCompile project(':gradle-plugins:metadata-annotations-test-models')
   testCompile project(':testing:core-models-testing')
   testCompile project(':testing:test-models')
   testCompile externalDependency.mysql

dao-impl/ebean-dao/src/main/java/com/linkedin/metadata/dao/utils/EBeanDAOUtils.java

@@ -4,7 +4,11 @@
 import com.linkedin.data.schema.RecordDataSchema;
 import com.linkedin.data.template.DataTemplateUtil;
 import com.linkedin.data.template.RecordTemplate;
+import com.linkedin.metadata.annotations.AlwaysAllowList;
 import com.linkedin.metadata.annotations.DeltaEntityAnnotationArrayMap;
+import com.linkedin.metadata.annotations.GmaAnnotation;
+import com.linkedin.metadata.annotations.GmaAnnotationParser;
+import com.linkedin.metadata.annotations.ModelType;
 import com.linkedin.metadata.aspect.AuditedAspect;
 import com.linkedin.metadata.aspect.SoftDeletedAspect;
 import com.linkedin.metadata.dao.EbeanMetadataAspect;
@@ -27,10 +31,13 @@
 import java.util.ArrayList;
 import java.util.List;
 import java.util.Map;
+import java.util.Objects;
+import java.util.Optional;
 import java.util.stream.Collectors;
 import javax.annotation.Nonnull;
 import javax.annotation.Nullable;
 import lombok.extern.slf4j.Slf4j;
+import org.apache.commons.lang.StringUtils;
 import org.json.simple.JSONObject;
 import org.json.simple.parser.JSONParser;
 import org.json.simple.parser.ParseException;
@@ -366,4 +373,47 @@ public static LocalRelationshipCriterion buildRelationshipFieldCriterion(LocalRe
         .setValue(localRelationshipValue)
         .setCondition(condition);
   }
+
+  /**
+   * Extract the non-null values from all top-level relationship fields of an aspect.
+   * @param aspect aspect (possibly with relationships) to be ingested
+   * @return a list of relationship arrays, with each array representing the relationship(s) present in each top-level relationship
+   * field in the aspect. An empty list means that there is no non-null relationship metadata attached to the given aspect.
+   */
+  @Nonnull
+  public static <RELATIONSHIP extends RecordTemplate, ASPECT extends RecordTemplate> List<List<RELATIONSHIP>> extractRelationshipsFromAspect(ASPECT aspect) {
+    return aspect.schema().getFields().stream().filter(field -> !field.getType().isPrimitive()).map(field -> {
+      String fieldName = field.getName();
+      Class<RecordTemplate> clazz = (Class<RecordTemplate>) aspect.getClass();
+      try {
+        Method getMethod = clazz.getMethod("get" + StringUtils.capitalize(fieldName));
+        Object obj = getMethod.invoke(aspect);
+        if (!(obj instanceof List) || ((List) obj).isEmpty() || !(((List) obj).get(0) instanceof RecordTemplate)) {
+          return null;
+        }
+        List<RecordTemplate> relationshipsList = (List<RecordTemplate>) obj;
+        ModelType modelType = parseModelTypeFromGmaAnnotation(relationshipsList.get(0));
+        if (modelType == ModelType.RELATIONSHIP) {
+          return (List<RELATIONSHIP>) relationshipsList;
+        }
+      } catch (ReflectiveOperationException e) {
+        throw new RuntimeException(e);
+      }
+      return null;
+    }).filter(Objects::nonNull).collect(Collectors.toList());
+  }
+
+  // Using the GmaAnnotationParser, extract the model type from the @gma.model annotation on any models.
+  private static ModelType parseModelTypeFromGmaAnnotation(RecordTemplate model) {
+    try {
+      final RecordDataSchema schema = (RecordDataSchema) DataTemplateUtil.getSchema(model.getClass());
+      final Optional<GmaAnnotation> gmaAnnotation = new GmaAnnotationParser(new AlwaysAllowList()).parse(schema);
+      if (!gmaAnnotation.isPresent() || !gmaAnnotation.get().hasModel()) {
+        return null;
+      }
+      return gmaAnnotation.get().getModel();
+    } catch (Exception e) {
+      throw new RuntimeException(String.format("Failed to parse the annotations for field %s", model.getClass().getCanonicalName()), e);
+    }
+  }
 }

dao-impl/ebean-dao/src/test/java/com/linkedin/metadata/dao/utils/EBeanDAOUtilsTest.java

@@ -1,6 +1,7 @@
 package com.linkedin.metadata.dao.utils;
 
 import com.google.common.io.Resources;
+import com.linkedin.data.template.RecordTemplate;
 import com.linkedin.data.template.StringArray;
 import com.linkedin.metadata.aspect.AuditedAspect;
 import com.linkedin.metadata.dao.EbeanLocalAccess;
@@ -13,7 +14,13 @@
 import com.linkedin.metadata.query.LocalRelationshipValue;
 import com.linkedin.metadata.query.RelationshipField;
 import com.linkedin.metadata.query.UrnField;
+import com.linkedin.testing.AnnotatedAspectBarWithRelationshipFields;
+import com.linkedin.testing.AnnotatedRelationshipBar;
+import com.linkedin.testing.AnnotatedRelationshipBarArray;
+import com.linkedin.testing.AnnotatedRelationshipFoo;
+import com.linkedin.testing.AnnotatedRelationshipFooArray;
 import com.linkedin.testing.AspectFoo;
+import com.linkedin.testing.AnnotatedAspectFooWithRelationshipField;
 import com.linkedin.testing.urn.BurgerUrn;
 import com.linkedin.testing.urn.FooUrn;
 import io.ebean.Ebean;
@@ -33,6 +40,7 @@
 import org.testng.annotations.Test;
 
 import static com.linkedin.testing.TestUtils.*;
+import static org.junit.Assert.*;
 import static org.mockito.Mockito.*;
 import static org.testng.Assert.assertFalse;
 import static org.testng.Assert.assertTrue;
@@ -567,4 +575,55 @@ public void testBuildRelationshipFieldCriterionWithAspectField() {
         relationshipField);
     assertEquals(relationshipField, filterCriterion.getField().getRelationshipField());
   }
+
+  @Test
+  public void testExtractRelationshipsFromAspect() {
+    // case 1: aspect model does not contain any relationship typed fields
+    // expected: return null
+    AspectFoo foo = new AspectFoo();
+    assertTrue(EBeanDAOUtils.extractRelationshipsFromAspect(foo).isEmpty());
+
+    // case 2: aspect model contains only a non-null relationship type field
+    // expected: return the non-null relationship
+    AnnotatedRelationshipFoo relationshipFoo = new AnnotatedRelationshipFoo();
+    AnnotatedRelationshipFooArray relationshipFoos = new AnnotatedRelationshipFooArray(relationshipFoo);
+    AnnotatedAspectFooWithRelationshipField fooWithRelationshipField = new AnnotatedAspectFooWithRelationshipField()
+        .setRelationshipFoo(relationshipFoos);
+
+    List<List<RecordTemplate>> results = EBeanDAOUtils.extractRelationshipsFromAspect(fooWithRelationshipField);
+    assertEquals(1, results.size());
+    assertEquals(1, results.get(0).size());
+    assertEquals(relationshipFoo, results.get(0).get(0));
+
+    // case 3: aspect model contains only a null relationship type field
+    // expected: return null
+    AnnotatedAspectFooWithRelationshipField fooWithNullRelationshipField = new AnnotatedAspectFooWithRelationshipField();
+    assertTrue(EBeanDAOUtils.extractRelationshipsFromAspect(fooWithNullRelationshipField).isEmpty());
+
+    // case 4: aspect model contains multiple relationship fields, some null and some non-null
+    // expected: return only the non-null relationships
+    relationshipFoos = new AnnotatedRelationshipFooArray(new AnnotatedRelationshipFoo(), new AnnotatedRelationshipFoo());
+    AnnotatedRelationshipBarArray relationshipBars = new AnnotatedRelationshipBarArray(new AnnotatedRelationshipBar());
+    // given:
+    // aspect = {
+    //     value -> "abc"
+    //     relationshipFoos -> [foo1, foo2]
+    //     relationshipBars -> [bar1]
+    //     moreRelationshipFoos -> not present
+    // }
+    // expect:
+    // [[foo1, foo2], [bar1]]
+    AnnotatedAspectBarWithRelationshipFields barWithRelationshipFields = new AnnotatedAspectBarWithRelationshipFields()
+        .setValue("abc")
+        .setRelationshipFoos(relationshipFoos)
+        .setRelationshipBars(relationshipBars); // don't set moreRelationshipFoos field
+
+    results = EBeanDAOUtils.extractRelationshipsFromAspect(barWithRelationshipFields);
+    assertEquals(2, results.size());
+    assertEquals(2, results.get(0).size());
+    assertEquals(1, results.get(1).size());
+    assertEquals(new AnnotatedRelationshipFoo(), results.get(0).get(0));
+    assertEquals(new AnnotatedRelationshipFoo(), results.get(0).get(1));
+    assertEquals(new AnnotatedRelationshipBar(), results.get(1).get(0));
+  }
 }

gradle-plugins/metadata-annotations-schema/src/main/pegasus/com/linkedin/metadata/annotations/GmaAnnotation.pdl

@@ -14,6 +14,15 @@ record GmaAnnotation {
    */
   delta: optional DeltaAnnotation
 
+  /**
+   * The type of model
+   */
+  model: optional enum ModelType {
+    ASSET,
+    ASPECT,
+    RELATIONSHIP
+  }
+
   /**
    * Information about GMA Search functionality.
    */

gradle-plugins/metadata-annotations-test-models/src/main/pegasus/com/linkedin/testing/AnnotatedAspectBarWithRelationshipFields.pdl

@@ -0,0 +1,28 @@
+namespace com.linkedin.testing
+
+/**
+ * For unit tests
+ */
+@gma.aspect.column.name = "annotatedaspectbarwithrelationshipfields"
+@gma.model = "ASPECT"
+record AnnotatedAspectBarWithRelationshipFields {
+  /**
+   * For unit tests
+   */
+  value: string
+
+  /**
+   * For unit tests
+   */
+  relationshipFoos: optional array[AnnotatedRelationshipFoo]
+
+  /**
+   * For unit tests
+   */
+  relationshipBars: optional array[AnnotatedRelationshipBar]
+
+  /**
+   * For unit tests
+   */
+  moreRelationshipFoos: optional array[AnnotatedRelationshipFoo]
+}

gradle-plugins/metadata-annotations-test-models/src/main/pegasus/com/linkedin/testing/AnnotatedAspectFooWithRelationshipField.pdl

@@ -0,0 +1,18 @@
+namespace com.linkedin.testing
+
+/**
+ * For unit tests
+ */
+@gma.aspect.column.name = "annotatedaspectfoowithrelationshipfield"
+@gma.model = "ASPECT"
+record AnnotatedAspectFooWithRelationshipField {
+  /**
+   * For unit tests
+   */
+  value: string
+
+  /**
+   * For unit tests
+   */
+  relationshipFoo: optional array[AnnotatedRelationshipFoo]
+}

gradle-plugins/metadata-annotations-test-models/src/main/pegasus/com/linkedin/testing/AnnotatedRelationshipBar.pdl

@@ -0,0 +1,24 @@
+namespace com.linkedin.testing
+
+import com.linkedin.common.Urn
+
+/**
+ * For unit tests
+ */
+@pairings = [ {
+  "destination" : "com.linkedin.common.urn.Urn",
+  "source" : "com.linkedin.common.urn.Urn"
+} ]
+@gma.model = "RELATIONSHIP"
+record AnnotatedRelationshipBar {
+
+  /**
+   * For unit tests
+   */
+  source: Urn
+
+  /**
+   * For unit tests
+   */
+  destination: Urn
+}

gradle-plugins/metadata-annotations-test-models/src/main/pegasus/com/linkedin/testing/AnnotatedRelationshipFoo.pdl

@@ -0,0 +1,24 @@
+namespace com.linkedin.testing
+
+import com.linkedin.common.Urn
+
+/**
+ * For unit tests
+ */
+@pairings = [ {
+  "destination" : "com.linkedin.common.urn.Urn",
+  "source" : "com.linkedin.common.urn.Urn"
+} ]
+@gma.model = "RELATIONSHIP"
+record AnnotatedRelationshipFoo {
+
+  /**
+   * For unit tests
+   */
+  source: Urn
+
+  /**
+   * For unit tests
+   */
+  destination: Urn
+}