feat(mxe auto-gen):Support aspect modeled as typeref when auto-genera…

…ting MXE (#150)
linkedin · Mar 7, 2022 · 9d15edb · 9d15edb
1 parent 660bea2
commit 9d15edb
Show file tree

Hide file tree

Showing 8 changed files with 182 additions and 19 deletions.
diff --git a/...data-annotations-lib/src/main/java/com/linkedin/metadata/annotations/AlwaysAllowList.java b/...data-annotations-lib/src/main/java/com/linkedin/metadata/annotations/AlwaysAllowList.java
@@ -1,14 +1,14 @@
 package com.linkedin.metadata.annotations;
 
 import com.google.common.annotations.VisibleForTesting;
-import com.linkedin.data.schema.RecordDataSchema;
+import com.linkedin.data.schema.DataSchema;
 import javax.annotation.Nonnull;
 
 
 @VisibleForTesting
 public class AlwaysAllowList implements GmaEntitiesAnnotationAllowList {
   @Override
-  public void check(@Nonnull RecordDataSchema schema, @Nonnull AspectAnnotation aspectAnnotation) {
+  public void check(@Nonnull DataSchema schema, @Nonnull AspectAnnotation aspectAnnotation) {
     // nothing
   }
 }
diff --git a/...adata-annotations-lib/src/main/java/com/linkedin/metadata/annotations/DataSchemaUtil.java b/...adata-annotations-lib/src/main/java/com/linkedin/metadata/annotations/DataSchemaUtil.java
@@ -0,0 +1,39 @@
+package com.linkedin.metadata.annotations;
+
+import com.linkedin.data.schema.DataSchema;
+import com.linkedin.data.schema.RecordDataSchema;
+import com.linkedin.data.schema.TyperefDataSchema;
+import javax.annotation.CheckForNull;
+import javax.annotation.Nonnull;
+
+
+/**
+ * Utility operations on DataSchema.
+ */
+public class DataSchemaUtil {
+
+  private DataSchemaUtil() {
+  }
+
+  @CheckForNull
+  public static String getFullName(@Nonnull DataSchema dataSchema) {
+    if (dataSchema instanceof RecordDataSchema) {
+      return ((RecordDataSchema) dataSchema).getFullName();
+    } else if (dataSchema instanceof TyperefDataSchema) {
+      return ((TyperefDataSchema) dataSchema).getFullName();
+    }
+
+    return null;
+  }
+
+  @CheckForNull
+  public static String getNamespace(@Nonnull DataSchema dataSchema) {
+    if (dataSchema instanceof RecordDataSchema) {
+      return ((RecordDataSchema) dataSchema).getNamespace();
+    } else if (dataSchema instanceof TyperefDataSchema) {
+      return ((TyperefDataSchema) dataSchema).getNamespace();
+    }
+
+    return null;
+  }
+}
diff --git a/...-annotations-lib/src/main/java/com/linkedin/metadata/annotations/GmaAnnotationParser.java b/...-annotations-lib/src/main/java/com/linkedin/metadata/annotations/GmaAnnotationParser.java
@@ -1,7 +1,7 @@
 package com.linkedin.metadata.annotations;
 
 import com.google.common.base.Joiner;
-import com.linkedin.data.schema.RecordDataSchema;
+import com.linkedin.data.schema.DataSchema;
 import com.linkedin.data.schema.validation.RequiredMode;
 import com.linkedin.data.schema.validation.UnrecognizedFieldMode;
 import com.linkedin.data.schema.validation.ValidateDataAgainstSchema;
@@ -37,7 +37,7 @@ public GmaAnnotationParseException(String message) {
    *
    * @throws GmaAnnotationParseException if the provided {@code @gma} annotation does not match the schema.
    */
-  public Optional<GmaAnnotation> parse(@Nonnull RecordDataSchema schema) {
+  public Optional<GmaAnnotation> parse(@Nonnull DataSchema schema) {
     final Map<String, Object> properties = schema.getProperties();
     final Object gmaObj = properties.get(GMA);
 
@@ -55,7 +55,7 @@ public Optional<GmaAnnotation> parse(@Nonnull RecordDataSchema schema) {
     final GmaAnnotation gmaAnnotation = DataTemplateUtil.wrap(gmaObj, GmaAnnotation.class);
 
     if (!result.isValid()) {
-      throw new GmaAnnotationParseException(String.format("Error parsing @gma on %s: \n%s", schema.getFullName(),
+      throw new GmaAnnotationParseException(String.format("Error parsing @gma on %s: \n%s", DataSchemaUtil.getFullName(schema),
           Joiner.on('\n').join(result.getMessages())));
     }
 

diff --git a/...s-lib/src/main/java/com/linkedin/metadata/annotations/GmaEntitiesAnnotationAllowList.java b/...s-lib/src/main/java/com/linkedin/metadata/annotations/GmaEntitiesAnnotationAllowList.java
@@ -1,6 +1,6 @@
 package com.linkedin.metadata.annotations;
 
-import com.linkedin.data.schema.RecordDataSchema;
+import com.linkedin.data.schema.DataSchema;
 import javax.annotation.Nonnull;
 
 
@@ -21,5 +21,5 @@ public AnnotationNotAllowedException(String message) {
     }
   }
 
-  void check(@Nonnull RecordDataSchema schema, @Nonnull AspectAnnotation aspectAnnotation);
+  void check(@Nonnull DataSchema schema, @Nonnull AspectAnnotation aspectAnnotation);
 }
diff --git a/...b/src/main/java/com/linkedin/metadata/annotations/GmaEntitiesAnnotationAllowListImpl.java b/...b/src/main/java/com/linkedin/metadata/annotations/GmaEntitiesAnnotationAllowListImpl.java
@@ -2,7 +2,7 @@
 
 import com.google.common.collect.ImmutableSetMultimap;
 import com.google.common.collect.Multimap;
-import com.linkedin.data.schema.RecordDataSchema;
+import com.linkedin.data.schema.DataSchema;
 import java.util.Collection;
 import javax.annotation.Nonnull;
 
@@ -40,18 +40,18 @@ public GmaEntitiesAnnotationAllowListImpl(Multimap<String, String> allowedModelF
   }
 
   @Override
-  public void check(@Nonnull RecordDataSchema schema, @Nonnull AspectAnnotation aspectAnnotation) {
-    final Collection<String> urns = _allowedModelFqcnsToAllowedUrns.get(schema.getFullName());
+  public void check(@Nonnull DataSchema schema, @Nonnull AspectAnnotation aspectAnnotation) {
+    final Collection<String> urns = _allowedModelFqcnsToAllowedUrns.get(DataSchemaUtil.getFullName(schema));
 
     if (urns.isEmpty()) {
       throw new AnnotationNotAllowedException(
-          String.format("@gma.aspect.entities not allowed on %s", schema.getFullName()));
+          String.format("@gma.aspect.entities not allowed on %s", DataSchemaUtil.getFullName(schema)));
     }
 
     for (AspectEntityAnnotation entity : aspectAnnotation.getEntities()) {
       if (!urns.contains(entity.getUrn())) {
         throw new AnnotationNotAllowedException(
-            String.format("URN %s is not allowed for entity %s", entity.getUrn(), schema.getFullName()));
+            String.format("URN %s is not allowed for entity %s", entity.getUrn(), DataSchemaUtil.getFullName(schema)));
       }
     }
   }

diff --git a/...enerator-lib/src/main/java/com/linkedin/metadata/generator/SchemaAnnotationRetriever.java b/...enerator-lib/src/main/java/com/linkedin/metadata/generator/SchemaAnnotationRetriever.java
@@ -1,8 +1,8 @@
 package com.linkedin.metadata.generator;
 
 import com.linkedin.data.schema.DataSchema;
-import com.linkedin.data.schema.RecordDataSchema;
 import com.linkedin.metadata.annotations.AspectEntityAnnotation;
+import com.linkedin.metadata.annotations.DataSchemaUtil;
 import com.linkedin.metadata.annotations.GmaAnnotationParser;
 import com.linkedin.metadata.annotations.GmaEntitiesAnnotationAllowList;
 import com.linkedin.pegasus.generator.DataSchemaParser;
@@ -41,15 +41,16 @@ public List<EventSpec> generate(@Nonnull String[] sources) throws IOException {
     final DataSchemaParser.ParseResult parseResult = _dataSchemaParser.parseSources(sources);
     final List<EventSpec> eventSpecs = new ArrayList<>();
     for (DataSchema dataSchema : parseResult.getSchemaAndLocations().keySet()) {
-      if (dataSchema.getType() == DataSchema.Type.RECORD) {
-        eventSpecs.addAll(generate((RecordDataSchema) dataSchema));
+      if (dataSchema.getType() == DataSchema.Type.RECORD
+          || dataSchema.getType() == DataSchema.Type.TYPEREF && dataSchema.getDereferencedType() == DataSchema.Type.RECORD) {
+        eventSpecs.addAll(generate(dataSchema));
       }
     }
     return eventSpecs;
   }
 
   @Nonnull
-  private List<EventSpec> generate(@Nonnull RecordDataSchema schema) {
+  private List<EventSpec> generate(@Nonnull DataSchema schema) {
     return _gmaAnnotationParser.parse(schema).map(gma -> {
       final List<EventSpec> eventSpecs = new ArrayList<>();
 
@@ -67,9 +68,9 @@ private List<EventSpec> generate(@Nonnull RecordDataSchema schema) {
 
       for (AspectEntityAnnotation entityAnnotation : entityAnnotations) {
         final EventSpec.EventSpecBuilder builder = EventSpec.builder();
-        builder.namespace(schema.getNamespace());
-        builder.fullValueType(schema.getFullName());
-        builder.valueType(SchemaGeneratorUtil.stripNamespace(schema.getFullName()));
+        builder.namespace(DataSchemaUtil.getNamespace(schema));
+        builder.fullValueType(DataSchemaUtil.getFullName(schema));
+        builder.valueType(SchemaGeneratorUtil.stripNamespace(DataSchemaUtil.getFullName(schema)));
         builder.urn(entityAnnotation.getUrn());
 
         builder.eventType(SchemaGeneratorConstants.MetadataEventType.CHANGE);

diff --git a/...rc/test/groovy/com/linkedin/metadata/gradle/MetadataEventsGeneratorPluginIntegTest.groovy b/...rc/test/groovy/com/linkedin/metadata/gradle/MetadataEventsGeneratorPluginIntegTest.groovy
@@ -186,6 +186,108 @@ class MetadataEventsGeneratorPluginIntegTest extends Specification {
          */
         metadataChangeEvent: MetadataChangeEvent
 
+        /**
+         * The error message or the stacktrace for the failure.
+         */
+        error: string
+      }'''.stripIndent()
+
+    projectFile('build/my-dummy-module/generateMetadataEventsSchema/com/linkedin/mxe/foo/testTyperefAspect/MetadataChangeEvent.pdl').text == '''\
+      namespace com.linkedin.mxe.foo.testTyperefAspect
+
+      import com.linkedin.avro2pegasus.events.KafkaAuditHeader
+      import com.linkedin.metadata.events.ChangeType
+      import com.linkedin.testing.FooUrn
+      import com.linkedin.metadata.test.aspects.TestTyperefAspect
+
+      /**
+       * MetadataChangeEvent for the FooUrn with TestTyperefAspect aspect.
+       */
+      @MetadataChangeEvent
+      record MetadataChangeEvent {
+
+        /**
+         * Kafka audit header. See go/kafkaauditheader for more info.
+         */
+        auditHeader: optional KafkaAuditHeader
+
+        /**
+         * FooUrn as the key for the MetadataChangeEvent.
+         */
+        urn: FooUrn
+
+        /**
+         * Value of the proposed TestTyperefAspect change.
+         */
+        proposedValue: optional TestTyperefAspect
+
+        /**
+         * Change type.
+         */
+        changeType: optional union[null, ChangeType] = null
+      }'''.stripIndent()
+
+    projectFile('build/my-dummy-module/generateMetadataEventsSchema/com/linkedin/mxe/foo/testTyperefAspect/MetadataAuditEvent.pdl').text == '''\
+      namespace com.linkedin.mxe.foo.testTyperefAspect
+
+      import com.linkedin.avro2pegasus.events.KafkaAuditHeader
+      import com.linkedin.metadata.events.ChangeType
+      import com.linkedin.testing.FooUrn
+      import com.linkedin.metadata.test.aspects.TestTyperefAspect
+
+      /**
+       * MetadataAuditEvent for the FooUrn with TestTyperefAspect aspect.
+       */
+      @MetadataAuditEvent
+      record MetadataAuditEvent {
+
+        /**
+         * Kafka audit header for the MetadataAuditEvent.
+         */
+        auditHeader: optional KafkaAuditHeader
+
+        /**
+         * FooUrn as the key for the MetadataAuditEvent.
+         */
+        urn: FooUrn
+
+        /**
+         * Aspect of the TestTyperefAspect before the update.
+         */
+        oldValue: optional TestTyperefAspect
+
+        /**
+         * Aspect of the TestTyperefAspect after the update.
+         */
+        newValue: TestTyperefAspect
+
+        /**
+         * Change type.
+         */
+        changeType: optional union[null, ChangeType] = null
+      }'''.stripIndent()
+
+    projectFile('build/my-dummy-module/generateMetadataEventsSchema/com/linkedin/mxe/foo/testTyperefAspect/FailedMetadataChangeEvent.pdl').text == '''\
+      namespace com.linkedin.mxe.foo.testTyperefAspect
+
+      import com.linkedin.avro2pegasus.events.KafkaAuditHeader
+
+      /**
+       * FailedMetadataChangeEvent for the FooUrn with TestTyperefAspect aspect.
+       */
+      @FailedMetadataChangeEvent
+      record FailedMetadataChangeEvent {
+
+        /**
+         * Kafka event for capturing a failure to process a MetadataChangeEvent.
+         */
+        auditHeader: optional KafkaAuditHeader
+
+        /**
+         * The event that failed to be processed.
+         */
+        metadataChangeEvent: MetadataChangeEvent
+
         /**
          * The error message or the stacktrace for the failure.
          */
@@ -213,6 +315,11 @@ class MetadataEventsGeneratorPluginIntegTest extends Specification {
     testAspectDir.resolve('MetadataChangeEvent.java').toFile().exists()
     testAspectDir.resolve('MetadataAuditEvent.java').toFile().exists()
     testAspectDir.resolve('FailedMetadataChangeEvent.java').toFile().exists()
+
+    Path testTyperefAspectDir = projectPath('my-dummy-module/src/mainGeneratedDataTemplate/java/com/linkedin/mxe/foo/testTyperefAspect/')
+    testTyperefAspectDir.resolve('MetadataChangeEvent.java').toFile().exists()
+    testTyperefAspectDir.resolve('MetadataAuditEvent.java').toFile().exists()
+    testTyperefAspectDir.resolve('FailedMetadataChangeEvent.java').toFile().exists()
   }
 
   def 'data template jar has java class events'() {
@@ -235,6 +342,11 @@ class MetadataEventsGeneratorPluginIntegTest extends Specification {
     jar.getEntry('com/linkedin/mxe/foo/testAspect/MetadataAuditEvent.class') != null
     jar.getEntry('com/linkedin/mxe/foo/testAspect/FailedMetadataChangeEvent.class') != null
     jar.getEntry('com/linkedin/metadata/test/aspects/TestAspect.class') != null
+
+    jar.getEntry('com/linkedin/mxe/foo/testTyperefAspect/MetadataChangeEvent.class') != null
+    jar.getEntry('com/linkedin/mxe/foo/testTyperefAspect/MetadataAuditEvent.class') != null
+    jar.getEntry('com/linkedin/mxe/foo/testTyperefAspect/FailedMetadataChangeEvent.class') != null
+    jar.getEntry('com/linkedin/metadata/test/aspects/TestTyperefAspect.class') != null
   }
 
   def 'data template jar has pdl events'() {
@@ -256,5 +368,10 @@ class MetadataEventsGeneratorPluginIntegTest extends Specification {
     jar.getEntry('pegasus/com/linkedin/mxe/foo/testAspect/MetadataAuditEvent.pdl') != null
     jar.getEntry('pegasus/com/linkedin/mxe/foo/testAspect/FailedMetadataChangeEvent.pdl') != null
     jar.getEntry('pegasus/com/linkedin/metadata/test/aspects/TestAspect.pdl') != null
+
+    jar.getEntry('pegasus/com/linkedin/mxe/foo/testTyperefAspect/MetadataChangeEvent.pdl') != null
+    jar.getEntry('pegasus/com/linkedin/mxe/foo/testTyperefAspect/MetadataAuditEvent.pdl') != null
+    jar.getEntry('pegasus/com/linkedin/mxe/foo/testTyperefAspect/FailedMetadataChangeEvent.pdl') != null
+    jar.getEntry('pegasus/com/linkedin/metadata/test/aspects/TestTyperefAspect.pdl') != null
   }
 }
diff --git a/...my-dummy-module/src/main/pegasus/com/linkedin/metadata/test/aspects/TestTyperefAspect.pdl b/...my-dummy-module/src/main/pegasus/com/linkedin/metadata/test/aspects/TestTyperefAspect.pdl
@@ -0,0 +1,6 @@
+namespace com.linkedin.metadata.test.aspects
+
+import com.linkedin.metadata.test.models.TestModel
+
+@gma.aspect.entity.urn = "com.linkedin.testing.FooUrn"
+typeref TestTyperefAspect = TestModel