From ed967272d1eff1c1061a130eea7370c44ed43640 Mon Sep 17 00:00:00 2001 From: Justin Donn Date: Tue, 12 Aug 2025 17:53:12 -0700 Subject: [PATCH] feat(collections): add annotations for handling collections --- .../annotations/GmaAnnotationParserTest.java | 30 +++++++++++++++++ .../annotations/CollectionAnnotation.pdl | 30 +++++++++++++++++ .../metadata/annotations/GmaAnnotation.pdl | 6 ++++ .../testing/CollectionAnnotatedAspectBar.pdl | 32 +++++++++++++++++++ 4 files changed, 98 insertions(+) create mode 100644 gradle-plugins/metadata-annotations-schema/src/main/pegasus/com/linkedin/metadata/annotations/CollectionAnnotation.pdl create mode 100644 gradle-plugins/metadata-annotations-test-models/src/main/pegasus/com/linkedin/testing/CollectionAnnotatedAspectBar.pdl diff --git a/gradle-plugins/metadata-annotations-lib/src/test/java/com/linkedin/metadata/annotations/GmaAnnotationParserTest.java b/gradle-plugins/metadata-annotations-lib/src/test/java/com/linkedin/metadata/annotations/GmaAnnotationParserTest.java index c1fed15b7..0fbb16ca5 100644 --- a/gradle-plugins/metadata-annotations-lib/src/test/java/com/linkedin/metadata/annotations/GmaAnnotationParserTest.java +++ b/gradle-plugins/metadata-annotations-lib/src/test/java/com/linkedin/metadata/annotations/GmaAnnotationParserTest.java @@ -5,10 +5,14 @@ import com.linkedin.data.schema.RecordDataSchema; import com.linkedin.data.template.DataTemplateUtil; import com.linkedin.data.template.StringArray; +import com.linkedin.data.template.StringArrayMap; +import com.linkedin.data.template.StringArrayMapMap; +import com.linkedin.data.template.StringArrayMapMapArray; import com.linkedin.testing.AnnotatedAspectBar; import com.linkedin.testing.AnnotatedAspectFoo; import com.linkedin.testing.BarAspect; import com.linkedin.testing.CommonAspect; +import com.linkedin.testing.CollectionAnnotatedAspectBar; import com.linkedin.testing.SearchAnnotatedAspectBar; import java.util.Optional; @@ -116,4 +120,30 @@ public void parseAspectWithOnlySearchIndexAnnotations() { } // TODO: if add support for disallowing certain search annotations, add tests for them + + @Test + public void parseAspectWithCollectionAnnotations() { + final Optional gma = + new GmaAnnotationParser().parse((RecordDataSchema) DataTemplateUtil.getSchema(CollectionAnnotatedAspectBar.class)); + + StringArray fooPaths = new StringArray("x.y", "x.z"); + StringArrayMap fooMap = new StringArrayMap(); + fooMap.put("paths", fooPaths); + StringArrayMapMap foo = new StringArrayMapMap(); + foo.put("foo", fooMap); + + StringArray barPaths = new StringArray("abc", "def"); + StringArrayMap barMap = new StringArrayMap(); + barMap.put("paths", barPaths); + StringArrayMapMap bar = new StringArrayMapMap(); + bar.put("bar", barMap); + + StringArrayMapMapArray primaryKeys = new StringArrayMapMapArray(foo, bar); + assertThat(gma).contains(new GmaAnnotation().setCollection( + new CollectionAnnotation() + .setIsCollection(true) + .setPrimaryKeys(primaryKeys) + .setDefaultUpdateBehavior(UpdateBehavior.REPLACE_BY_ACTOR) + )); + } } diff --git a/gradle-plugins/metadata-annotations-schema/src/main/pegasus/com/linkedin/metadata/annotations/CollectionAnnotation.pdl b/gradle-plugins/metadata-annotations-schema/src/main/pegasus/com/linkedin/metadata/annotations/CollectionAnnotation.pdl new file mode 100644 index 000000000..e829d0c2a --- /dev/null +++ b/gradle-plugins/metadata-annotations-schema/src/main/pegasus/com/linkedin/metadata/annotations/CollectionAnnotation.pdl @@ -0,0 +1,30 @@ +namespace com.linkedin.metadata.annotations + +/** + * Annotation for collections-related attributes + */ +record CollectionAnnotation { + /** + * Whether a model is a collection + */ + isCollection: optional boolean + + /** + * A list of key-values representing the primary keys to use for updating collection-typed fields. The keys are + * strings representing collections field names. The values are lists of field masks indicating the key(s) to be used for updates on that field. + * + * e.g. [{"my_struct.foo": {"paths": ["x.y", "x.z"]}}] -> the foo field of my_struct is a list, which has primary keys of x.y (where x is a struct + * and y is a field of x) and x.z (where z is also a field of x). A real world example could be lineage where x is an Upstream struct, + * x.y is the dataset urn, and x.z is the lastmodifiedby value. + */ + primaryKeys: optional array[map[string, map[string, array[string]]]] + + /** + * How to handle updates. The default/implicit behavior is REPLACE_BY_KEY + */ + defaultUpdateBehavior: optional enum UpdateBehavior { + REPLACE_BY_KEY, + REPLACE_BY_ACTOR, + REPLACE_ALL + } +} diff --git a/gradle-plugins/metadata-annotations-schema/src/main/pegasus/com/linkedin/metadata/annotations/GmaAnnotation.pdl b/gradle-plugins/metadata-annotations-schema/src/main/pegasus/com/linkedin/metadata/annotations/GmaAnnotation.pdl index bf4fa9f4f..83f506ef4 100644 --- a/gradle-plugins/metadata-annotations-schema/src/main/pegasus/com/linkedin/metadata/annotations/GmaAnnotation.pdl +++ b/gradle-plugins/metadata-annotations-schema/src/main/pegasus/com/linkedin/metadata/annotations/GmaAnnotation.pdl @@ -27,4 +27,10 @@ record GmaAnnotation { * Information about GMA Search functionality. */ search: optional SearchAnnotation + + /** + * Information about collections + */ + collection: optional CollectionAnnotation + } \ No newline at end of file diff --git a/gradle-plugins/metadata-annotations-test-models/src/main/pegasus/com/linkedin/testing/CollectionAnnotatedAspectBar.pdl b/gradle-plugins/metadata-annotations-test-models/src/main/pegasus/com/linkedin/testing/CollectionAnnotatedAspectBar.pdl new file mode 100644 index 000000000..236ac3128 --- /dev/null +++ b/gradle-plugins/metadata-annotations-test-models/src/main/pegasus/com/linkedin/testing/CollectionAnnotatedAspectBar.pdl @@ -0,0 +1,32 @@ +namespace com.linkedin.testing + +/** + * For unit tests + */ +@gma.collection = { + "isCollection" : true, + "primaryKeys" : [ { + "foo" : { + "paths" : [ "x.y", "x.z" ] + } + }, { + "bar" : { + "paths" : [ "abc", "def" ] + } + } ], + "defaultUpdateBehavior" : "REPLACE_BY_ACTOR" +} +record CollectionAnnotatedAspectBar { + + /** For unit tests */ + stringField: string + + /** For unit tests */ + boolField: boolean + + /** For unit tests */ + longField: long + + /** For unit tests */ + arrayField: array[string] +} \ No newline at end of file