diff --git a/metadata-ingestion/src/datahub/ingestion/autogenerated/lineage.json b/metadata-ingestion/src/datahub/ingestion/autogenerated/lineage.json index 40ab8e3ba0..ed176636cb 100644 --- a/metadata-ingestion/src/datahub/ingestion/autogenerated/lineage.json +++ b/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" } \ No newline at end of file diff --git a/metadata-models/src/main/pegasus/com/linkedin/chart/ChartInfo.pdl b/metadata-models/src/main/pegasus/com/linkedin/chart/ChartInfo.pdl index 2d76417eb8..650cdaee7d 100644 --- a/metadata-models/src/main/pegasus/com/linkedin/chart/ChartInfo.pdl +++ b/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" diff --git a/metadata-service/services/src/main/java/com/linkedin/metadata/service/LineageService.java b/metadata-service/services/src/main/java/com/linkedin/metadata/service/LineageService.java index 08c17945c7..9792812f8e 100644 --- a/metadata-service/services/src/main/java/com/linkedin/metadata/service/LineageService.java +++ b/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 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 upstreamUrnsToRemove, @Nonnull final Urn actor) throws Exception { - // ensure all upstream urns are dataset urns and they exist - validateDatasetUrns(opContext, upstreamUrnsToAdd); + // ensure all upstream urns are either dataset or chart urns and they exist + validateChartUpstreamUrns(opContext, upstreamUrnsToAdd); // TODO: add permissions check here for entity type - or have one overall permissions check // above diff --git a/metadata-service/services/src/test/java/com/linkedin/metadata/graph/LineageGraphFiltersTest.java b/metadata-service/services/src/test/java/com/linkedin/metadata/graph/LineageGraphFiltersTest.java index a3d2f79bbc..a221a57880 100644 --- a/metadata-service/services/src/test/java/com/linkedin/metadata/graph/LineageGraphFiltersTest.java +++ b/metadata-service/services/src/test/java/com/linkedin/metadata/graph/LineageGraphFiltersTest.java @@ -158,13 +158,19 @@ public class LineageGraphFiltersTest { List> 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( diff --git a/metadata-service/services/src/test/java/com/linkedin/metadata/service/LineageServiceTest.java b/metadata-service/services/src/test/java/com/linkedin/metadata/service/LineageServiceTest.java index 8730538797..5e7bb1beba 100644 --- a/metadata-service/services/src/test/java/com/linkedin/metadata/service/LineageServiceTest.java +++ b/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 - // dataset2 + // Adds upstream for chart1 to dataset3 and chart2 and removes edge to dataset1 while keeping edge + // 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 upstreamUrnsToAdd = Collections.singletonList(datasetUrn3); + final List upstreamUrnsToAdd = Arrays.asList(datasetUrn3, chartUrn2); final List 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 - final List upstreamUrnsToAdd = Collections.singletonList(chartUrn2); + // charts can't have datajobs upstream of them + final List upstreamUrnsToAdd = Collections.singletonList(datajobUrn1); final List upstreamUrnsToRemove = Collections.emptyList(); assertThrows( RuntimeException.class,