From af8812b8cdf22b03b2115b03db6267414716d565 Mon Sep 17 00:00:00 2001 From: Justin Donn Date: Wed, 13 Nov 2024 10:09:17 -0800 Subject: [PATCH] fix(relationship 2.0): handle the case where aspects can have multiple fields with the same relationship type (#476) --- .../linkedin/metadata/dao/EbeanLocalDAO.java | 15 +++++---- .../metadata/dao/utils/EBeanDAOUtils.java | 26 ++++++++------- .../metadata/dao/utils/EBeanDAOUtilsTest.java | 33 +++++++++++-------- 3 files changed, 42 insertions(+), 32 deletions(-) diff --git a/dao-impl/ebean-dao/src/main/java/com/linkedin/metadata/dao/EbeanLocalDAO.java b/dao-impl/ebean-dao/src/main/java/com/linkedin/metadata/dao/EbeanLocalDAO.java index c37189965..7bb9271d1 100644 --- a/dao-impl/ebean-dao/src/main/java/com/linkedin/metadata/dao/EbeanLocalDAO.java +++ b/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 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 relationships = extractRelationshipsFromAspect(oldValue).stream() - .flatMap(List::stream) + List 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 List } List localRelationshipUpdates = Collections.emptyList(); if (_relationshipSource == RelationshipSource.ASPECT_METADATA) { - List> 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, Set> 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)) { diff --git a/dao-impl/ebean-dao/src/main/java/com/linkedin/metadata/dao/utils/EBeanDAOUtils.java b/dao-impl/ebean-dao/src/main/java/com/linkedin/metadata/dao/utils/EBeanDAOUtils.java index ade574786..510411ff7 100644 --- a/dao-impl/ebean-dao/src/main/java/com/linkedin/metadata/dao/utils/EBeanDAOUtils.java +++ b/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 List> extractRelationshipsFromAspect(ASPECT aspect) { - return aspect.schema().getFields().stream().filter(field -> !field.getType().isPrimitive()).map(field -> { + public static Map, Set> + extractRelationshipsFromAspect(ASPECT aspect) { + Map, Set> relationshipMap = new HashMap<>(); + aspect.schema().getFields().stream().filter(field -> !field.getType().isPrimitive()).forEach(field -> { String fieldName = field.getName(); Class clazz = (Class) aspect.getClass(); try { @@ -396,23 +399,24 @@ public static new HashSet<>()).add((RELATIONSHIP) obj); + return; } } else if (!(obj instanceof List) || ((List) obj).isEmpty() || !(((List) obj).get(0) instanceof RecordTemplate)) { - return null; + return; } List relationshipsList = (List) 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) relationshipsList; + relationshipMap.computeIfAbsent(relationshipsList.get(0).getClass(), k -> new HashSet<>()).addAll((List) 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. diff --git a/dao-impl/ebean-dao/src/test/java/com/linkedin/metadata/dao/utils/EBeanDAOUtilsTest.java b/dao-impl/ebean-dao/src/test/java/com/linkedin/metadata/dao/utils/EBeanDAOUtilsTest.java index 132a77b8f..c2d7ee2ed 100644 --- a/dao-impl/ebean-dao/src/test/java/com/linkedin/metadata/dao/utils/EBeanDAOUtilsTest.java +++ b/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> results = EBeanDAOUtils.extractRelationshipsFromAspect(fooWithRelationshipField); + Map, Set> 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())); } } \ No newline at end of file