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

This commit is contained in:
Michael Maltese 2025-08-27 12:37:08 -04:00 committed by GitHub
parent d1c36b27f2
commit c2f9fc812c
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
5 changed files with 40 additions and 14 deletions

View File

@ -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"
}

View File

@ -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"

View File

@ -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
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

View File

@ -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(

View File

@ -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<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
final List<Urn> upstreamUrnsToAdd = Collections.singletonList(chartUrn2);
// charts can't have datajobs upstream of them
final List<Urn> upstreamUrnsToAdd = Collections.singletonList(datajobUrn1);
final List<Urn> upstreamUrnsToRemove = Collections.emptyList();
assertThrows(
RuntimeException.class,