fix(relationship 2.0): handle the case where aspects can have multipl…

…e fields with the same relationship type (linkedin#476)

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

@@ -48,6 +48,7 @@
 import java.net.URISyntaxException;
 import java.sql.Timestamp;
 import java.util.ArrayList;
+import java.util.Arrays;
 import java.util.Collections;
 import java.util.HashMap;
 import java.util.List;
@@ -650,8 +651,8 @@ protected <ASPECT extends RecordTemplate> long saveLatest(@Nonnull URN urn, @Non
     // If the aspect is to be soft deleted and we are deriving relationships from aspect metadata, remove any relationships
     // associated with the previous aspect value.
     if (newValue == null && _relationshipSource == RelationshipSource.ASPECT_METADATA && oldValue != null) {
-      List<RecordTemplate> relationships = extractRelationshipsFromAspect(oldValue).stream()
-          .flatMap(List::stream)
+      List<RecordTemplate> relationships = extractRelationshipsFromAspect(oldValue).values().stream()
+          .flatMap(Set::stream)
           .collect(Collectors.toList());
       _localRelationshipWriterDAO.removeRelationshipsV2(relationships, urn);
     // Otherwise, add any local relationships that are derived from the aspect.
@@ -892,11 +893,11 @@ public <ASPECT extends RecordTemplate, RELATIONSHIP extends RecordTemplate> List
     }
     List<LocalRelationshipUpdates> localRelationshipUpdates = Collections.emptyList();
     if (_relationshipSource == RelationshipSource.ASPECT_METADATA) {
-      List<List<RELATIONSHIP>> allRelationships = EBeanDAOUtils.extractRelationshipsFromAspect(aspect);
-      localRelationshipUpdates = allRelationships.stream()
-          .filter(relationships -> !relationships.isEmpty()) // ensure at least 1 relationship in sublist to avoid index out of bounds
-          .map(relationships -> new LocalRelationshipUpdates(
-            relationships, relationships.get(0).getClass(), BaseGraphWriterDAO.RemovalOption.REMOVE_ALL_EDGES_FROM_SOURCE))
+      Map<Class<?>, Set<RELATIONSHIP>> allRelationships = EBeanDAOUtils.extractRelationshipsFromAspect(aspect);
+      localRelationshipUpdates = allRelationships.entrySet().stream()
+          .filter(entry -> !entry.getValue().isEmpty()) // ensure at least 1 relationship in sublist to avoid index out of bounds
+          .map(entry -> new LocalRelationshipUpdates(
+              Arrays.asList(entry.getValue().toArray()), entry.getKey(), BaseGraphWriterDAO.RemovalOption.REMOVE_ALL_EDGES_FROM_SOURCE))
           .collect(Collectors.toList());
     } else if (_relationshipSource == RelationshipSource.RELATIONSHIP_BUILDERS) {
       if (_localRelationshipBuilderRegistry != null && _localRelationshipBuilderRegistry.isRegistered(aspectClass)) {

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

@@ -29,11 +29,12 @@
 import java.sql.ResultSet;
 import java.sql.Timestamp;
 import java.util.ArrayList;
-import java.util.Collections;
+import java.util.HashMap;
+import java.util.HashSet;
 import java.util.List;
 import java.util.Map;
-import java.util.Objects;
 import java.util.Optional;
+import java.util.Set;
 import java.util.stream.Collectors;
 import javax.annotation.Nonnull;
 import javax.annotation.Nullable;
@@ -378,12 +379,14 @@ public static LocalRelationshipCriterion buildRelationshipFieldCriterion(LocalRe
   /**
    * 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.
+   * @return a map of relationship class to relationship sets, with each set representing the relationship(s) of a particular type present in
+   *     all the top-level relationship fields in the aspect. An empty map 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 -> {
+  public static <RELATIONSHIP extends RecordTemplate, ASPECT extends RecordTemplate> Map<Class<?>, Set<RELATIONSHIP>>
+  extractRelationshipsFromAspect(ASPECT aspect) {
+    Map<Class<?>, Set<RELATIONSHIP>> relationshipMap = new HashMap<>();
+    aspect.schema().getFields().stream().filter(field -> !field.getType().isPrimitive()).forEach(field -> {
       String fieldName = field.getName();
       Class<RecordTemplate> clazz = (Class<RecordTemplate>) aspect.getClass();
       try {
@@ -396,23 +399,24 @@ public static <RELATIONSHIP extends RecordTemplate, ASPECT extends RecordTemplat
           if (modelType == ModelType.RELATIONSHIP) {
             log.debug("Found {} relationship(s) of type {} for field {} of aspect class {}.",
                 1, obj.getClass(), fieldName, clazz.getName());
-            return Collections.singletonList((RELATIONSHIP) obj);
+            relationshipMap.computeIfAbsent(obj.getClass(), k -> new HashSet<>()).add((RELATIONSHIP) obj);
+            return;
           }
         } else if (!(obj instanceof List) || ((List) obj).isEmpty() || !(((List) obj).get(0) instanceof RecordTemplate)) {
-          return null;
+          return;
         }
         List<RecordTemplate> relationshipsList = (List<RecordTemplate>) obj;
         ModelType modelType = parseModelTypeFromGmaAnnotation(relationshipsList.get(0));
         if (modelType == ModelType.RELATIONSHIP) {
           log.debug("Found {} relationships of type {} for field {} of aspect class {}.",
               relationshipsList.size(), relationshipsList.get(0).getClass(), fieldName, clazz.getName());
-          return (List<RELATIONSHIP>) relationshipsList;
+          relationshipMap.computeIfAbsent(relationshipsList.get(0).getClass(), k -> new HashSet<>()).addAll((List<RELATIONSHIP>) relationshipsList);
         }
       } catch (ReflectiveOperationException e) {
         throw new RuntimeException(e);
       }
-      return null;
-    }).filter(Objects::nonNull).collect(Collectors.toList());
+    });
+    return relationshipMap;
   }
 
   // Using the GmaAnnotationParser, extract the model type from the @gma.model annotation on any models.

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.common.urn.Urn;
 import com.linkedin.data.template.IntegerArray;
 import com.linkedin.data.template.RecordTemplate;
 import com.linkedin.data.template.StringArray;
@@ -42,6 +43,8 @@
 import java.util.ArrayList;
 import java.util.Collections;
 import java.util.List;
+import java.util.Map;
+import java.util.Set;
 import javax.annotation.Nonnull;
 import org.testng.annotations.Test;
 
@@ -601,7 +604,7 @@ public void testBuildRelationshipFieldCriterionWithAspectField() {
   }
 
   @Test
-  public void testExtractRelationshipsFromAspect() {
+  public void testExtractRelationshipsFromAspect() throws URISyntaxException {
     // case 1: aspect model does not contain any relationship typed fields
     // expected: return null
     AspectFoo foo = new AspectFoo();
@@ -614,10 +617,10 @@ public void testExtractRelationshipsFromAspect() {
     AnnotatedAspectFooWithRelationshipField fooWithRelationshipField = new AnnotatedAspectFooWithRelationshipField()
         .setRelationshipFoo(relationshipFoos);
 
-    List<List<RecordTemplate>> results = EBeanDAOUtils.extractRelationshipsFromAspect(fooWithRelationshipField);
+    Map<Class<?>, Set<RecordTemplate>> results = EBeanDAOUtils.extractRelationshipsFromAspect(fooWithRelationshipField);
     assertEquals(1, results.size());
-    assertEquals(1, results.get(0).size());
-    assertEquals(relationshipFoo, results.get(0).get(0));
+    assertEquals(1, results.get(AnnotatedRelationshipFoo.class).size());
+    assertTrue(results.get(AnnotatedRelationshipFoo.class).contains(relationshipFoo));
 
     // case 3: aspect model contains only a null relationship type field
     // expected: return null
@@ -627,7 +630,10 @@ public void testExtractRelationshipsFromAspect() {
     // case 4: aspect model contains multiple singleton and array-type relationship fields, some null and some non-null, as well as array fields
     //         containing non-Relationship objects
     // expected: return only the non-null relationships
-    relationshipFoos = new AnnotatedRelationshipFooArray(new AnnotatedRelationshipFoo(), new AnnotatedRelationshipFoo());
+    AnnotatedRelationshipFoo test1 = new AnnotatedRelationshipFoo().setDestination(Urn.createFromString("urn:li:test:1"));
+    AnnotatedRelationshipFoo test2 = new AnnotatedRelationshipFoo().setDestination(Urn.createFromString("urn:li:test:2"));
+    AnnotatedRelationshipFoo test3 = new AnnotatedRelationshipFoo().setDestination(Urn.createFromString("urn:li:test:3"));
+    relationshipFoos = new AnnotatedRelationshipFooArray(test1, test2);
     AnnotatedRelationshipBarArray relationshipBars = new AnnotatedRelationshipBarArray(new AnnotatedRelationshipBar());
     // given:
     // aspect = {
@@ -646,19 +652,18 @@ public void testExtractRelationshipsFromAspect() {
         .setValue("abc")
         .setIntegers(new IntegerArray(1))
         .setNonRelationshipStructs(new CommonAspectArray(new CommonAspect()))
-        .setRelationshipFoo1(new AnnotatedRelationshipFoo())
+        .setRelationshipFoo1(test3)
         // don't set relationshipFoo2 fields
         .setRelationshipFoos(relationshipFoos)
         .setRelationshipBars(relationshipBars); // don't set moreRelationshipFoos field
 
     results = EBeanDAOUtils.extractRelationshipsFromAspect(barWithRelationshipFields);
-    assertEquals(3, results.size());
-    assertEquals(1, results.get(0).size()); // relationshipFoo1
-    assertEquals(2, results.get(1).size()); // relationshipFoos
-    assertEquals(1, results.get(2).size()); // relationshipBars
-    assertEquals(new AnnotatedRelationshipFoo(), results.get(0).get(0));
-    assertEquals(new AnnotatedRelationshipFoo(), results.get(1).get(0));
-    assertEquals(new AnnotatedRelationshipFoo(), results.get(1).get(1));
-    assertEquals(new AnnotatedRelationshipBar(), results.get(2).get(0));
+    assertEquals(2, results.size());
+    assertEquals(3, results.get(AnnotatedRelationshipFoo.class).size()); // relationshipFoo1 (1) + relationshipFoos (2)
+    assertEquals(1, results.get(AnnotatedRelationshipBar.class).size()); // relationshipBars
+    assertTrue(results.get(AnnotatedRelationshipFoo.class).contains(test1));
+    assertTrue(results.get(AnnotatedRelationshipFoo.class).contains(test2));
+    assertTrue(results.get(AnnotatedRelationshipFoo.class).contains(test3));
+    assertTrue(results.get(AnnotatedRelationshipBar.class).contains(new AnnotatedRelationshipBar()));
   }
 }