feat: Add ingestion mode annotation (#359)

* feat: Add ingestion mode annotation

* remove unused import

* add doc

* fix

* add filter

* fix

* fix comment

---------

Co-authored-by: Jesse Jia <[email protected]>

dao-impl/ebean-dao/build.gradle

@@ -7,6 +7,7 @@ configurations {
 }
 
 dependencies {
+  compile project(':gradle-plugins:metadata-annotations-lib')
   compile project(':core-models-utils')
   compile project(':dao-api')
   compile externalDependency.ebean

dao-impl/ebean-dao/src/main/java/com/linkedin/metadata/dao/utils/EBeanDAOUtils.java

@@ -1,6 +1,11 @@
 package com.linkedin.metadata.dao.utils;
 
+import com.linkedin.data.schema.RecordDataSchema;
+import com.linkedin.data.template.DataTemplateUtil;
 import com.linkedin.data.template.RecordTemplate;
+import com.linkedin.metadata.annotations.AspectIngestionAnnotationArray;
+import com.linkedin.metadata.annotations.GmaAnnotation;
+import com.linkedin.metadata.annotations.GmaAnnotationParser;
 import com.linkedin.metadata.aspect.AuditedAspect;
 import com.linkedin.metadata.aspect.SoftDeletedAspect;
 import com.linkedin.metadata.dao.EbeanMetadataAspect;
@@ -22,6 +27,7 @@
 import java.util.ArrayList;
 import java.util.List;
 import java.util.Map;
+import java.util.Optional;
 import java.util.stream.Collectors;
 import javax.annotation.Nonnull;
 import javax.annotation.Nullable;
@@ -48,6 +54,26 @@ private EBeanDAOUtils() {
     // Utils class
   }
 
+  /**
+   * Parse the ingestion mode annotation given an aspect class.
+   */
+  @Nonnull
+  public static AspectIngestionAnnotationArray parseIngestionModeFromAnnotation(@Nonnull final String aspectCanonicalName) {
+    try {
+      final RecordDataSchema schema = (RecordDataSchema) DataTemplateUtil.getSchema(ClassUtils.loadClass(aspectCanonicalName));
+      final Optional<GmaAnnotation> gmaAnnotation = new GmaAnnotationParser().parse(schema);
+
+      // Return empty array if user did not specify any ingestion annotation on the aspect.
+      if (!gmaAnnotation.isPresent() || !gmaAnnotation.get().hasAspect() || !gmaAnnotation.get().getAspect().hasIngestion()) {
+        return new AspectIngestionAnnotationArray();
+      }
+
+      return gmaAnnotation.get().getAspect().getIngestion();
+    } catch (Exception e) {
+      throw new RuntimeException(String.format("Failed to parse the annotations for aspect %s", aspectCanonicalName), e);
+    }
+  }
+
   /**
    * Given urn string and Urn class, return Urn instance.
    * @param urn urn string

gradle-plugins/metadata-annotations-lib/src/main/java/com/linkedin/metadata/annotations/GmaAnnotationParser.java

@@ -21,6 +21,7 @@
 public class GmaAnnotationParser {
   private static final String GMA = "gma";
   private static final String SEARCH = "search";
+  private static final String INDEX = "index";
 
   private final GmaEntitiesAnnotationAllowList _gmaEntitiesAnnotationAllowList;
 
@@ -117,7 +118,7 @@ private Optional<GmaAnnotation> parseTopLevelAnnotations(@Nonnull DataSchema sch
           continue;
         }
 
-        final Object indexObj = ((DataMap) searchObj).get("index");
+        final Object indexObj = ((DataMap) searchObj).get(INDEX);
         if (indexObj == null) {
           continue;
         }

gradle-plugins/metadata-annotations-lib/src/test/java/com/linkedin/metadata/annotations/GmaAnnotationParserTest.java

@@ -19,12 +19,20 @@
 public class GmaAnnotationParserTest {
   @Test
   public void parseBar() {
+    AspectIngestionAnnotationArray ingestionAnnotations = new AspectIngestionAnnotationArray(
+        new AspectIngestionAnnotation()
+            .setUrn("com.linkedin.testing.BarUrn")
+            .setMode(Mode.FORCE_UPDATE)
+            .setFilter(new UrnFilterArray(new UrnFilter().setPath("/platform").setValue("hdfs"))));
+
     // has both @gma.aspect.entity.urn and @gma.aspect.column.name annotations
     final Optional<GmaAnnotation> gma =
         new GmaAnnotationParser().parse((RecordDataSchema) DataTemplateUtil.getSchema(AnnotatedAspectBar.class));
-    assertThat(gma).contains(new GmaAnnotation().setAspect(
-        new AspectAnnotation().setEntity(new AspectEntityAnnotation().setUrn("com.linkedin.testing.BarUrn"))
-            .setColumn(new ColumnNameAnnotation().setName("barurn"))));
+    assertThat(gma).contains(new GmaAnnotation()
+        .setAspect(new AspectAnnotation()
+            .setEntity(new AspectEntityAnnotation().setUrn("com.linkedin.testing.BarUrn"))
+            .setColumn(new ColumnNameAnnotation().setName("barurn"))
+            .setIngestion(ingestionAnnotations)));
   }
 
   @Test

gradle-plugins/metadata-annotations-schema/src/main/pegasus/com/linkedin/metadata/annotations/AspectAnnotation.pdl

@@ -14,6 +14,14 @@ record AspectAnnotation {
    */
   column: optional ColumnNameAnnotation
 
+  /**
+   * Information about how the aspect should be ingested.
+   * Example use cases:
+   * 1. Skip equality check.
+   * 2. Skip sementic versioning check.
+   */
+  ingestion: optional array[AspectIngestionAnnotation]
+
   /**
    * Information on multiple entities this aspect is associated with. Deprecated; it is an error to use outside of a limited set of allowed models. Used to help migrate existing "common aspects" to v5 only.
    */

gradle-plugins/metadata-annotations-schema/src/main/pegasus/com/linkedin/metadata/annotations/AspectIngestionAnnotation.pdl

@@ -0,0 +1,41 @@
+namespace com.linkedin.metadata.annotations
+
+record AspectIngestionAnnotation {
+
+  /**
+   * Ingestion Mode applied this aspect. This can be overridden by IngestionMode set in (union) MCE.
+   */
+  mode: optional enum Mode {
+    /**
+     * Skip any check in data access layer. Update the aspect in database irrespectively.
+     */
+    FORCE_UPDATE
+
+    /**
+     * Honor all existing checks in data access layer.
+     */
+    DEFAULT
+  }
+
+  /**
+   * The FQCN of the URN that identifies this entity. The ingestion mode will be applied to the aspect associated with this entity.
+   */
+  urn: optional string
+
+  /**
+   * Filter on the URN so that this ingestion mode is only applicable to a subset of entities.
+   */
+  filter: optional array[record UrnFilter {
+
+    /**
+     * Path extracted by UrnPathExtractor
+     * https://github.com/linkedin/datahub-gma/blob/master/dao-impl/ebean-dao/src/main/java/com/linkedin/metadata/dao/scsi/UrnPathExtractor.java
+     */
+    path: optional string,
+
+    /**
+     * The target value lead to by by the path.
+     */
+    value: optional string
+  }]
+}

gradle-plugins/metadata-annotations-test-models/src/main/pegasus/com/linkedin/testing/AnnotatedAspectBar.pdl

@@ -3,6 +3,11 @@ namespace com.linkedin.testing
 /**
  * For unit tests
  */
+@gma.aspect.ingestion = [
+  {"mode": "FORCE_UPDATE", "urn": "com.linkedin.testing.BarUrn", "filter": [
+    {"path": "/platform", "value": "hdfs"}
+  ]}
+]
 @gma.aspect.column.name = "barurn"
 @gma.aspect.entity = {
   "urn": "com.linkedin.testing.BarUrn"