Skip to content

Commit

Permalink
refactor(search): refactor field type detection (datahub-project#11395)
Browse files Browse the repository at this point in the history
  • Loading branch information
david-leifker authored Sep 17, 2024
1 parent c2977d8 commit 6a165a8
Show file tree
Hide file tree
Showing 47 changed files with 277 additions and 223 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -6,22 +6,19 @@
import com.datahub.authentication.Authentication;
import com.fasterxml.jackson.core.StreamReadConstraints;
import com.fasterxml.jackson.databind.ObjectMapper;
import com.google.common.collect.ImmutableSet;
import com.linkedin.common.urn.UrnUtils;
import com.linkedin.data.template.StringArray;
import com.linkedin.datahub.graphql.QueryContext;
import com.linkedin.datahub.graphql.exception.ValidationException;
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;
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.search.utils.ESUtils;
import com.linkedin.metadata.service.ViewService;
import com.linkedin.view.DataHubViewInfo;
import graphql.schema.DataFetchingEnvironment;
Expand All @@ -39,8 +36,6 @@

public class ResolverUtils {

private static final Set<String> KEYWORD_EXCLUDED_FILTERS =
ImmutableSet.of("runId", "_entityType");
private static final ObjectMapper MAPPER = new ObjectMapper();

static {
Expand Down Expand Up @@ -111,11 +106,10 @@ public static Map<String, String> buildFacetFilters(
return facetFilters;
}

public static List<Criterion> criterionListFromAndFilter(
List<FacetFilterInput> andFilters, @Nullable AspectRetriever aspectRetriever) {
public static List<Criterion> criterionListFromAndFilter(List<FacetFilterInput> andFilters) {
return andFilters != null && !andFilters.isEmpty()
? andFilters.stream()
.map(filter -> criterionFromFilter(filter, aspectRetriever))
.map(filter -> criterionFromFilter(filter))
.collect(Collectors.toList())
: Collections.emptyList();
}
Expand All @@ -124,24 +118,21 @@ public static List<Criterion> criterionListFromAndFilter(
// conjunctive criterion
// arrays, rather than just one for the AND case.
public static ConjunctiveCriterionArray buildConjunctiveCriterionArrayWithOr(
@Nonnull List<AndFilterInput> orFilters, @Nullable AspectRetriever aspectRetriever) {
@Nonnull List<AndFilterInput> orFilters) {
return new ConjunctiveCriterionArray(
orFilters.stream()
.map(
orFilter -> {
CriterionArray andCriterionForOr =
new CriterionArray(
criterionListFromAndFilter(orFilter.getAnd(), aspectRetriever));
new CriterionArray(criterionListFromAndFilter(orFilter.getAnd()));
return new ConjunctiveCriterion().setAnd(andCriterionForOr);
})
.collect(Collectors.toList()));
}

@Nullable
public static Filter buildFilter(
@Nullable List<FacetFilterInput> andFilters,
@Nullable List<AndFilterInput> orFilters,
@Nullable AspectRetriever aspectRetriever) {
@Nullable List<FacetFilterInput> andFilters, @Nullable List<AndFilterInput> orFilters) {
if ((andFilters == null || andFilters.isEmpty())
&& (orFilters == null || orFilters.isEmpty())) {
return null;
Expand All @@ -150,34 +141,21 @@ 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, aspectRetriever));
return new Filter().setOr(buildConjunctiveCriterionArrayWithOr(orFilters));
}

// If or filters are not set, someone may be using the legacy and filters
final List<Criterion> andCriterions = criterionListFromAndFilter(andFilters, aspectRetriever);
final List<Criterion> andCriterions = criterionListFromAndFilter(andFilters);
return new Filter()
.setOr(
new ConjunctiveCriterionArray(
new ConjunctiveCriterion().setAnd(new CriterionArray(andCriterions))));
}

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,
@Nullable AspectRetriever aspectRetriever) {
public static Criterion criterionFromFilter(final FacetFilterInput filter) {
Criterion result = new Criterion();

if (skipKeywordSuffix) {
result.setField(filter.getField());
} else {
result.setField(getFilterField(filter.getField(), skipKeywordSuffix, aspectRetriever));
}
result.setField(filter.getField());

// `value` is deprecated in place of `values`- this is to support old query patterns. If values
// is provided,
Expand Down Expand Up @@ -210,16 +188,6 @@ public static Criterion criterionFromFilter(
return result;
}

private static String getFilterField(
final String originalField,
final boolean skipKeywordSuffix,
@Nullable AspectRetriever aspectRetriever) {
if (KEYWORD_EXCLUDED_FILTERS.contains(originalField)) {
return originalField;
}
return ESUtils.toKeywordField(originalField, skipKeywordSuffix, aspectRetriever);
}

public static Filter viewFilter(
OperationContext opContext, ViewService viewService, String viewUrn) {
if (viewUrn == null) {
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -12,6 +12,7 @@
import com.linkedin.datahub.graphql.generated.AssertionRunStatus;
import com.linkedin.datahub.graphql.generated.FacetFilterInput;
import com.linkedin.datahub.graphql.generated.FilterInput;
import com.linkedin.datahub.graphql.resolvers.ResolverUtils;
import com.linkedin.datahub.graphql.types.dataset.mappers.AssertionRunEventMapper;
import com.linkedin.entity.client.EntityClient;
import com.linkedin.metadata.Constants;
Expand Down Expand Up @@ -147,7 +148,7 @@ public static Filter buildFilter(
.setAnd(
new CriterionArray(
facetFilters.stream()
.map(filter -> criterionFromFilter(filter, true, aspectRetriever))
.map(ResolverUtils::criterionFromFilter)
.collect(Collectors.toList())))));
}
}
Original file line number Diff line number Diff line change
Expand Up @@ -71,10 +71,7 @@ public CompletableFuture<ListAccessTokenResult> get(DataFetchingEnvironment envi
.withSearchFlags(flags -> flags.setFulltext(true)),
Constants.ACCESS_TOKEN_ENTITY_NAME,
"",
buildFilter(
filters,
Collections.emptyList(),
context.getOperationContext().getAspectRetriever()),
buildFilter(filters, Collections.emptyList()),
sortCriteria,
start,
count);
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -74,9 +74,7 @@ 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(), context.getOperationContext().getAspectRetriever());
final Filter inputFilter = ResolverUtils.buildFilter(null, input.getOrFilters());

BrowseResultV2 browseResults =
_entityClient.browseV2(
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -151,11 +151,7 @@ public CompletableFuture<SearchResults> get(DataFetchingEnvironment environment)
filters.removeIf(f -> f.getField().equals(OUTPUT_PORTS_FILTER_FIELD));
}
// add urns from the aspect to our filters
final Filter baseFilter =
ResolverUtils.buildFilter(
filters,
input.getOrFilters(),
context.getOperationContext().getAspectRetriever());
final Filter baseFilter = ResolverUtils.buildFilter(filters, input.getOrFilters());
final Filter finalFilter =
buildFilterWithUrns(
context.getDataHubAppConfig(), new HashSet<>(urnsToFilterOn), baseFilter);
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -79,9 +79,7 @@ public CompletableFuture<SearchResults> get(final DataFetchingEnvironment enviro
.getFilters()
.forEach(
filter -> {
criteria.add(
criterionFromFilter(
filter, true, context.getOperationContext().getAspectRetriever()));
criteria.add(criterionFromFilter(filter));
});
}

Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -33,9 +33,7 @@ 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, context.getOperationContext().getAspectRetriever());
final DynamicFormAssignment formAssignment = FormUtils.mapDynamicFormAssignment(input);

return GraphQLConcurrencyUtils.supplyAsync(
() -> {
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -68,10 +68,7 @@ public CompletableFuture<ListIngestionSourcesResult> get(
.withSearchFlags(flags -> flags.setFulltext(true)),
Constants.INGESTION_SOURCE_ENTITY_NAME,
query,
buildFilter(
filters,
Collections.emptyList(),
context.getOperationContext().getAspectRetriever()),
buildFilter(filters, Collections.emptyList()),
null,
start,
count);
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -9,9 +9,9 @@
import com.linkedin.datahub.graphql.generated.Entity;
import com.linkedin.datahub.graphql.generated.FilterInput;
import com.linkedin.datahub.graphql.generated.TimeSeriesAspect;
import com.linkedin.datahub.graphql.resolvers.ResolverUtils;
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;
Expand Down Expand Up @@ -120,7 +120,7 @@ public CompletableFuture<List<TimeSeriesAspect>> get(DataFetchingEnvironment env
maybeStartTimeMillis,
maybeEndTimeMillis,
maybeLimit,
buildFilters(maybeFilters, context.getOperationContext().getAspectRetriever()),
buildFilters(maybeFilters),
maybeSort);

// Step 2: Bind profiles into GraphQL strong types.
Expand All @@ -135,8 +135,7 @@ public CompletableFuture<List<TimeSeriesAspect>> get(DataFetchingEnvironment env
"get");
}

private Filter buildFilters(
@Nullable FilterInput maybeFilters, @Nullable AspectRetriever aspectRetriever) {
private Filter buildFilters(@Nullable FilterInput maybeFilters) {
if (maybeFilters == null) {
return null;
}
Expand All @@ -147,7 +146,7 @@ private Filter buildFilters(
.setAnd(
new CriterionArray(
maybeFilters.getAnd().stream()
.map(filter -> criterionFromFilter(filter, true, aspectRetriever))
.map(ResolverUtils::criterionFromFilter)
.collect(Collectors.toList())))));
}
}
Original file line number Diff line number Diff line change
Expand Up @@ -18,7 +18,6 @@
import com.linkedin.form.FormPromptType;
import com.linkedin.form.FormType;
import com.linkedin.form.StructuredPropertyParams;
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;
Expand All @@ -31,7 +30,6 @@
import java.util.UUID;
import java.util.stream.Collectors;
import javax.annotation.Nonnull;
import javax.annotation.Nullable;

public class FormUtils {

Expand Down Expand Up @@ -61,16 +59,13 @@ public static PrimitivePropertyValueArray getStructuredPropertyValuesFromInput(
/** Map a GraphQL CreateDynamicFormAssignmentInput to the GMS DynamicFormAssignment aspect */
@Nonnull
public static DynamicFormAssignment mapDynamicFormAssignment(
@Nonnull final CreateDynamicFormAssignmentInput input,
@Nullable AspectRetriever aspectRetriever) {
@Nonnull final CreateDynamicFormAssignmentInput input) {
Objects.requireNonNull(input, "input must not be null");

final DynamicFormAssignment result = new DynamicFormAssignment();
final Filter filter =
new Filter()
.setOr(
ResolverUtils.buildConjunctiveCriterionArrayWithOr(
input.getOrFilters(), aspectRetriever));
.setOr(ResolverUtils.buildConjunctiveCriterionArrayWithOr(input.getOrFilters()));
result.setFilter(filter);
return result;
}
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -63,10 +63,7 @@ public CompletableFuture<ListOwnershipTypesResult> get(DataFetchingEnvironment e
context.getOperationContext().withSearchFlags(flags -> flags.setFulltext(true)),
Constants.OWNERSHIP_TYPE_ENTITY_NAME,
query,
buildFilter(
filters,
Collections.emptyList(),
context.getOperationContext().getAspectRetriever()),
buildFilter(filters, Collections.emptyList()),
Collections.singletonList(DEFAULT_SORT_CRITERION),
start,
count);
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -59,11 +59,7 @@ 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(),
context.getOperationContext().getAspectRetriever());
final Filter filter = ResolverUtils.buildFilter(facetFilters, Collections.emptyList());

return _policyFetcher
.fetchPolicies(context.getOperationContext(), start, query, count, filter)
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -15,7 +15,6 @@
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;
Expand Down Expand Up @@ -74,7 +73,7 @@ public CompletableFuture<ListQueriesResult> get(final DataFetchingEnvironment en
flags -> flags.setFulltext(true).setSkipHighlighting(true)),
QUERY_ENTITY_NAME,
query,
buildFilters(input, context.getOperationContext().getAspectRetriever()),
buildFilters(input),
sortCriteria,
start,
count);
Expand Down Expand Up @@ -111,8 +110,7 @@ private List<QueryEntity> mapUnresolvedQueries(final List<Urn> queryUrns) {
}

@Nullable
private Filter buildFilters(
@Nonnull final ListQueriesInput input, @Nullable AspectRetriever aspectRetriever) {
private Filter buildFilters(@Nonnull final ListQueriesInput input) {
final AndFilterInput criteria = new AndFilterInput();
List<FacetFilterInput> andConditions = new ArrayList<>();

Expand All @@ -139,6 +137,6 @@ private Filter buildFilters(
}

criteria.setAnd(andConditions);
return buildFilter(Collections.emptyList(), ImmutableList.of(criteria), aspectRetriever);
return buildFilter(Collections.emptyList(), ImmutableList.of(criteria));
}
}
Original file line number Diff line number Diff line change
Expand Up @@ -17,6 +17,7 @@
import com.linkedin.datahub.graphql.generated.RecommendationRenderType;
import com.linkedin.datahub.graphql.generated.RecommendationRequestContext;
import com.linkedin.datahub.graphql.generated.SearchParams;
import com.linkedin.datahub.graphql.resolvers.ResolverUtils;
import com.linkedin.datahub.graphql.types.common.mappers.UrnToEntityMapper;
import com.linkedin.datahub.graphql.types.entitytype.EntityTypeMapper;
import com.linkedin.metadata.query.filter.CriterionArray;
Expand Down Expand Up @@ -105,9 +106,7 @@ private com.linkedin.metadata.recommendation.RecommendationRequestContext mapReq
searchRequestContext.setFilters(
new CriterionArray(
requestContext.getSearchRequestContext().getFilters().stream()
.map(
facetField ->
criterionFromFilter(facetField, opContext.getAspectRetriever()))
.map(ResolverUtils::criterionFromFilter)
.collect(Collectors.toList())));
}
mappedRequestContext.setSearchRequestContext(searchRequestContext);
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -64,9 +64,7 @@ public CompletableFuture<AggregateResults> get(DataFetchingEnvironment environme
UrnUtils.getUrn(input.getViewUrn()))
: null;

final Filter inputFilter =
ResolverUtils.buildFilter(
null, input.getOrFilters(), context.getOperationContext().getAspectRetriever());
final Filter inputFilter = ResolverUtils.buildFilter(null, input.getOrFilters());

final SearchFlags searchFlags = mapInputFlags(context, input.getSearchFlags());

Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -51,11 +51,7 @@ public CompletableFuture<AutoCompleteResults> get(DataFetchingEnvironment enviro
throw new ValidationException("'query' parameter can not be null or empty");
}

final Filter filter =
ResolverUtils.buildFilter(
input.getFilters(),
input.getOrFilters(),
context.getOperationContext().getRetrieverContext().orElseThrow().getAspectRetriever());
final Filter filter = ResolverUtils.buildFilter(input.getFilters(), input.getOrFilters());
final int limit = input.getLimit() != null ? input.getLimit() : DEFAULT_LIMIT;
return GraphQLConcurrencyUtils.supplyAsync(
() -> {
Expand Down
Loading

0 comments on commit 6a165a8

Please sign in to comment.