From ea9061a6316fefb023812d3a507b1ec3ee98fc7b Mon Sep 17 00:00:00 2001 From: Justin Donn Date: Mon, 21 Oct 2024 16:11:56 -0700 Subject: [PATCH] feat(relationship): add support for extracting relationships from aspect metadata (#452) --- .../linkedin/metadata/dao/EbeanLocalDAO.java | 27 ++++++----- .../metadata/dao/utils/EBeanDAOUtils.java | 19 +++++--- .../metadata/dao/EbeanLocalDAOTest.java | 45 +++++++++++++++++-- .../metadata/dao/utils/EBeanDAOUtilsTest.java | 18 +++++--- ...notatedAspectBarWithRelationshipFields.pdl | 10 +++++ .../localrelationship/AspectFooBar.pdl | 1 + .../testing/localrelationship/BelongsTo.pdl | 1 + 7 files changed, 95 insertions(+), 26 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 2f70cac78..82c6cdac3 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 @@ -8,10 +8,12 @@ import com.linkedin.data.template.RecordTemplate; import com.linkedin.data.template.SetMode; import com.linkedin.data.template.UnionTemplate; +import com.linkedin.metadata.dao.builder.BaseLocalRelationshipBuilder; import com.linkedin.metadata.dao.builder.BaseLocalRelationshipBuilder.LocalRelationshipUpdates; import com.linkedin.metadata.dao.builder.LocalRelationshipBuilderRegistry; import com.linkedin.metadata.dao.exception.ModelConversionException; import com.linkedin.metadata.dao.exception.RetryLimitReached; +import com.linkedin.metadata.dao.internal.BaseGraphWriterDAO; import com.linkedin.metadata.dao.producer.BaseMetadataEventProducer; import com.linkedin.metadata.dao.producer.BaseTrackingMetadataEventProducer; import com.linkedin.metadata.dao.retention.TimeBasedRetention; @@ -136,6 +138,7 @@ public boolean isChangeLogEnabled() { /** * Set where the relationships should be derived from during ingestion. Either from aspect models or from relationship * builders. The default is relationship builders. This should only be used during DAO instantiation i.e. in the DAO factory. + * One limitation when setting this is that all aspects for a particular entity type must use the same relationship source config. * @param relationshipSource {@link RelationshipSource ASPECT_METADATA or RELATIONSHIP_BUILDERS} */ public void setRelationshipSource(RelationshipSource relationshipSource) { @@ -609,7 +612,7 @@ protected long saveLatest(@Nonnull URN urn, @Non } if (_relationshipSource == RelationshipSource.ASPECT_METADATA) { - // TODO: not yet implemented + // TODO: not yet implemented -> add support for removing relationships when the aspect is to be soft-deleted throw new UnsupportedOperationException("This method has not been implemented yet to support the " + "ASPECT_METADATA RelationshipSource type yet."); } @@ -882,24 +885,26 @@ protected void insert(@Nonnull URN urn, @Nullabl * @param isTestMode Whether the test mode is enabled or not * @return List of LocalRelationshipUpdates that were executed */ - public List addRelationshipsIfAny(@Nonnull URN urn, @Nullable ASPECT aspect, - @Nonnull Class aspectClass, boolean isTestMode) { + public List addRelationshipsIfAny( + @Nonnull URN urn, @Nullable ASPECT aspect, @Nonnull Class aspectClass, boolean isTestMode) { + List localRelationshipUpdates = Collections.emptyList(); if (_relationshipSource == RelationshipSource.ASPECT_METADATA) { - // TODO: not yet implemented - throw new UnsupportedOperationException("This method has not been implemented yet to support the " - + "ASPECT_METADATA RelationshipSource type yet."); + 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 BaseLocalRelationshipBuilder.LocalRelationshipUpdates( + relationships, relationships.get(0).getClass(), BaseGraphWriterDAO.RemovalOption.REMOVE_NONE)) + .collect(Collectors.toList()); } else if (_relationshipSource == RelationshipSource.RELATIONSHIP_BUILDERS) { if (_localRelationshipBuilderRegistry != null && _localRelationshipBuilderRegistry.isRegistered(aspectClass)) { - List localRelationshipUpdates = - _localRelationshipBuilderRegistry.getLocalRelationshipBuilder(aspect).buildRelationships(urn, aspect); - _localRelationshipWriterDAO.processLocalRelationshipUpdates(urn, localRelationshipUpdates, isTestMode); - return localRelationshipUpdates; + localRelationshipUpdates = _localRelationshipBuilderRegistry.getLocalRelationshipBuilder(aspect).buildRelationships(urn, aspect); } } else { throw new UnsupportedOperationException("Please ensure that the RelationshipSource enum is properly set using " + "setRelationshipSource method."); } - return Collections.emptyList(); + _localRelationshipWriterDAO.processLocalRelationshipUpdates(urn, localRelationshipUpdates, isTestMode); + return localRelationshipUpdates; } @Nonnull 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 83ce335e0..ade574786 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,6 +29,7 @@ import java.sql.ResultSet; import java.sql.Timestamp; import java.util.ArrayList; +import java.util.Collections; import java.util.List; import java.util.Map; import java.util.Objects; @@ -42,7 +43,7 @@ import org.json.simple.parser.JSONParser; import org.json.simple.parser.ParseException; -import static com.linkedin.metadata.annotations.GmaAnnotationParser.parseDeltaFields; +import static com.linkedin.metadata.annotations.GmaAnnotationParser.*; /** @@ -388,15 +389,23 @@ public static relationshipsList = (List) obj; ModelType modelType = parseModelTypeFromGmaAnnotation(relationshipsList.get(0)); if (modelType == ModelType.RELATIONSHIP) { - log.debug(String.format("Found {%d} relationships of type {%s} for field {%s} of aspect class {%s}.", - relationshipsList.size(), relationshipsList.get(0).getClass(), fieldName, clazz.getName())); + log.debug("Found {} relationships of type {} for field {} of aspect class {}.", + relationshipsList.size(), relationshipsList.get(0).getClass(), fieldName, clazz.getName()); return (List) relationshipsList; } } catch (ReflectiveOperationException e) { diff --git a/dao-impl/ebean-dao/src/test/java/com/linkedin/metadata/dao/EbeanLocalDAOTest.java b/dao-impl/ebean-dao/src/test/java/com/linkedin/metadata/dao/EbeanLocalDAOTest.java index 426c118e9..e412f93b5 100644 --- a/dao-impl/ebean-dao/src/test/java/com/linkedin/metadata/dao/EbeanLocalDAOTest.java +++ b/dao-impl/ebean-dao/src/test/java/com/linkedin/metadata/dao/EbeanLocalDAOTest.java @@ -25,13 +25,13 @@ import com.linkedin.metadata.dao.producer.BaseTrackingMetadataEventProducer; import com.linkedin.metadata.dao.retention.TimeBasedRetention; import com.linkedin.metadata.dao.retention.VersionBasedRetention; -import com.linkedin.metadata.dao.urnpath.UrnPathExtractor; import com.linkedin.metadata.dao.storage.LocalDAOStorageConfig; import com.linkedin.metadata.dao.tracking.BaseTrackingManager; +import com.linkedin.metadata.dao.urnpath.UrnPathExtractor; import com.linkedin.metadata.dao.utils.BarUrnPathExtractor; import com.linkedin.metadata.dao.utils.EbeanServerUtils; -import com.linkedin.metadata.dao.utils.FooUrnPathExtractor; import com.linkedin.metadata.dao.utils.EmbeddedMariaInstance; +import com.linkedin.metadata.dao.utils.FooUrnPathExtractor; import com.linkedin.metadata.dao.utils.ModelUtils; import com.linkedin.metadata.dao.utils.RecordUtils; import com.linkedin.metadata.dao.utils.SQLSchemaUtils; @@ -64,10 +64,10 @@ import com.linkedin.testing.MixedRecord; import com.linkedin.testing.localrelationship.AspectFooBar; import com.linkedin.testing.localrelationship.BelongsTo; +import com.linkedin.testing.localrelationship.BelongsToArray; import com.linkedin.testing.urn.BarUrn; import com.linkedin.testing.urn.BurgerUrn; import com.linkedin.testing.urn.FooUrn; - import io.ebean.Ebean; import io.ebean.EbeanServer; import io.ebean.ExpressionList; @@ -100,7 +100,6 @@ import javax.annotation.Nullable; import javax.persistence.OptimisticLockException; import javax.persistence.RollbackException; - import org.mockito.ArgumentMatchers; import org.mockito.InOrder; import org.mockito.MockedStatic; @@ -3067,6 +3066,44 @@ public void testAddWithLocalRelationshipBuilder() throws URISyntaxException { assertEquals(aspects.size(), 1); } + @Test + public void testAddRelationshipsFromAspect() throws URISyntaxException { + EbeanLocalDAO fooDao = createDao(FooUrn.class); + EbeanLocalDAO barDao = createDao(BarUrn.class); + + fooDao.setRelationshipSource(EbeanLocalDAO.RelationshipSource.ASPECT_METADATA); + + FooUrn fooUrn = makeFooUrn(1); + BarUrn barUrn1 = BarUrn.createFromString("urn:li:bar:1"); + BelongsTo belongsTo1 = new BelongsTo().setSource(barUrn1).setDestination(fooUrn); + BarUrn barUrn2 = BarUrn.createFromString("urn:li:bar:2"); + BelongsTo belongsTo2 = new BelongsTo().setSource(barUrn2).setDestination(fooUrn); + BarUrn barUrn3 = BarUrn.createFromString("urn:li:bar:3"); + BelongsTo belongsTo3 = new BelongsTo().setSource(barUrn3).setDestination(fooUrn); + BelongsToArray belongsToArray = new BelongsToArray(belongsTo1, belongsTo2, belongsTo3); + AspectFooBar aspectFooBar = new AspectFooBar().setBars(new BarUrnArray(barUrn1, barUrn2, barUrn3)).setBelongsTos(belongsToArray); + AuditStamp auditStamp = makeAuditStamp("foo", System.currentTimeMillis()); + + fooDao.add(fooUrn, aspectFooBar, auditStamp); + barDao.add(barUrn1, new AspectFoo().setValue("1"), auditStamp); + barDao.add(barUrn2, new AspectFoo().setValue("2"), auditStamp); + barDao.add(barUrn3, new AspectFoo().setValue("3"), auditStamp); + + // Verify local relationships and entity are added. + EbeanLocalRelationshipQueryDAO ebeanLocalRelationshipQueryDAO = new EbeanLocalRelationshipQueryDAO(_server); + ebeanLocalRelationshipQueryDAO.setSchemaConfig(_schemaConfig); + + List relationships = + ebeanLocalRelationshipQueryDAO.findRelationships(BarSnapshot.class, EMPTY_FILTER, FooSnapshot.class, + EMPTY_FILTER, BelongsTo.class, OUTGOING_FILTER, 0, 10); + + AspectKey key = new AspectKey<>(AspectFooBar.class, fooUrn, 0L); + List aspects = fooDao.batchGetHelper(Collections.singletonList(key), 1, 0); + + assertEquals(relationships.size(), 3); + assertEquals(aspects.size(), 1); + } + @Test public void testNewSchemaFilterByArray() { if (_schemaConfig == SchemaConfig.NEW_SCHEMA_ONLY) { 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 a1272e581..34372907e 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 @@ -602,7 +602,7 @@ public void testExtractRelationshipsFromAspect() { AnnotatedAspectFooWithRelationshipField fooWithNullRelationshipField = new AnnotatedAspectFooWithRelationshipField(); assertTrue(EBeanDAOUtils.extractRelationshipsFromAspect(fooWithNullRelationshipField).isEmpty()); - // case 4: aspect model contains multiple relationship fields, some null and some non-null, as well as array fields + // 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()); @@ -612,6 +612,8 @@ public void testExtractRelationshipsFromAspect() { // value -> "abc" // integers -> [1] // nonRelationshipStructs -> [commonAspect1] + // relationshipFoo1 -> foo1 + // relationshipFoo2 -> not present // relationshipFoos -> [foo1, foo2] // relationshipBars -> [bar1] // moreRelationshipFoos -> not present @@ -622,16 +624,20 @@ public void testExtractRelationshipsFromAspect() { .setValue("abc") .setIntegers(new IntegerArray(1)) .setNonRelationshipStructs(new CommonAspectArray(new CommonAspect())) + .setRelationshipFoo1(new AnnotatedRelationshipFoo()) + // don't set relationshipFoo2 fields .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(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(0).get(1)); - assertEquals(new AnnotatedRelationshipBar(), results.get(1).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)); } diff --git a/gradle-plugins/metadata-annotations-test-models/src/main/pegasus/com/linkedin/testing/AnnotatedAspectBarWithRelationshipFields.pdl b/gradle-plugins/metadata-annotations-test-models/src/main/pegasus/com/linkedin/testing/AnnotatedAspectBarWithRelationshipFields.pdl index 55466e65d..4bbc36855 100644 --- a/gradle-plugins/metadata-annotations-test-models/src/main/pegasus/com/linkedin/testing/AnnotatedAspectBarWithRelationshipFields.pdl +++ b/gradle-plugins/metadata-annotations-test-models/src/main/pegasus/com/linkedin/testing/AnnotatedAspectBarWithRelationshipFields.pdl @@ -21,6 +21,16 @@ record AnnotatedAspectBarWithRelationshipFields { */ nonRelationshipStructs: optional array[CommonAspect] + /** + * For unit tests + */ + relationshipFoo1: optional AnnotatedRelationshipFoo + + /** + * For unit tests + */ + relationshipFoo2: optional AnnotatedRelationshipFoo + /** * For unit tests */ diff --git a/testing/test-models/src/main/pegasus/com/linkedin/testing/localrelationship/AspectFooBar.pdl b/testing/test-models/src/main/pegasus/com/linkedin/testing/localrelationship/AspectFooBar.pdl index a0acbf474..d2517349d 100644 --- a/testing/test-models/src/main/pegasus/com/linkedin/testing/localrelationship/AspectFooBar.pdl +++ b/testing/test-models/src/main/pegasus/com/linkedin/testing/localrelationship/AspectFooBar.pdl @@ -5,4 +5,5 @@ import com.linkedin.testing.BarUrn @gma.aspect.column.name = "aspectfoobar" record AspectFooBar { bars: array[BarUrn] + belongsTos: optional array[BelongsTo] } \ No newline at end of file diff --git a/testing/test-models/src/main/pegasus/com/linkedin/testing/localrelationship/BelongsTo.pdl b/testing/test-models/src/main/pegasus/com/linkedin/testing/localrelationship/BelongsTo.pdl index c37252c08..fcb03ffdd 100644 --- a/testing/test-models/src/main/pegasus/com/linkedin/testing/localrelationship/BelongsTo.pdl +++ b/testing/test-models/src/main/pegasus/com/linkedin/testing/localrelationship/BelongsTo.pdl @@ -4,5 +4,6 @@ namespace com.linkedin.testing.localrelationship "destination": "com.linkedin.testing.urn.FooUrn", "source": "com.linkedin.testing.urn.BarUrn" } ] +@gma.model = "RELATIONSHIP" record BelongsTo includes BaseRelationship { } \ No newline at end of file