feat(models): support charts in edges fields of other charts (#14339)

metadata-ingestion/src/datahub/ingestion/autogenerated/lineage.json

@@ -192,7 +192,8 @@
             "relationship": {
               "name": "Consumes",
               "entityTypes": [
-                "dataset"
+                "dataset",
+                "chart"
               ],
               "isLineage": true
             }
@@ -397,5 +398,5 @@
     }
   },
   "generated_by": "metadata-ingestion/scripts/modeldocgen.py",
-  "generated_at": "2025-07-01T10:49:03.713749+00:00"
+  "generated_at": "2025-08-05T19:29:49.306404+00:00"
 }

metadata-models/src/main/pegasus/com/linkedin/chart/ChartInfo.pdl

@@ -66,7 +66,7 @@ record ChartInfo includes CustomProperties, ExternalReference {
   @Relationship = {
     "/*/destinationUrn": {
       "name": "Consumes",
-      "entityTypes": [ "dataset" ],
+      "entityTypes": [ "dataset", "chart" ],
       "isLineage": true,
       "createdOn": "inputEdges/*/created/time"
       "createdActor": "inputEdges/*/created/actor"

metadata-service/services/src/main/java/com/linkedin/metadata/service/LineageService.java

@@ -76,6 +76,24 @@ public class LineageService {
     }
   }
 
+  /**
+   * Validates that a given list of urns are all either datasets or charts and that they exist.
+   * Otherwise, throw an error.
+   */
+  public void validateChartUpstreamUrns(
+      @Nonnull OperationContext opContext, @Nonnull final List<Urn> urns) throws Exception {
+    for (final Urn urn : urns) {
+      if (!urn.getEntityType().equals(Constants.DATASET_ENTITY_NAME)
+          && !urn.getEntityType().equals(Constants.CHART_ENTITY_NAME)) {
+        throw new IllegalArgumentException(
+            String.format(
+                "Tried to add an upstream to a chart that isn't a chart or dataset. Upstream urn: %s",
+                urn));
+      }
+      validateUrnExists(opContext, urn);
+    }
+  }
+
   /** Validates that a given urn exists using the entityService */
   public void validateUrnExists(@Nonnull OperationContext opContext, @Nonnull final Urn urn)
       throws Exception {
@@ -174,8 +192,8 @@ public class LineageService {
       @Nonnull final List<Urn> upstreamUrnsToRemove,
       @Nonnull final Urn actor)
       throws Exception {
-    // ensure all upstream urns are dataset urns and they exist
+    // ensure all upstream urns are either dataset or chart urns and they exist
-    validateDatasetUrns(opContext, upstreamUrnsToAdd);
+    validateChartUpstreamUrns(opContext, upstreamUrnsToAdd);
     // TODO: add permissions check here for entity type - or have one overall permissions check
     // above
 

metadata-service/services/src/test/java/com/linkedin/metadata/graph/LineageGraphFiltersTest.java

@@ -158,13 +158,19 @@ public class LineageGraphFiltersTest {
 
     List<Pair<String, LineageRegistry.EdgeInfo>> streamResult = filters.streamEdgeInfo().toList();
 
-    assertEquals(streamResult.size(), 1);
+    assertEquals(streamResult.size(), 2);
     assertTrue(
         streamResult.contains(
             Pair.of(
                 "chart",
                 new LineageRegistry.EdgeInfo(
                     "Consumes", RelationshipDirection.OUTGOING, "dataset"))));
+    assertTrue(
+        streamResult.contains(
+            Pair.of(
+                "chart",
+                new LineageRegistry.EdgeInfo(
+                    "Consumes", RelationshipDirection.OUTGOING, "chart"))));
 
     assertTrue(filters.containsEdgeInfo("chart", streamResult.get(0).getValue()));
     assertFalse(

metadata-service/services/src/test/java/com/linkedin/metadata/service/LineageServiceTest.java

@@ -186,11 +186,12 @@ public class LineageServiceTest {
                 opContext, datasetUrn1, upstreamUrnsToAdd, upstreamUrnsToRemove, actorUrn));
   }
 
-  // Adds upstream for chart1 to dataset3 and removes edge to dataset1 while keeping edge to
+  // Adds upstream for chart1 to dataset3 and chart2 and removes edge to dataset1 while keeping edge
-  // dataset2
+  // to dataset2
   @Test
   public void testUpdateChartLineage() throws Exception {
     Mockito.when(_mockClient.exists(any(OperationContext.class), eq(chartUrn1))).thenReturn(true);
+    Mockito.when(_mockClient.exists(any(OperationContext.class), eq(chartUrn2))).thenReturn(true);
     Mockito.when(_mockClient.exists(any(OperationContext.class), eq(datasetUrn1))).thenReturn(true);
     Mockito.when(_mockClient.exists(any(OperationContext.class), eq(datasetUrn2))).thenReturn(true);
     Mockito.when(_mockClient.exists(any(OperationContext.class), eq(datasetUrn3))).thenReturn(true);
@@ -215,17 +216,17 @@ public class LineageServiceTest {
                             Constants.CHART_INFO_ASPECT_NAME,
                             new EnvelopedAspect().setValue(new Aspect(chartInfo.data()))))));
 
-    final List<Urn> upstreamUrnsToAdd = Collections.singletonList(datasetUrn3);
+    final List<Urn> upstreamUrnsToAdd = Arrays.asList(datasetUrn3, chartUrn2);
     final List<Urn> upstreamUrnsToRemove = Collections.singletonList(datasetUrn2);
     _lineageService.updateChartLineage(
         opContext, chartUrn1, upstreamUrnsToAdd, upstreamUrnsToRemove, actorUrn);
 
-    // chartInfo with dataset1 in inputs and dataset3 in inputEdges
+    // chartInfo with dataset1 in inputs, dataset3 and chart2 in inputEdges
     ChartInfo updatedChartInfo =
         createChartInfo(
             chartUrn1,
             Collections.singletonList(datasetUrn1),
-            Collections.singletonList(datasetUrn3));
+            Arrays.asList(datasetUrn3, chartUrn2));
 
     final MetadataChangeProposal proposal = new MetadataChangeProposal();
     proposal.setEntityUrn(chartUrn1);
@@ -253,10 +254,10 @@ public class LineageServiceTest {
 
   @Test
   public void testFailUpdateChartWithInvalidEdge() throws Exception {
-    Mockito.when(_mockClient.exists(opContext, chartUrn2)).thenReturn(true);
+    Mockito.when(_mockClient.exists(opContext, datajobUrn1)).thenReturn(true);
 
-    // charts can't have charts upstream of them
+    // charts can't have datajobs upstream of them
-    final List<Urn> upstreamUrnsToAdd = Collections.singletonList(chartUrn2);
+    final List<Urn> upstreamUrnsToAdd = Collections.singletonList(datajobUrn1);
     final List<Urn> upstreamUrnsToRemove = Collections.emptyList();
     assertThrows(
         RuntimeException.class,