Merge branch 'datahub-project:master' into master

datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/ResolverUtils.java

@@ -15,6 +15,7 @@
 import com.linkedin.datahub.graphql.generated.AndFilterInput;
 import com.linkedin.datahub.graphql.generated.FacetFilterInput;
 import com.linkedin.datahub.graphql.resolvers.search.SearchUtils;
+import com.linkedin.metadata.aspect.AspectRetriever;
 import com.linkedin.metadata.query.filter.Condition;
 import com.linkedin.metadata.query.filter.ConjunctiveCriterion;
 import com.linkedin.metadata.query.filter.ConjunctiveCriterionArray;
@@ -112,10 +113,11 @@ public static Map<String, String> buildFacetFilters(
     return facetFilters;
   }
 
-  public static List<Criterion> criterionListFromAndFilter(List<FacetFilterInput> andFilters) {
+  public static List<Criterion> criterionListFromAndFilter(
+      List<FacetFilterInput> andFilters, @Nullable AspectRetriever aspectRetriever) {
     return andFilters != null && !andFilters.isEmpty()
         ? andFilters.stream()
-            .map(filter -> criterionFromFilter(filter))
+            .map(filter -> criterionFromFilter(filter, aspectRetriever))
             .collect(Collectors.toList())
         : Collections.emptyList();
   }
@@ -124,21 +126,24 @@ public static List<Criterion> criterionListFromAndFilter(List<FacetFilterInput>
   // conjunctive criterion
   // arrays, rather than just one for the AND case.
   public static ConjunctiveCriterionArray buildConjunctiveCriterionArrayWithOr(
-      @Nonnull List<AndFilterInput> orFilters) {
+      @Nonnull List<AndFilterInput> orFilters, @Nullable AspectRetriever aspectRetriever) {
     return new ConjunctiveCriterionArray(
         orFilters.stream()
             .map(
                 orFilter -> {
                   CriterionArray andCriterionForOr =
-                      new CriterionArray(criterionListFromAndFilter(orFilter.getAnd()));
+                      new CriterionArray(
+                          criterionListFromAndFilter(orFilter.getAnd(), aspectRetriever));
                   return new ConjunctiveCriterion().setAnd(andCriterionForOr);
                 })
             .collect(Collectors.toList()));
   }
 
   @Nullable
   public static Filter buildFilter(
-      @Nullable List<FacetFilterInput> andFilters, @Nullable List<AndFilterInput> orFilters) {
+      @Nullable List<FacetFilterInput> andFilters,
+      @Nullable List<AndFilterInput> orFilters,
+      @Nullable AspectRetriever aspectRetriever) {
     if ((andFilters == null || andFilters.isEmpty())
         && (orFilters == null || orFilters.isEmpty())) {
       return null;
@@ -147,30 +152,33 @@ public static Filter buildFilter(
     // Or filters are the new default. We will check them first.
     // If we have OR filters, we need to build a series of CriterionArrays
     if (orFilters != null && !orFilters.isEmpty()) {
-      return new Filter().setOr(buildConjunctiveCriterionArrayWithOr(orFilters));
+      return new Filter().setOr(buildConjunctiveCriterionArrayWithOr(orFilters, aspectRetriever));
     }
 
     // If or filters are not set, someone may be using the legacy and filters
-    final List<Criterion> andCriterions = criterionListFromAndFilter(andFilters);
+    final List<Criterion> andCriterions = criterionListFromAndFilter(andFilters, aspectRetriever);
     return new Filter()
         .setOr(
             new ConjunctiveCriterionArray(
                 new ConjunctiveCriterion().setAnd(new CriterionArray(andCriterions))));
   }
 
-  public static Criterion criterionFromFilter(final FacetFilterInput filter) {
-    return criterionFromFilter(filter, false);
+  public static Criterion criterionFromFilter(
+      final FacetFilterInput filter, @Nullable AspectRetriever aspectRetriever) {
+    return criterionFromFilter(filter, false, aspectRetriever);
   }
 
   // Translates a FacetFilterInput (graphql input class) into Criterion (our internal model)
   public static Criterion criterionFromFilter(
-      final FacetFilterInput filter, final Boolean skipKeywordSuffix) {
+      final FacetFilterInput filter,
+      final Boolean skipKeywordSuffix,
+      @Nullable AspectRetriever aspectRetriever) {
     Criterion result = new Criterion();
 
     if (skipKeywordSuffix) {
       result.setField(filter.getField());
     } else {
-      result.setField(getFilterField(filter.getField(), skipKeywordSuffix));
+      result.setField(getFilterField(filter.getField(), skipKeywordSuffix, aspectRetriever));
     }
 
     // `value` is deprecated in place of `values`- this is to support old query patterns. If values
@@ -205,11 +213,13 @@ public static Criterion criterionFromFilter(
   }
 
   private static String getFilterField(
-      final String originalField, final boolean skipKeywordSuffix) {
+      final String originalField,
+      final boolean skipKeywordSuffix,
+      @Nullable AspectRetriever aspectRetriever) {
     if (KEYWORD_EXCLUDED_FILTERS.contains(originalField)) {
       return originalField;
     }
-    return ESUtils.toKeywordField(originalField, skipKeywordSuffix);
+    return ESUtils.toKeywordField(originalField, skipKeywordSuffix, aspectRetriever);
   }
 
   public static Filter buildFilterWithUrns(@Nonnull Set<Urn> urns, @Nullable Filter inputFilters) {

datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/assertion/AssertionRunEventResolver.java

@@ -15,6 +15,7 @@
 import com.linkedin.datahub.graphql.types.dataset.mappers.AssertionRunEventMapper;
 import com.linkedin.entity.client.EntityClient;
 import com.linkedin.metadata.Constants;
+import com.linkedin.metadata.aspect.AspectRetriever;
 import com.linkedin.metadata.aspect.EnvelopedAspect;
 import com.linkedin.metadata.query.filter.ConjunctiveCriterion;
 import com.linkedin.metadata.query.filter.ConjunctiveCriterionArray;
@@ -67,7 +68,10 @@ public CompletableFuture<AssertionRunEventsResult> get(DataFetchingEnvironment e
                     maybeStartTimeMillis,
                     maybeEndTimeMillis,
                     maybeLimit,
-                    buildFilter(maybeFilters, maybeStatus));
+                    buildFilter(
+                        maybeFilters,
+                        maybeStatus,
+                        context.getOperationContext().getAspectRetriever()));
 
             // Step 2: Bind profiles into GraphQL strong types.
             List<AssertionRunEvent> runEvents =
@@ -120,7 +124,9 @@ public CompletableFuture<AssertionRunEventsResult> get(DataFetchingEnvironment e
 
   @Nullable
   public static Filter buildFilter(
-      @Nullable FilterInput filtersInput, @Nullable final String status) {
+      @Nullable FilterInput filtersInput,
+      @Nullable final String status,
+      @Nullable AspectRetriever aspectRetriever) {
     if (filtersInput == null && status == null) {
       return null;
     }
@@ -141,7 +147,7 @@ public static Filter buildFilter(
                     .setAnd(
                         new CriterionArray(
                             facetFilters.stream()
-                                .map(filter -> criterionFromFilter(filter, true))
+                                .map(filter -> criterionFromFilter(filter, true, aspectRetriever))
                                 .collect(Collectors.toList())))));
   }
 }

datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/auth/ListAccessTokensResolver.java

@@ -70,7 +70,10 @@ public CompletableFuture<ListAccessTokenResult> get(DataFetchingEnvironment envi
                           .withSearchFlags(flags -> flags.setFulltext(true)),
                       Constants.ACCESS_TOKEN_ENTITY_NAME,
                       "",
-                      buildFilter(filters, Collections.emptyList()),
+                      buildFilter(
+                          filters,
+                          Collections.emptyList(),
+                          context.getOperationContext().getAspectRetriever()),
                       sortCriterion,
                       start,
                       count);

datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/chart/BrowseV2Resolver.java

@@ -74,7 +74,9 @@ public CompletableFuture<BrowseResultsV2> get(DataFetchingEnvironment environmen
                     ? BROWSE_PATH_V2_DELIMITER
                         + String.join(BROWSE_PATH_V2_DELIMITER, input.getPath())
                     : "";
-            final Filter inputFilter = ResolverUtils.buildFilter(null, input.getOrFilters());
+            final Filter inputFilter =
+                ResolverUtils.buildFilter(
+                    null, input.getOrFilters(), context.getOperationContext().getAspectRetriever());
 
             BrowseResultV2 browseResults =
                 _entityClient.browseV2(

datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/dataproduct/ListDataProductAssetsResolver.java

@@ -132,7 +132,10 @@ public CompletableFuture<SearchResults> get(DataFetchingEnvironment environment)
 
           // add urns from the aspect to our filters
           final Filter baseFilter =
-              ResolverUtils.buildFilter(input.getFilters(), input.getOrFilters());
+              ResolverUtils.buildFilter(
+                  input.getFilters(),
+                  input.getOrFilters(),
+                  context.getOperationContext().getAspectRetriever());
           final Filter finalFilter = buildFilterWithUrns(new HashSet<>(assetUrns), baseFilter);
 
           final SearchFlags searchFlags;

datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/domain/DomainEntitiesResolver.java

@@ -78,7 +78,9 @@ public CompletableFuture<SearchResults> get(final DataFetchingEnvironment enviro
                   .getFilters()
                   .forEach(
                       filter -> {
-                        criteria.add(criterionFromFilter(filter, true));
+                        criteria.add(
+                            criterionFromFilter(
+                                filter, true, context.getOperationContext().getAspectRetriever()));
                       });
             }
 

datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/form/CreateDynamicFormAssignmentResolver.java

@@ -33,7 +33,9 @@ public CompletableFuture<Boolean> get(final DataFetchingEnvironment environment)
     final CreateDynamicFormAssignmentInput input =
         bindArgument(environment.getArgument("input"), CreateDynamicFormAssignmentInput.class);
     final Urn formUrn = UrnUtils.getUrn(input.getFormUrn());
-    final DynamicFormAssignment formAssignment = FormUtils.mapDynamicFormAssignment(input);
+    final DynamicFormAssignment formAssignment =
+        FormUtils.mapDynamicFormAssignment(
+            input, context.getOperationContext().getAspectRetriever());
 
     return GraphQLConcurrencyUtils.supplyAsync(
         () -> {

datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/ingest/source/ListIngestionSourcesResolver.java

@@ -68,7 +68,10 @@ public CompletableFuture<ListIngestionSourcesResult> get(
                           .withSearchFlags(flags -> flags.setFulltext(true)),
                       Constants.INGESTION_SOURCE_ENTITY_NAME,
                       query,
-                      buildFilter(filters, Collections.emptyList()),
+                      buildFilter(
+                          filters,
+                          Collections.emptyList(),
+                          context.getOperationContext().getAspectRetriever()),
                       null,
                       start,
                       count);

datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/load/TimeSeriesAspectResolver.java

@@ -11,6 +11,7 @@
 import com.linkedin.datahub.graphql.generated.TimeSeriesAspect;
 import com.linkedin.entity.client.EntityClient;
 import com.linkedin.metadata.Constants;
+import com.linkedin.metadata.aspect.AspectRetriever;
 import com.linkedin.metadata.aspect.EnvelopedAspect;
 import com.linkedin.metadata.authorization.PoliciesConfig;
 import com.linkedin.metadata.query.filter.ConjunctiveCriterion;
@@ -120,7 +121,7 @@ public CompletableFuture<List<TimeSeriesAspect>> get(DataFetchingEnvironment env
                     maybeStartTimeMillis,
                     maybeEndTimeMillis,
                     maybeLimit,
-                    buildFilters(maybeFilters),
+                    buildFilters(maybeFilters, context.getOperationContext().getAspectRetriever()),
                     maybeSort);
 
             // Step 2: Bind profiles into GraphQL strong types.
@@ -135,7 +136,8 @@ public CompletableFuture<List<TimeSeriesAspect>> get(DataFetchingEnvironment env
         "get");
   }
 
-  private Filter buildFilters(@Nullable FilterInput maybeFilters) {
+  private Filter buildFilters(
+      @Nullable FilterInput maybeFilters, @Nullable AspectRetriever aspectRetriever) {
     if (maybeFilters == null) {
       return null;
     }
@@ -146,7 +148,7 @@ private Filter buildFilters(@Nullable FilterInput maybeFilters) {
                     .setAnd(
                         new CriterionArray(
                             maybeFilters.getAnd().stream()
-                                .map(filter -> criterionFromFilter(filter, true))
+                                .map(filter -> criterionFromFilter(filter, true, aspectRetriever))
                                 .collect(Collectors.toList())))));
   }
 }

datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/mutate/util/FormUtils.java

@@ -6,6 +6,7 @@
 import com.linkedin.datahub.graphql.resolvers.ResolverUtils;
 import com.linkedin.form.DynamicFormAssignment;
 import com.linkedin.form.FormInfo;
+import com.linkedin.metadata.aspect.AspectRetriever;
 import com.linkedin.metadata.query.filter.Condition;
 import com.linkedin.metadata.query.filter.ConjunctiveCriterion;
 import com.linkedin.metadata.query.filter.ConjunctiveCriterionArray;
@@ -15,6 +16,7 @@
 import com.linkedin.structured.PrimitivePropertyValueArray;
 import java.util.Objects;
 import javax.annotation.Nonnull;
+import javax.annotation.Nullable;
 
 public class FormUtils {
 
@@ -44,13 +46,16 @@ public static PrimitivePropertyValueArray getStructuredPropertyValuesFromInput(
   /** Map a GraphQL CreateDynamicFormAssignmentInput to the GMS DynamicFormAssignment aspect */
   @Nonnull
   public static DynamicFormAssignment mapDynamicFormAssignment(
-      @Nonnull final CreateDynamicFormAssignmentInput input) {
+      @Nonnull final CreateDynamicFormAssignmentInput input,
+      @Nullable AspectRetriever aspectRetriever) {
     Objects.requireNonNull(input, "input must not be null");
 
     final DynamicFormAssignment result = new DynamicFormAssignment();
     final Filter filter =
         new Filter()
-            .setOr(ResolverUtils.buildConjunctiveCriterionArrayWithOr(input.getOrFilters()));
+            .setOr(
+                ResolverUtils.buildConjunctiveCriterionArrayWithOr(
+                    input.getOrFilters(), aspectRetriever));
     result.setFilter(filter);
     return result;
   }

datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/ownership/ListOwnershipTypesResolver.java

@@ -63,7 +63,10 @@ public CompletableFuture<ListOwnershipTypesResult> get(DataFetchingEnvironment e
                     context.getOperationContext().withSearchFlags(flags -> flags.setFulltext(true)),
                     Constants.OWNERSHIP_TYPE_ENTITY_NAME,
                     query,
-                    buildFilter(filters, Collections.emptyList()),
+                    buildFilter(
+                        filters,
+                        Collections.emptyList(),
+                        context.getOperationContext().getAspectRetriever()),
                     DEFAULT_SORT_CRITERION,
                     start,
                     count);

datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/policy/ListPoliciesResolver.java

@@ -59,7 +59,11 @@ public CompletableFuture<ListPoliciesResult> get(final DataFetchingEnvironment e
       log.debug(
           "User {} listing policies with filters {}", context.getActorUrn(), filters.toString());
 
-      final Filter filter = ResolverUtils.buildFilter(facetFilters, Collections.emptyList());
+      final Filter filter =
+          ResolverUtils.buildFilter(
+              facetFilters,
+              Collections.emptyList(),
+              context.getOperationContext().getAspectRetriever());
 
       return _policyFetcher
           .fetchPolicies(context.getOperationContext(), start, query, count, filter)

datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/query/ListQueriesResolver.java

@@ -15,6 +15,7 @@
 import com.linkedin.datahub.graphql.generated.ListQueriesResult;
 import com.linkedin.datahub.graphql.generated.QueryEntity;
 import com.linkedin.entity.client.EntityClient;
+import com.linkedin.metadata.aspect.AspectRetriever;
 import com.linkedin.metadata.query.filter.Filter;
 import com.linkedin.metadata.query.filter.SortCriterion;
 import com.linkedin.metadata.query.filter.SortOrder;
@@ -72,7 +73,7 @@ public CompletableFuture<ListQueriesResult> get(final DataFetchingEnvironment en
                             flags -> flags.setFulltext(true).setSkipHighlighting(true)),
                     QUERY_ENTITY_NAME,
                     query,
-                    buildFilters(input),
+                    buildFilters(input, context.getOperationContext().getAspectRetriever()),
                     sortCriterion,
                     start,
                     count);
@@ -109,7 +110,8 @@ private List<QueryEntity> mapUnresolvedQueries(final List<Urn> queryUrns) {
   }
 
   @Nullable
-  private Filter buildFilters(@Nonnull final ListQueriesInput input) {
+  private Filter buildFilters(
+      @Nonnull final ListQueriesInput input, @Nullable AspectRetriever aspectRetriever) {
     final AndFilterInput criteria = new AndFilterInput();
     List<FacetFilterInput> andConditions = new ArrayList<>();
 
@@ -136,6 +138,6 @@ private Filter buildFilters(@Nonnull final ListQueriesInput input) {
     }
 
     criteria.setAnd(andConditions);
-    return buildFilter(Collections.emptyList(), ImmutableList.of(criteria));
+    return buildFilter(Collections.emptyList(), ImmutableList.of(criteria), aspectRetriever);
   }
 }

datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/recommendation/ListRecommendationsResolver.java

@@ -26,13 +26,15 @@
 import com.linkedin.metadata.service.ViewService;
 import graphql.schema.DataFetcher;
 import graphql.schema.DataFetchingEnvironment;
+import io.datahubproject.metadata.context.OperationContext;
 import io.opentelemetry.extension.annotations.WithSpan;
 import java.net.URISyntaxException;
 import java.util.Collections;
 import java.util.List;
 import java.util.Optional;
 import java.util.concurrent.CompletableFuture;
 import java.util.stream.Collectors;
+import javax.annotation.Nonnull;
 import javax.annotation.Nullable;
 import lombok.RequiredArgsConstructor;
 import lombok.extern.slf4j.Slf4j;
@@ -62,7 +64,7 @@ public CompletableFuture<ListRecommendationsResult> get(DataFetchingEnvironment
             List<com.linkedin.metadata.recommendation.RecommendationModule> modules =
                 _recommendationsService.listRecommendations(
                     context.getOperationContext(),
-                    mapRequestContext(input.getRequestContext()),
+                    mapRequestContext(context.getOperationContext(), input.getRequestContext()),
                     viewFilter(context.getOperationContext(), _viewService, input.getViewUrn()),
                     input.getLimit());
             return ListRecommendationsResult.builder()
@@ -83,7 +85,7 @@ public CompletableFuture<ListRecommendationsResult> get(DataFetchingEnvironment
   }
 
   private com.linkedin.metadata.recommendation.RecommendationRequestContext mapRequestContext(
-      RecommendationRequestContext requestContext) {
+      @Nonnull OperationContext opContext, RecommendationRequestContext requestContext) {
     com.linkedin.metadata.recommendation.ScenarioType mappedScenarioType;
     try {
       mappedScenarioType =
@@ -103,7 +105,9 @@ private com.linkedin.metadata.recommendation.RecommendationRequestContext mapReq
         searchRequestContext.setFilters(
             new CriterionArray(
                 requestContext.getSearchRequestContext().getFilters().stream()
-                    .map(facetField -> criterionFromFilter(facetField))
+                    .map(
+                        facetField ->
+                            criterionFromFilter(facetField, opContext.getAspectRetriever()))
                     .collect(Collectors.toList())));
       }
       mappedRequestContext.setSearchRequestContext(searchRequestContext);