feat(relationship): add support for extracting relationships from asp…

…ect metadata (#452)

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 <ASPECT extends RecordTemplate> 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 <ASPECT extends RecordTemplate> void insert(@Nonnull URN urn, @Nullabl
    * @param isTestMode Whether the test mode is enabled or not
    * @return List of LocalRelationshipUpdates that were executed
    */
-  public <ASPECT extends RecordTemplate> List<LocalRelationshipUpdates> addRelationshipsIfAny(@Nonnull URN urn, @Nullable ASPECT aspect,
-      @Nonnull Class<ASPECT> aspectClass, boolean isTestMode) {
+  public <ASPECT extends RecordTemplate, RELATIONSHIP extends RecordTemplate> List<LocalRelationshipUpdates> addRelationshipsIfAny(
+      @Nonnull URN urn, @Nullable ASPECT aspect, @Nonnull Class<ASPECT> aspectClass, boolean isTestMode) {
+    List<LocalRelationshipUpdates> 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<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 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> 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

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 <RELATIONSHIP extends RecordTemplate, ASPECT extends RecordTemplat
       try {
         Method getMethod = clazz.getMethod("get" + StringUtils.capitalize(fieldName)); // getFieldName
         Object obj = getMethod.invoke(aspect); // invoke getFieldName
-        // all relationship fields will be represented as Arrays so filter out any non-lists, empty lists, and lists that don't contain RecordTemplates
-        if (!(obj instanceof List) || ((List) obj).isEmpty() || !(((List) obj).get(0) instanceof RecordTemplate)) {
+        // all relationship fields will be represented as either singleton fields or Arrays so filter out any non-RecordTemplates,
+        // empty lists, and lists that don't contain RecordTemplates
+        if (obj instanceof RecordTemplate) {
+          ModelType modelType = parseModelTypeFromGmaAnnotation((RecordTemplate) obj);
+          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);
+          }
+        } else 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) {
-          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<RELATIONSHIP>) relationshipsList;
         }
       } catch (ReflectiveOperationException e) {

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<EntityAspectUnion, FooUrn> fooDao = createDao(FooUrn.class);
+    EbeanLocalDAO<EntityAspectUnion, BarUrn> 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<BelongsTo> relationships =
+        ebeanLocalRelationshipQueryDAO.findRelationships(BarSnapshot.class, EMPTY_FILTER, FooSnapshot.class,
+            EMPTY_FILTER, BelongsTo.class, OUTGOING_FILTER, 0, 10);
+
+    AspectKey<FooUrn, AspectFooBar> key = new AspectKey<>(AspectFooBar.class, fooUrn, 0L);
+    List<EbeanMetadataAspect> 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) {

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));
   }
 
 

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
    */

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]
 }

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 {
 }