From 17de393b9ed02350d079c3eb01279d405bb85e19 Mon Sep 17 00:00:00 2001 From: Chris Collins Date: Mon, 3 Mar 2025 18:36:47 -0500 Subject: [PATCH] fix(lineage) Support views and sorting in impact analysis (#12769) --- .../datahub/graphql/GmsGraphQLEngine.java | 3 +- .../search/SearchAcrossLineageResolver.java | 29 ++- .../src/main/resources/search.graphql | 10 + .../ScrollAcrossLineageResolverTest.java | 3 +- .../SearchAcrossLineageResolverTest.java | 9 +- .../utils/elasticsearch/FilterUtils.java | 130 +++++++++++ .../metadata/utils/FilterUtilsTest.java | 213 ++++++++++++++++++ 7 files changed, 389 insertions(+), 8 deletions(-) create mode 100644 metadata-utils/src/main/java/com/linkedin/metadata/utils/elasticsearch/FilterUtils.java create mode 100644 metadata-utils/src/test/java/com/linkedin/metadata/utils/FilterUtilsTest.java diff --git a/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/GmsGraphQLEngine.java b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/GmsGraphQLEngine.java index 3927ca900e62a4..dd68fdf9e6bc06 100644 --- a/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/GmsGraphQLEngine.java +++ b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/GmsGraphQLEngine.java @@ -1026,7 +1026,8 @@ private void configureQueryResolvers(final RuntimeWiring.Builder builder) { new ScrollAcrossEntitiesResolver(this.entityClient, this.viewService)) .dataFetcher( "searchAcrossLineage", - new SearchAcrossLineageResolver(this.entityClient, this.entityRegistry)) + new SearchAcrossLineageResolver( + this.entityClient, this.entityRegistry, this.viewService)) .dataFetcher( "scrollAcrossLineage", new ScrollAcrossLineageResolver(this.entityClient)) .dataFetcher( diff --git a/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/search/SearchAcrossLineageResolver.java b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/search/SearchAcrossLineageResolver.java index c90be924ac69f5..a62fb94dbda96d 100644 --- a/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/search/SearchAcrossLineageResolver.java +++ b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/search/SearchAcrossLineageResolver.java @@ -6,6 +6,7 @@ import com.google.common.collect.ImmutableSet; import com.linkedin.common.urn.Urn; +import com.linkedin.common.urn.UrnUtils; import com.linkedin.datahub.graphql.QueryContext; import com.linkedin.datahub.graphql.concurrency.GraphQLConcurrencyUtils; import com.linkedin.datahub.graphql.generated.AndFilterInput; @@ -25,8 +26,12 @@ import com.linkedin.metadata.query.LineageFlags; import com.linkedin.metadata.query.SearchFlags; import com.linkedin.metadata.query.filter.Filter; +import com.linkedin.metadata.query.filter.SortCriterion; import com.linkedin.metadata.search.LineageSearchResult; +import com.linkedin.metadata.service.ViewService; +import com.linkedin.metadata.utils.elasticsearch.FilterUtils; import com.linkedin.r2.RemoteInvocationException; +import com.linkedin.view.DataHubViewInfo; import graphql.VisibleForTesting; import graphql.schema.DataFetcher; import graphql.schema.DataFetchingEnvironment; @@ -53,10 +58,13 @@ public class SearchAcrossLineageResolver private final EntityRegistry _entityRegistry; + private final ViewService _viewService; + @VisibleForTesting final Set _allEntities; private final List _allowedEntities; - public SearchAcrossLineageResolver(EntityClient entityClient, EntityRegistry entityRegistry) { + public SearchAcrossLineageResolver( + EntityClient entityClient, EntityRegistry entityRegistry, final ViewService viewService) { this._entityClient = entityClient; this._entityRegistry = entityRegistry; this._allEntities = @@ -68,6 +76,8 @@ public SearchAcrossLineageResolver(EntityClient entityClient, EntityRegistry ent this._allEntities.stream() .filter(e -> !TRANSIENT_ENTITIES.contains(e)) .collect(Collectors.toList()); + + this._viewService = viewService; } private List getEntityNamesFromInput(List inputTypes) { @@ -127,6 +137,13 @@ public CompletableFuture get(DataFetchingEnvironment com.linkedin.metadata.graph.LineageDirection.valueOf(lineageDirection.toString()); return GraphQLConcurrencyUtils.supplyAsync( () -> { + final DataHubViewInfo maybeResolvedView = + (input.getViewUrn() != null) + ? resolveView( + context.getOperationContext(), + _viewService, + UrnUtils.getUrn(input.getViewUrn())) + : null; try { log.debug( "Executing search across relationships: source urn {}, direction {}, entity types {}, query {}, filters: {}, start: {}, count: {}", @@ -138,8 +155,13 @@ public CompletableFuture get(DataFetchingEnvironment start, count); - final Filter filter = + final Filter baseFilter = ResolverUtils.buildFilter(input.getFilters(), input.getOrFilters()); + Filter filter = + maybeResolvedView != null + ? FilterUtils.combineFilters( + baseFilter, maybeResolvedView.getDefinition().getFilter()) + : baseFilter; final SearchFlags searchFlags; com.linkedin.datahub.graphql.generated.SearchFlags inputFlags = input.getSearchFlags(); if (inputFlags != null) { @@ -150,6 +172,7 @@ public CompletableFuture get(DataFetchingEnvironment } else { searchFlags = new SearchFlags().setFulltext(true).setSkipHighlighting(true); } + List sortCriteria = SearchUtils.getSortCriteria(input.getSortInput()); LineageSearchResult salResults = _entityClient.searchAcrossLineage( context @@ -162,7 +185,7 @@ public CompletableFuture get(DataFetchingEnvironment sanitizedQuery, maxHops, filter, - null, + sortCriteria, start, count); diff --git a/datahub-graphql-core/src/main/resources/search.graphql b/datahub-graphql-core/src/main/resources/search.graphql index c0bec68cc23c5b..4e76309cd34074 100644 --- a/datahub-graphql-core/src/main/resources/search.graphql +++ b/datahub-graphql-core/src/main/resources/search.graphql @@ -417,6 +417,16 @@ input SearchAcrossLineageInput { Flags controlling the lineage query """ lineageFlags: LineageFlags + + """ + Optional - A View to apply when generating results + """ + viewUrn: String + + """ + Optional - Information on how to sort this search result + """ + sortInput: SearchSortInput } """ diff --git a/datahub-graphql-core/src/test/java/com/linkedin/datahub/graphql/resolvers/search/ScrollAcrossLineageResolverTest.java b/datahub-graphql-core/src/test/java/com/linkedin/datahub/graphql/resolvers/search/ScrollAcrossLineageResolverTest.java index a12f593253b533..0a9d6234eb4b8c 100644 --- a/datahub-graphql-core/src/test/java/com/linkedin/datahub/graphql/resolvers/search/ScrollAcrossLineageResolverTest.java +++ b/datahub-graphql-core/src/test/java/com/linkedin/datahub/graphql/resolvers/search/ScrollAcrossLineageResolverTest.java @@ -31,6 +31,7 @@ import com.linkedin.metadata.search.LineageSearchEntityArray; import com.linkedin.metadata.search.MatchedFieldArray; import com.linkedin.metadata.search.SearchResultMetadata; +import com.linkedin.metadata.service.ViewService; import graphql.schema.DataFetchingEnvironment; import io.datahubproject.metadata.context.OperationContext; import java.io.InputStream; @@ -76,7 +77,7 @@ public void testAllEntitiesInitialization() { InputStream inputStream = ClassLoader.getSystemResourceAsStream("entity-registry.yml"); EntityRegistry entityRegistry = new ConfigEntityRegistry(inputStream); SearchAcrossLineageResolver resolver = - new SearchAcrossLineageResolver(_entityClient, entityRegistry); + new SearchAcrossLineageResolver(_entityClient, entityRegistry, mock(ViewService.class)); assertTrue(resolver._allEntities.contains("dataset")); assertTrue(resolver._allEntities.contains("dataFlow")); // Test for case sensitivity diff --git a/datahub-graphql-core/src/test/java/com/linkedin/datahub/graphql/resolvers/search/SearchAcrossLineageResolverTest.java b/datahub-graphql-core/src/test/java/com/linkedin/datahub/graphql/resolvers/search/SearchAcrossLineageResolverTest.java index 153e98149ff1a5..d7c1a230811945 100644 --- a/datahub-graphql-core/src/test/java/com/linkedin/datahub/graphql/resolvers/search/SearchAcrossLineageResolverTest.java +++ b/datahub-graphql-core/src/test/java/com/linkedin/datahub/graphql/resolvers/search/SearchAcrossLineageResolverTest.java @@ -22,6 +22,7 @@ import com.linkedin.metadata.search.LineageSearchResult; import com.linkedin.metadata.search.MatchedFieldArray; import com.linkedin.metadata.search.SearchResultMetadata; +import com.linkedin.metadata.service.ViewService; import graphql.schema.DataFetchingEnvironment; import io.datahubproject.metadata.context.OperationContext; import java.io.InputStream; @@ -48,6 +49,7 @@ public class SearchAcrossLineageResolverTest { private SearchAcrossLineageResolver _resolver; private EntityRegistry _entityRegistry; + private ViewService _viewService; @BeforeMethod public void setupTest() { @@ -56,7 +58,8 @@ public void setupTest() { _authentication = mock(Authentication.class); _entityRegistry = mock(EntityRegistry.class); - _resolver = new SearchAcrossLineageResolver(_entityClient, _entityRegistry); + _viewService = mock(ViewService.class); + _resolver = new SearchAcrossLineageResolver(_entityClient, _entityRegistry, _viewService); } @Test @@ -64,7 +67,7 @@ public void testAllEntitiesInitialization() { InputStream inputStream = ClassLoader.getSystemResourceAsStream("entity-registry.yml"); EntityRegistry entityRegistry = new ConfigEntityRegistry(inputStream); SearchAcrossLineageResolver resolver = - new SearchAcrossLineageResolver(_entityClient, entityRegistry); + new SearchAcrossLineageResolver(_entityClient, entityRegistry, _viewService); assertTrue(resolver._allEntities.contains("dataset")); assertTrue(resolver._allEntities.contains("dataFlow")); // Test for case sensitivity @@ -115,7 +118,7 @@ public void testSearchAcrossLineage() throws Exception { eq(QUERY), eq(null), any(), - eq(null), + eq(Collections.emptyList()), eq(START), eq(COUNT))) .thenReturn(lineageSearchResult); diff --git a/metadata-utils/src/main/java/com/linkedin/metadata/utils/elasticsearch/FilterUtils.java b/metadata-utils/src/main/java/com/linkedin/metadata/utils/elasticsearch/FilterUtils.java new file mode 100644 index 00000000000000..d803d120d7a302 --- /dev/null +++ b/metadata-utils/src/main/java/com/linkedin/metadata/utils/elasticsearch/FilterUtils.java @@ -0,0 +1,130 @@ +package com.linkedin.metadata.utils.elasticsearch; + +import com.google.common.collect.ImmutableList; +import com.linkedin.data.template.StringArray; +import com.linkedin.metadata.query.filter.Condition; +import com.linkedin.metadata.query.filter.ConjunctiveCriterion; +import com.linkedin.metadata.query.filter.ConjunctiveCriterionArray; +import com.linkedin.metadata.query.filter.Criterion; +import com.linkedin.metadata.query.filter.CriterionArray; +import com.linkedin.metadata.query.filter.Filter; +import com.linkedin.metadata.utils.CriterionUtils; +import java.util.ArrayList; +import java.util.Collections; +import java.util.List; +import javax.annotation.Nonnull; +import javax.annotation.Nullable; + +public class FilterUtils { + + /** + * Combines two {@link Filter} instances in a conjunction and returns a new instance of {@link + * Filter} in disjunctive normal form. + * + * @param baseFilter the filter to apply the view to + * @param viewFilter the view filter, null if it doesn't exist + * @return a new instance of {@link Filter} representing the applied view. + */ + @Nonnull + public static Filter combineFilters( + @Nullable final Filter baseFilter, @Nonnull final Filter viewFilter) { + final Filter finalBaseFilter = + baseFilter == null + ? new Filter().setOr(new ConjunctiveCriterionArray(Collections.emptyList())) + : baseFilter; + + // Join the filter conditions in Disjunctive Normal Form. + return combineFiltersInConjunction(finalBaseFilter, viewFilter); + } + + /** + * Joins two filters in conjunction by reducing to Disjunctive Normal Form. + * + * @param filter1 the first filter in the pair + * @param filter2 the second filter in the pair + *

This method supports either Filter format, where the "or" field is used, instead of + * criteria. If the criteria filter is used, then it will be converted into an "OR" before + * returning the new filter. + * @return the result of joining the 2 filters in a conjunction (AND) + *

How does it work? It basically cross-products the conjunctions inside of each Filter + * clause. + *

Example Inputs: filter1 -> { or: [ { and: [ { field: tags, condition: EQUAL, values: + * ["urn:li:tag:tag"] } ] }, { and: [ { field: glossaryTerms, condition: EQUAL, values: + * ["urn:li:glossaryTerm:term"] } ] } ] } filter2 -> { or: [ { and: [ { field: domain, + * condition: EQUAL, values: ["urn:li:domain:domain"] }, ] }, { and: [ { field: glossaryTerms, + * condition: EQUAL, values: ["urn:li:glossaryTerm:term2"] } ] } ] } Example Output: { or: [ { + * and: [ { field: tags, condition: EQUAL, values: ["urn:li:tag:tag"] }, { field: domain, + * condition: EQUAL, values: ["urn:li:domain:domain"] } ] }, { and: [ { field: tags, + * condition: EQUAL, values: ["urn:li:tag:tag"] }, { field: glossaryTerms, condition: EQUAL, + * values: ["urn:li:glosaryTerm:term2"] } ] }, { and: [ { field: glossaryTerm, condition: + * EQUAL, values: ["urn:li:glossaryTerm:term"] }, { field: domain, condition: EQUAL, values: + * ["urn:li:domain:domain"] } ] }, { and: [ { field: glossaryTerm, condition: EQUAL, values: + * ["urn:li:glossaryTerm:term"] }, { field: glossaryTerms, condition: EQUAL, values: + * ["urn:li:glosaryTerm:term2"] } ] }, ] } + */ + @Nonnull + private static Filter combineFiltersInConjunction( + @Nonnull final Filter filter1, @Nonnull final Filter filter2) { + + final Filter finalFilter1 = convertToV2Filter(filter1); + final Filter finalFilter2 = convertToV2Filter(filter2); + + // If either filter is empty, simply return the other filter. + if (!finalFilter1.hasOr() || finalFilter1.getOr().size() == 0) { + return finalFilter2; + } + if (!finalFilter2.hasOr() || finalFilter2.getOr().size() == 0) { + return finalFilter1; + } + + // Iterate through the base filter, then cross-product with filter 2 conditions. + final Filter result = new Filter(); + final List newDisjunction = new ArrayList<>(); + for (ConjunctiveCriterion conjunction1 : finalFilter1.getOr()) { + for (ConjunctiveCriterion conjunction2 : finalFilter2.getOr()) { + final List joinedCriterion = new ArrayList<>(conjunction1.getAnd()); + joinedCriterion.addAll(conjunction2.getAnd()); + ConjunctiveCriterion newConjunction = + new ConjunctiveCriterion().setAnd(new CriterionArray(joinedCriterion)); + newDisjunction.add(newConjunction); + } + } + result.setOr(new ConjunctiveCriterionArray(newDisjunction)); + return result; + } + + @Nonnull + private static Filter convertToV2Filter(@Nonnull Filter filter) { + if (filter.hasOr()) { + return filter; + } else if (filter.hasCriteria()) { + // Convert criteria to an OR + return new Filter() + .setOr( + new ConjunctiveCriterionArray( + ImmutableList.of(new ConjunctiveCriterion().setAnd(filter.getCriteria())))); + } + throw new IllegalArgumentException( + String.format( + "Illegal filter provided! Neither 'or' nor 'criteria' fields were populated for filter %s", + filter)); + } + + @Nonnull + public static Filter createValuesFilter( + @Nonnull final String fieldName, @Nonnull final List values) { + Filter filter = new Filter(); + CriterionArray criterionArray = new CriterionArray(); + + StringArray valuesArray = new StringArray(); + valuesArray.addAll(values); + Criterion criterion = CriterionUtils.buildCriterion(fieldName, Condition.EQUAL, valuesArray); + + criterionArray.add(criterion); + filter.setOr( + new ConjunctiveCriterionArray( + ImmutableList.of(new ConjunctiveCriterion().setAnd(criterionArray)))); + + return filter; + } +} diff --git a/metadata-utils/src/test/java/com/linkedin/metadata/utils/FilterUtilsTest.java b/metadata-utils/src/test/java/com/linkedin/metadata/utils/FilterUtilsTest.java new file mode 100644 index 00000000000000..b0243dbc44e8b4 --- /dev/null +++ b/metadata-utils/src/test/java/com/linkedin/metadata/utils/FilterUtilsTest.java @@ -0,0 +1,213 @@ +package com.linkedin.metadata.utils; + +import static com.linkedin.metadata.utils.CriterionUtils.buildCriterion; + +import com.google.common.collect.ImmutableList; +import com.linkedin.metadata.query.filter.Condition; +import com.linkedin.metadata.query.filter.ConjunctiveCriterion; +import com.linkedin.metadata.query.filter.ConjunctiveCriterionArray; +import com.linkedin.metadata.query.filter.CriterionArray; +import com.linkedin.metadata.query.filter.Filter; +import com.linkedin.metadata.utils.elasticsearch.FilterUtils; +import org.testng.Assert; +import org.testng.annotations.Test; + +public class FilterUtilsTest { + + @Test + public static void testApplyViewToFilterNullBaseFilter() { + + Filter viewFilter = + new Filter() + .setOr( + new ConjunctiveCriterionArray( + new ConjunctiveCriterion() + .setAnd( + new CriterionArray( + ImmutableList.of( + buildCriterion("field", Condition.EQUAL, "test")))))); + + Filter result = FilterUtils.combineFilters(null, viewFilter); + Assert.assertEquals(viewFilter, result); + } + + @Test + public static void testApplyViewToFilterComplexBaseFilter() { + Filter baseFilter = + new Filter() + .setOr( + new ConjunctiveCriterionArray( + ImmutableList.of( + new ConjunctiveCriterion() + .setAnd( + new CriterionArray( + ImmutableList.of( + buildCriterion("field1", Condition.EQUAL, "test1"), + buildCriterion("field2", Condition.EQUAL, "test2")))), + new ConjunctiveCriterion() + .setAnd( + new CriterionArray( + ImmutableList.of( + buildCriterion("field3", Condition.EQUAL, "test3"), + buildCriterion("field4", Condition.EQUAL, "test4"))))))); + + Filter viewFilter = + new Filter() + .setOr( + new ConjunctiveCriterionArray( + new ConjunctiveCriterion() + .setAnd( + new CriterionArray( + ImmutableList.of( + buildCriterion("field", Condition.EQUAL, "test")))))); + + Filter result = FilterUtils.combineFilters(baseFilter, viewFilter); + + Filter expectedResult = + new Filter() + .setOr( + new ConjunctiveCriterionArray( + ImmutableList.of( + new ConjunctiveCriterion() + .setAnd( + new CriterionArray( + ImmutableList.of( + buildCriterion("field1", Condition.EQUAL, "test1"), + buildCriterion("field2", Condition.EQUAL, "test2"), + buildCriterion("field", Condition.EQUAL, "test")))), + new ConjunctiveCriterion() + .setAnd( + new CriterionArray( + ImmutableList.of( + buildCriterion("field3", Condition.EQUAL, "test3"), + buildCriterion("field4", Condition.EQUAL, "test4"), + buildCriterion("field", Condition.EQUAL, "test"))))))); + + Assert.assertEquals(expectedResult, result); + } + + @Test + public static void testApplyViewToFilterComplexViewFilter() { + Filter baseFilter = + new Filter() + .setOr( + new ConjunctiveCriterionArray( + ImmutableList.of( + new ConjunctiveCriterion() + .setAnd( + new CriterionArray( + ImmutableList.of( + buildCriterion("field1", Condition.EQUAL, "test1"), + buildCriterion("field2", Condition.EQUAL, "test2")))), + new ConjunctiveCriterion() + .setAnd( + new CriterionArray( + ImmutableList.of( + buildCriterion("field3", Condition.EQUAL, "test3"), + buildCriterion("field4", Condition.EQUAL, "test4"))))))); + + Filter viewFilter = + new Filter() + .setOr( + new ConjunctiveCriterionArray( + ImmutableList.of( + new ConjunctiveCriterion() + .setAnd( + new CriterionArray( + ImmutableList.of( + buildCriterion("viewField1", Condition.EQUAL, "viewTest1"), + buildCriterion( + "viewField2", Condition.EQUAL, "viewTest2")))), + new ConjunctiveCriterion() + .setAnd( + new CriterionArray( + ImmutableList.of( + buildCriterion("viewField3", Condition.EQUAL, "viewTest3"), + buildCriterion( + "viewField4", Condition.EQUAL, "viewTest4"))))))); + + Filter result = FilterUtils.combineFilters(baseFilter, viewFilter); + + Filter expectedResult = + new Filter() + .setOr( + new ConjunctiveCriterionArray( + ImmutableList.of( + new ConjunctiveCriterion() + .setAnd( + new CriterionArray( + ImmutableList.of( + buildCriterion("field1", Condition.EQUAL, "test1"), + buildCriterion("field2", Condition.EQUAL, "test2"), + buildCriterion("viewField1", Condition.EQUAL, "viewTest1"), + buildCriterion( + "viewField2", Condition.EQUAL, "viewTest2")))), + new ConjunctiveCriterion() + .setAnd( + new CriterionArray( + ImmutableList.of( + buildCriterion("field1", Condition.EQUAL, "test1"), + buildCriterion("field2", Condition.EQUAL, "test2"), + buildCriterion("viewField3", Condition.EQUAL, "viewTest3"), + buildCriterion( + "viewField4", Condition.EQUAL, "viewTest4")))), + new ConjunctiveCriterion() + .setAnd( + new CriterionArray( + ImmutableList.of( + buildCriterion("field3", Condition.EQUAL, "test3"), + buildCriterion("field4", Condition.EQUAL, "test4"), + buildCriterion("viewField1", Condition.EQUAL, "viewTest1"), + buildCriterion( + "viewField2", Condition.EQUAL, "viewTest2")))), + new ConjunctiveCriterion() + .setAnd( + new CriterionArray( + ImmutableList.of( + buildCriterion("field3", Condition.EQUAL, "test3"), + buildCriterion("field4", Condition.EQUAL, "test4"), + buildCriterion("viewField3", Condition.EQUAL, "viewTest3"), + buildCriterion( + "viewField4", Condition.EQUAL, "viewTest4"))))))); + + Assert.assertEquals(expectedResult, result); + } + + @Test + public static void testApplyViewToFilterV1Filter() { + Filter baseFilter = + new Filter() + .setCriteria( + new CriterionArray( + ImmutableList.of( + buildCriterion("field1", Condition.EQUAL, "test1"), + buildCriterion("field2", Condition.EQUAL, "test2")))); + + Filter viewFilter = + new Filter() + .setCriteria( + new CriterionArray( + ImmutableList.of( + buildCriterion("viewField1", Condition.EQUAL, "viewTest1"), + buildCriterion("viewField2", Condition.EQUAL, "viewTest2")))); + + Filter result = FilterUtils.combineFilters(baseFilter, viewFilter); + + Filter expectedResult = + new Filter() + .setOr( + new ConjunctiveCriterionArray( + ImmutableList.of( + new ConjunctiveCriterion() + .setAnd( + new CriterionArray( + ImmutableList.of( + buildCriterion("field1", Condition.EQUAL, "test1"), + buildCriterion("field2", Condition.EQUAL, "test2"), + buildCriterion("viewField1", Condition.EQUAL, "viewTest1"), + buildCriterion( + "viewField2", Condition.EQUAL, "viewTest2"))))))); + + Assert.assertEquals(expectedResult, result); + } +}