fix(siblingsHook): add logic to account for non dbt upstreams (#9154)

Co-authored-by: Ethan Cartwright <ethan.cartwright@acryl.io>
datahub-project · Dec 4, 2023 · 7857944 · 7857944
1 parent f7f25bd
commit 7857944
Show file tree

Hide file tree

Showing 2 changed files with 23 additions and 2 deletions.
diff --git a/...sumer/src/main/java/com/linkedin/metadata/kafka/hook/siblings/SiblingAssociationHook.java b/...sumer/src/main/java/com/linkedin/metadata/kafka/hook/siblings/SiblingAssociationHook.java
@@ -205,9 +205,8 @@ private void handleSourceDatasetEvent(MetadataChangeLog event, DatasetUrn source
         // We're assuming a data asset (eg. snowflake table) will only ever be downstream of 1 dbt model
         if (dbtUpstreams.size() == 1) {
           setSiblingsAndSoftDeleteSibling(dbtUpstreams.get(0).getDataset(), sourceUrn);
-        } else {
+        } else if (dbtUpstreams.size() > 1) {
           log.error("{} has an unexpected number of dbt upstreams: {}. Not adding any as siblings.", sourceUrn.toString(), dbtUpstreams.size());
-
         }
       }
     }

diff --git a/...r/src/test/java/com/linkedin/metadata/kafka/hook/siblings/SiblingAssociationHookTest.java b/...r/src/test/java/com/linkedin/metadata/kafka/hook/siblings/SiblingAssociationHookTest.java
@@ -305,6 +305,28 @@ public void testInvokeWhenSourceUrnHasTwoUpstreamsOneDbt() throws Exception {
     Mockito.verify(_mockEntityClient, Mockito.times(2)).ingestProposal(Mockito.any(), eq(true));
   }
 
+  @Test
+  public void testInvokeWhenSourceUrnHasTwoUpstreamsNoDbt() throws Exception {
+
+    MetadataChangeLog event = createEvent(DATASET_ENTITY_NAME, UPSTREAM_LINEAGE_ASPECT_NAME, ChangeType.UPSERT);
+    final UpstreamLineage upstreamLineage = new UpstreamLineage();
+    final UpstreamArray upstreamArray = new UpstreamArray();
+    Upstream snowflakeUpstream1 = 
+            createUpstream("urn:li:dataset:(urn:li:dataPlatform:snowflake,my-proj.jaffle_shop1.customers,PROD)", DatasetLineageType.TRANSFORMED);
+    Upstream snowflakeUpstream2 =
+            createUpstream("urn:li:dataset:(urn:li:dataPlatform:snowflake,my-proj.jaffle_shop2.customers,PROD)", DatasetLineageType.TRANSFORMED);
+    upstreamArray.add(snowflakeUpstream1);
+    upstreamArray.add(snowflakeUpstream2);
+    upstreamLineage.setUpstreams(upstreamArray);
+
+    event.setAspect(GenericRecordUtils.serializeAspect(upstreamLineage));
+    event.setEntityUrn(Urn.createFromString("urn:li:dataset:(urn:li:dataPlatform:bigquery,my-proj.jaffle_shop.customers,PROD)"));
+    _siblingAssociationHook.invoke(event);
+
+
+    Mockito.verify(_mockEntityClient, Mockito.times(0)).ingestProposal(Mockito.any(), eq(true));
+  }
+
   private MetadataChangeLog createEvent(String entityType, String aspectName, ChangeType changeType) {
     MetadataChangeLog event = new MetadataChangeLog();
     event.setEntityType(entityType);