fix(lineage): fix lighting cache dataJob platform (datahub-project#10233

)
sleeperdeep · Jun 25, 2024 · c1e9013 · c1e9013
1 parent e5c33ab
commit c1e9013
Show file tree

Hide file tree

Showing 2 changed files with 35 additions and 2 deletions.
diff --git a/metadata-io/src/main/java/com/linkedin/metadata/search/LineageSearchService.java b/metadata-io/src/main/java/com/linkedin/metadata/search/LineageSearchService.java
@@ -10,6 +10,7 @@
 import com.google.common.collect.Lists;
 import com.linkedin.common.UrnArrayArray;
 import com.linkedin.common.urn.Urn;
+import com.linkedin.common.urn.UrnUtils;
 import com.linkedin.data.template.LongMap;
 import com.linkedin.data.template.StringArray;
 import com.linkedin.metadata.Constants;
@@ -469,10 +470,15 @@ private AggregationMetadata constructAggMetadata(String displayName, String name
         .setFilterValues(new FilterValueArray());
   }
 
-  private String getPlatform(String entityType, Urn entityUrn) {
+  @VisibleForTesting
+  String getPlatform(String entityType, Urn entityUrn) {
     String platform = null;
     if (PLATFORM_ENTITY_TYPES.contains(entityType)) {
-      platform = entityUrn.getEntityKey().get(0);
+      if (DATA_JOB_ENTITY_NAME.equals(entityType)) {
+        platform = UrnUtils.getUrn(entityUrn.getEntityKey().get(0)).getEntityKey().get(0);
+      } else {
+        platform = entityUrn.getEntityKey().get(0);
+      }
     }
     if ((platform != null) && (!platform.startsWith("urn:li:dataPlatform"))) {
       platform = "urn:li:dataPlatform:" + platform;

diff --git a/metadata-io/src/test/java/com/linkedin/metadata/search/LineageServiceTestBase.java b/metadata-io/src/test/java/com/linkedin/metadata/search/LineageServiceTestBase.java
@@ -1338,4 +1338,31 @@ public void testCanDoLightning() throws Exception {
     filter = new Filter().setOr(conCritArr);
     Assert.assertTrue(lineageSearchService.canDoLightning(lineageRelationships, "*", filter, null));
   }
+
+  @Test
+  public void testPlatform() {
+    assertEquals(
+        lineageSearchService.getPlatform(
+            "dataset",
+            UrnUtils.getUrn(
+                "urn:li:dataset:(urn:li:dataPlatform:custom,file:///custom/path,PROD)")),
+        "urn:li:dataPlatform:custom");
+    assertEquals(
+        lineageSearchService.getPlatform(
+            "chart", UrnUtils.getUrn("urn:li:chart:(looker,foobar.1234)")),
+        "urn:li:dataPlatform:looker");
+    assertEquals(
+        lineageSearchService.getPlatform(
+            "dashboard", UrnUtils.getUrn("urn:li:dashboard:(looker,dashboards.1234)")),
+        "urn:li:dataPlatform:looker");
+    assertEquals(
+        lineageSearchService.getPlatform(
+            "dataFlow", UrnUtils.getUrn("urn:li:dataFlow:(airflow,foobar,PROD)")),
+        "urn:li:dataPlatform:airflow");
+    assertEquals(
+        lineageSearchService.getPlatform(
+            "dataJob",
+            UrnUtils.getUrn("urn:li:dataJob:(urn:li:dataFlow:(airflow,foobar,PROD),End)")),
+        "urn:li:dataPlatform:airflow");
+  }
 }