Fixed bug which stopped all current DatasetAccountableOwnership inges…

…tion (#424) * Fixed the bug * Added unit test * Updating * Fixed * Updated unit test --------- Co-authored-by: Rakhi Agrawal <[email protected]>
linkedin · Sep 12, 2024 · 2dd1a36 · 2dd1a36
1 parent 7c07e1e
commit 2dd1a36
Show file tree

Hide file tree

Showing 3 changed files with 41 additions and 15 deletions.
diff --git a/.github/workflows/build-and-test.yml b/.github/workflows/build-and-test.yml
@@ -32,7 +32,7 @@ jobs:
 
       - name: Upload artifact
         if: always()
-        uses: actions/upload-artifact@v2
+        uses: actions/upload-artifact@v3
         with:
           name: test-reports
           path: /home/runner/work/datahub-gma/datahub-gma/dao-impl/ebean-dao/build/reports/tests/test/**

diff --git a/restli-resources/src/main/java/com/linkedin/metadata/restli/BaseAspectRoutingResource.java b/restli-resources/src/main/java/com/linkedin/metadata/restli/BaseAspectRoutingResource.java
@@ -72,7 +72,6 @@ public abstract class BaseAspectRoutingResource<
   private final Class<INTERNAL_SNAPSHOT> _internalSnapshotClass;
   private final Class<INTERNAL_ASPECT_UNION> _internalAspectUnionClass;
   private final Class<ASSET> _assetClass;
-  private RestliPreUpdateAspectRegistry _restliPreUpdateAspectRegistry = null;
   private static final List<String> SKIP_INGESTION_FOR_ASPECTS = Collections.singletonList("com.linkedin.dataset.DatasetAccountableOwnership");
 
   public BaseAspectRoutingResource(@Nullable Class<SNAPSHOT> snapshotClass,
@@ -316,13 +315,13 @@ protected Task<Void> ingestInternal(@Nonnull SNAPSHOT snapshot,
               RestliPreUpdateAspectRegistry registry = getLocalDAO().getRestliPreUpdateAspectRegistry();
               if (registry != null && registry.isRegistered(aspect.getClass())) {
                 aspect = preUpdateRouting(urn, aspect, registry);
-              }
-              // Get the fqcn of the aspect class
-              String aspectFQCN = aspect.getClass().getCanonicalName();
-              //TODO: META-21112: Remove this check after adding annotations at model level; to handle SKIP/PROCEED for preUpdateRouting
-              if (SKIP_INGESTION_FOR_ASPECTS.contains(aspectFQCN)) {
-                return;
-              }
+                // Get the fqcn of the aspect class
+                String aspectFQCN = aspect.getClass().getCanonicalName();
+                //TODO: META-21112: Remove this check after adding annotations at model level; to handle SKIP/PROCEED for preUpdateRouting
+                if (SKIP_INGESTION_FOR_ASPECTS.contains(aspectFQCN)) {
+                  return;
+                }
+              } 
               if (trackingContext != null) {
                 getAspectRoutingGmsClientManager().getRoutingGmsClient(aspect.getClass()).ingestWithTracking(urn, aspect, trackingContext, ingestionParams);
               } else {
@@ -361,12 +360,12 @@ protected Task<Void> ingestInternalAsset(@Nonnull ASSET asset,
               RestliPreUpdateAspectRegistry registry = getLocalDAO().getRestliPreUpdateAspectRegistry();
               if (registry != null && registry.isRegistered(aspect.getClass())) {
                 aspect = preUpdateRouting(urn, aspect, registry);
-              }
-              // Get the fqcn of the aspect class
-              String aspectFQCN = aspect.getClass().getCanonicalName();
-              //TODO: META-21112: Remove this check after adding annotations at model level; to handle SKIP/PROCEED for preUpdateRouting
-              if (SKIP_INGESTION_FOR_ASPECTS.contains(aspectFQCN)) {
-                return;
+                // Get the fqcn of the aspect class
+                String aspectFQCN = aspect.getClass().getCanonicalName();
+                //TODO: META-21112: Remove this check after adding annotations at model level; to handle SKIP/PROCEED for preUpdateRouting
+                if (SKIP_INGESTION_FOR_ASPECTS.contains(aspectFQCN)) {
+                  return;
+                }
               }
               if (trackingContext != null) {
                 getAspectRoutingGmsClientManager().getRoutingGmsClient(aspect.getClass())

diff --git a/...i-resources/src/test/java/com/linkedin/metadata/restli/BaseAspectRoutingResourceTest.java b/...i-resources/src/test/java/com/linkedin/metadata/restli/BaseAspectRoutingResourceTest.java
@@ -612,4 +612,31 @@ public void testPreUpdateRoutingWithSkipIngestion() throws NoSuchFieldException,
     verifyNoMoreInteractions(_mockLocalDAO);
   }
 
+  //Testing the case when aspect has no pre lambda but skipIngestion contains the aspect, so it should not skip ingestion
+  @Test
+  public void testPreUpdateRoutingWithSkipIngestionNoPreLambda() throws NoSuchFieldException, IllegalAccessException {
+
+    Field skipIngestionField = BaseAspectRoutingResource.class.getDeclaredField("SKIP_INGESTION_FOR_ASPECTS");
+    skipIngestionField.setAccessible(true);
+    Field modifiersField = Field.class.getDeclaredField("modifiers");
+    modifiersField.setAccessible(true);
+    modifiersField.setInt(skipIngestionField, skipIngestionField.getModifiers() & ~Modifier.FINAL);
+    List<String> newSkipIngestionList = Arrays.asList("com.linkedin.testing.AspectFoo");
+    skipIngestionField.set(null, newSkipIngestionList);
+
+    FooUrn urn = makeFooUrn(1);
+    AspectFoo foo = new AspectFoo().setValue("foo");
+    List<EntityAspectUnion> aspects = Arrays.asList(ModelUtils.newAspectUnion(EntityAspectUnion.class, foo));
+    EntitySnapshot snapshot = ModelUtils.newSnapshot(EntitySnapshot.class, urn, aspects);
+
+    runAndWait(_resource.ingest(snapshot));
+    // Should not skip ingestion
+    verify(_mockAspectFooGmsClient, times(1)).ingest(eq(urn), eq(foo));
+    // Should check for pre lambda
+    verify(_mockLocalDAO, times(1)).getRestliPreUpdateAspectRegistry();
+    // Should not add to localDAO
+    verify(_mockLocalDAO, times(0)).add(any(), any(), any(), any(), any());
+    verifyNoMoreInteractions(_mockLocalDAO);
+  }
+
 }