add support for column lineage

atlanhq · May 8, 2024 · 9cb3abf · 9cb3abf
1 parent 890677a
commit 9cb3abf
Showing 1 changed file with 8 additions and 1 deletion.
diff --git a/repository/src/main/java/org/apache/atlas/discovery/EntityLineageService.java b/repository/src/main/java/org/apache/atlas/discovery/EntityLineageService.java
@@ -1470,7 +1470,14 @@ private void processEdge(final AtlasEdge edge, final Map<String, AtlasEntityHead
         AtlasPerfMetrics.MetricRecorder metricRecorder = RequestContext.get().startMetricRecord("processEdge");
         AtlasVertex inVertex     = edge.getInVertex();
         AtlasVertex outVertex    = edge.getOutVertex();
-        boolean     inIsProcess  = Objects.equals(AtlasGraphUtilsV2.getTypeName(inVertex), PROCESS_SUPER_TYPE);
+
+        String inTypeName = AtlasGraphUtilsV2.getTypeName(inVertex);
+        AtlasEntityType inEntityType = atlasTypeRegistry.getEntityTypeByName(inTypeName);
+        if (inEntityType == null) {
+            throw new AtlasBaseException(AtlasErrorCode.TYPE_NAME_NOT_FOUND, inTypeName);
+        }
+        boolean inIsProcess = inEntityType.getTypeAndAllSuperTypes().contains(PROCESS_SUPER_TYPE);
+
         String      inGuid       = AtlasGraphUtilsV2.getIdFromVertex(inVertex);
         String      outGuid      = AtlasGraphUtilsV2.getIdFromVertex(outVertex);
         String      relationGuid = AtlasGraphUtilsV2.getEncodedProperty(edge, RELATIONSHIP_GUID_PROPERTY_KEY, String.class);