Skip to content

Commit

Permalink
feat(model 2.0): implement extractRelationshipsFromAspect util method
Browse files Browse the repository at this point in the history
  • Loading branch information
jsdonn committed Oct 14, 2024
1 parent c6be657 commit aab78b0
Show file tree
Hide file tree
Showing 8 changed files with 214 additions and 0 deletions.
2 changes: 2 additions & 0 deletions dao-impl/ebean-dao/build.gradle
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand All @@ -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;
Expand Down Expand Up @@ -366,4 +373,47 @@ public static LocalRelationshipCriterion buildRelationshipFieldCriterion(LocalRe
.setValue(localRelationshipValue)
.setCondition(condition);
}

/**
* Extract the non-null values from all 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 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);
}
}
}
Original file line number Diff line number Diff line change
@@ -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;
Expand All @@ -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;
Expand All @@ -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;
Expand Down Expand Up @@ -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));
}
}
Original file line number Diff line number Diff line change
Expand Up @@ -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.
*/
Expand Down
Original file line number Diff line number Diff line change
@@ -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]
}
Original file line number Diff line number Diff line change
@@ -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]
}
Original file line number Diff line number Diff line change
@@ -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
}
Original file line number Diff line number Diff line change
@@ -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
}

0 comments on commit aab78b0

Please sign in to comment.