yujunjun/datahub
2
0
Fork 0
mirror of https://github.com/datahub-project/datahub.git synced 2025-12-03 14:16:28 +00:00
Code Issues Packages Projects Releases Wiki Activity

fix(autocomplete): fix autocomplete duplicate field (#12558)

Browse Source
This commit is contained in:
david-leifker 2025-02-05 20:04:53 -06:00 committed by GitHub
parent 65376ee2d3
commit 52f71dd111
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
8 changed files with 173 additions and 132 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

4
metadata-io/src/main/java/com/linkedin/metadata/search/elasticsearch/query/ESBrowseDAO.java
View File

@ -642,7 +642,7 @@ public class ESBrowseDAO {
EntitySpec entitySpec = opContext.getEntityRegistry().getEntitySpec(entityName);
QueryBuilder query =
SearchRequestHandler.getBuilder(
opContext.getEntityRegistry(),
opContext,
entitySpec,
searchConfiguration,
customSearchConfiguration,
@ -683,7 +683,7 @@ public class ESBrowseDAO {
QueryBuilder query =
SearchRequestHandler.getBuilder(
finalOpContext.getEntityRegistry(),
finalOpContext,
entitySpecs,
searchConfiguration,
customSearchConfiguration,

18
metadata-io/src/main/java/com/linkedin/metadata/search/elasticsearch/query/ESSearchDAO.java
View File

@ -150,7 +150,7 @@ public class ESSearchDAO {
return transformIndexIntoEntityName(
opContext.getSearchContext().getIndexConvention(),
SearchRequestHandler.getBuilder(
opContext.getEntityRegistry(),
opContext,
entitySpec,
searchConfiguration,
customSearchConfiguration,
@ -257,7 +257,7 @@ public class ESSearchDAO {
return transformIndexIntoEntityName(
opContext.getSearchContext().getIndexConvention(),
SearchRequestHandler.getBuilder(
opContext.getEntityRegistry(),
opContext,
entitySpecs,
searchConfiguration,
customSearchConfiguration,
@ -311,7 +311,7 @@ public class ESSearchDAO {
"searchRequest",
() ->
SearchRequestHandler.getBuilder(
opContext.getEntityRegistry(),
opContext,
entitySpecs,
searchConfiguration,
customSearchConfiguration,
@ -357,7 +357,7 @@ public class ESSearchDAO {
Filter transformedFilters = transformFilterForEntities(filters, indexConvention);
final SearchRequest searchRequest =
SearchRequestHandler.getBuilder(
opContext.getEntityRegistry(),
opContext,
entitySpec,
searchConfiguration,
customSearchConfiguration,
@ -395,7 +395,11 @@ public class ESSearchDAO {
IndexConvention indexConvention = opContext.getSearchContext().getIndexConvention();
AutocompleteRequestHandler builder =
AutocompleteRequestHandler.getBuilder(
entitySpec, customSearchConfiguration, queryFilterRewriteChain, searchConfiguration);
opContext,
entitySpec,
customSearchConfiguration,
queryFilterRewriteChain,
searchConfiguration);
SearchRequest req =
builder.getSearchRequest(
opContext,
@ -441,7 +445,7 @@ public class ESSearchDAO {
IndexConvention indexConvention = opContext.getSearchContext().getIndexConvention();
final SearchRequest searchRequest =
SearchRequestHandler.getBuilder(
opContext.getEntityRegistry(),
opContext,
entitySpecs,
searchConfiguration,
customSearchConfiguration,
@ -578,7 +582,7 @@ public class ESSearchDAO {
}
return SearchRequestHandler.getBuilder(
opContext.getEntityRegistry(),
opContext,
entitySpecs,
searchConfiguration,
customSearchConfiguration,

67
metadata-io/src/main/java/com/linkedin/metadata/search/elasticsearch/query/request/AutocompleteRequestHandler.java
View File

@ -23,6 +23,7 @@ import com.linkedin.metadata.search.utils.ESUtils;
import io.datahubproject.metadata.context.OperationContext;
import java.net.URISyntaxException;
import java.util.ArrayList;
import java.util.Collection;
import java.util.Collections;
import java.util.HashSet;
import java.util.List;
@ -45,9 +46,9 @@ import org.opensearch.search.builder.SearchSourceBuilder;
import org.opensearch.search.fetch.subphase.highlight.HighlightBuilder;
@Slf4j
public class AutocompleteRequestHandler {
public class AutocompleteRequestHandler extends BaseRequestHandler {
private final List<Pair> _defaultAutocompleteFields;
private final List<Pair<String, String>> _defaultAutocompleteFields;
private final Map<String, Set<SearchableAnnotation.FieldType>> searchableFieldTypes;
private static final Map<EntitySpec, AutocompleteRequestHandler>
@ -58,8 +59,10 @@ public class AutocompleteRequestHandler {
private final EntitySpec entitySpec;
private final QueryFilterRewriteChain queryFilterRewriteChain;
private final SearchConfiguration searchConfiguration;
@Nonnull private final HighlightBuilder highlights;
public AutocompleteRequestHandler(
@Nonnull OperationContext systemOperationContext,
@Nonnull EntitySpec entitySpec,
@Nullable CustomSearchConfiguration customSearchConfiguration,
@Nonnull QueryFilterRewriteChain queryFilterRewriteChain,
@ -79,6 +82,7 @@ public class AutocompleteRequestHandler {
Double.toString(searchableAnnotation.getBoostScore()))),
Stream.of(Pair.of("urn", "1.0")))
.collect(Collectors.toList());
this.highlights = getDefaultHighlights(systemOperationContext);
searchableFieldTypes =
fieldSpecs.stream()
.collect(
@ -98,6 +102,7 @@ public class AutocompleteRequestHandler {
}
public static AutocompleteRequestHandler getBuilder(
@Nonnull OperationContext systemOperationContext,
@Nonnull EntitySpec entitySpec,
@Nullable CustomSearchConfiguration customSearchConfiguration,
@Nonnull QueryFilterRewriteChain queryFilterRewriteChain,
@ -106,6 +111,7 @@ public class AutocompleteRequestHandler {
entitySpec,
k ->
new AutocompleteRequestHandler(
systemOperationContext,
entitySpec,
customSearchConfiguration,
queryFilterRewriteChain,
@ -165,7 +171,8 @@ public class AutocompleteRequestHandler {
ESUtils.buildSortOrder(searchSourceBuilder, null, List.of(entitySpec));
// wire inner non-scored query
searchSourceBuilder.highlighter(getHighlights(field));
searchSourceBuilder.highlighter(
field == null || field.isEmpty() ? highlights : getHighlights(opContext, List.of(field)));
searchRequest.source(searchSourceBuilder);
return searchRequest;
}
@ -181,7 +188,7 @@ public class AutocompleteRequestHandler {
public BoolQueryBuilder getQuery(
@Nonnull ObjectMapper objectMapper,
@Nullable AutocompleteConfiguration customAutocompleteConfig,
List<Pair> autocompleteFields,
List<Pair<String, String>> autocompleteFields,
@Nonnull String query) {
BoolQueryBuilder finalQuery =
@ -201,7 +208,7 @@ public class AutocompleteRequestHandler {
private Optional<QueryBuilder> getAutocompleteQuery(
@Nullable AutocompleteConfiguration customConfig,
List<Pair> autocompleteFields,
List<Pair<String, String>> autocompleteFields,
@Nonnull String query) {
Optional<QueryBuilder> result = Optional.empty();
@ -212,7 +219,8 @@ public class AutocompleteRequestHandler {
return result;
}
private BoolQueryBuilder defaultQuery(List<Pair> autocompleteFields, @Nonnull String query) {
private BoolQueryBuilder defaultQuery(
List<Pair<String, String>> autocompleteFields, @Nonnull String query) {
BoolQueryBuilder finalQuery = QueryBuilders.boolQuery().minimumShouldMatch(1);
// Search for exact matches with higher boost and ngram matches
@ -248,38 +256,25 @@ public class AutocompleteRequestHandler {
return finalQuery;
}
// Get HighlightBuilder to highlight the matched field
private HighlightBuilder getHighlights(@Nullable String field) {
HighlightBuilder highlightBuilder =
new HighlightBuilder()
// Don't set tags to get the original field value
.preTags("")
.postTags("")
.numOfFragments(1);
// Check for each field name and any subfields
getAutocompleteFields(field)
.forEach(
pair -> {
final String fieldName = (String) pair.getLeft();
highlightBuilder
.field(fieldName)
.field(fieldName + ".*")
.field(fieldName + ".ngram")
.field(fieldName + ".delimited");
if (!fieldName.equalsIgnoreCase("urn")) {
highlightBuilder.field(fieldName + ".keyword");
}
});
// set field match req false for ngram
highlightBuilder.fields().stream()
.filter(f -> f.name().contains("ngram"))
.forEach(f -> f.requireFieldMatch(false).noMatchSize(200));
return highlightBuilder;
@Override
public Collection<String> getDefaultQueryFieldNames() {
return _defaultAutocompleteFields.stream().map(Pair::getKey).collect(Collectors.toList());
}
private List<Pair> getAutocompleteFields(@Nullable String field) {
@Override
protected Collection<String> getValidQueryFieldNames() {
return searchableFieldTypes.keySet();
}
@Override
protected Stream<String> highlightFieldExpansion(
@Nonnull OperationContext opContext, @Nonnull String fieldName) {
return Stream.concat(
Stream.of(fieldName, fieldName + ".*", fieldName + ".ngram", fieldName + ".delimited"),
Stream.of(ESUtils.toKeywordField(fieldName, false, opContext.getAspectRetriever())));
}
private List<Pair<String, String>> getAutocompleteFields(@Nullable String field) {
if (field != null && !field.isEmpty() && !field.equalsIgnoreCase("urn")) {
return ImmutableList.of(Pair.of(field, "10.0"));
}

69
metadata-io/src/main/java/com/linkedin/metadata/search/elasticsearch/query/request/BaseRequestHandler.java Normal file
View File

@ -0,0 +1,69 @@
package com.linkedin.metadata.search.elasticsearch.query.request;
import com.google.common.annotations.VisibleForTesting;
import io.datahubproject.metadata.context.OperationContext;
import java.util.Collection;
import java.util.Objects;
import java.util.stream.Stream;
import javax.annotation.Nonnull;
import javax.annotation.Nullable;
import org.opensearch.search.fetch.subphase.highlight.HighlightBuilder;
public abstract class BaseRequestHandler {
/**
* Provide the fields which are queried by default
*
* @return collection of field names
*/
protected abstract Collection<String> getDefaultQueryFieldNames();
protected abstract Collection<String> getValidQueryFieldNames();
protected abstract Stream<String> highlightFieldExpansion(
@Nonnull OperationContext opContext, @Nonnull String fieldName);
@VisibleForTesting
public HighlightBuilder getDefaultHighlights(@Nonnull OperationContext opContext) {
return getHighlights(opContext, null);
}
@VisibleForTesting
public HighlightBuilder getHighlights(
@Nonnull OperationContext opContext, @Nullable Collection<String> fieldsToHighlight) {
HighlightBuilder highlightBuilder =
new HighlightBuilder()
// Don't set tags to get the original field value
.preTags("")
.postTags("")
.numOfFragments(1);
final Stream<String> fieldStream;
if (fieldsToHighlight == null || fieldsToHighlight.isEmpty()) {
fieldStream = getDefaultQueryFieldNames().stream();
} else {
// filter for valid names
fieldStream =
fieldsToHighlight.stream()
.filter(Objects::nonNull)
.filter(fieldName -> !fieldName.isEmpty())
.filter(getValidQueryFieldNames()::contains);
}
fieldStream
.flatMap(fieldName -> highlightFieldExpansion(opContext, fieldName))
.distinct()
.map(HighlightBuilder.Field::new)
.map(
field -> {
if (field.name().endsWith("ngram")) {
return field.requireFieldMatch(false).noMatchSize(200);
} else {
return field;
}
})
.forEach(highlightBuilder::field);
return highlightBuilder;
}
}

60
metadata-io/src/main/java/com/linkedin/metadata/search/elasticsearch/query/request/SearchRequestHandler.java
View File

@ -78,7 +78,7 @@ import org.opensearch.search.fetch.subphase.highlight.HighlightField;
import org.opensearch.search.suggest.term.TermSuggestion;
@Slf4j
public class SearchRequestHandler {
public class SearchRequestHandler extends BaseRequestHandler {
private static final Map<List<EntitySpec>, SearchRequestHandler> REQUEST_HANDLER_BY_ENTITY_NAME =
new ConcurrentHashMap<>();
@ -94,13 +94,13 @@ public class SearchRequestHandler {
private final QueryFilterRewriteChain queryFilterRewriteChain;
private SearchRequestHandler(
@Nonnull EntityRegistry entityRegistry,
@Nonnull OperationContext opContext,
@Nonnull EntitySpec entitySpec,
@Nonnull SearchConfiguration configs,
@Nullable CustomSearchConfiguration customSearchConfiguration,
@Nonnull QueryFilterRewriteChain queryFilterRewriteChain) {
this(
entityRegistry,
opContext,
ImmutableList.of(entitySpec),
configs,
customSearchConfiguration,
@ -108,7 +108,7 @@ public class SearchRequestHandler {
}
private SearchRequestHandler(
@Nonnull EntityRegistry entityRegistry,
@Nonnull OperationContext opContext,
@Nonnull List<EntitySpec> entitySpecs,
@Nonnull SearchConfiguration configs,
@Nullable CustomSearchConfiguration customSearchConfiguration,
@ -121,16 +121,17 @@ public class SearchRequestHandler {
.flatMap(List::stream)
.collect(Collectors.toList());
defaultQueryFieldNames = getDefaultQueryFieldNames(annotations);
highlights = getHighlights();
highlights = getDefaultHighlights(opContext);
searchQueryBuilder = new SearchQueryBuilder(configs, customSearchConfiguration);
aggregationQueryBuilder = new AggregationQueryBuilder(configs, entitySearchAnnotations);
this.configs = configs;
this.searchableFieldTypes = buildSearchableFieldTypes(entityRegistry, entitySpecs);
this.searchableFieldTypes =
buildSearchableFieldTypes(opContext.getEntityRegistry(), entitySpecs);
this.queryFilterRewriteChain = queryFilterRewriteChain;
}
public static SearchRequestHandler getBuilder(
@Nonnull EntityRegistry entityRegistry,
@Nonnull OperationContext systemOperationContext,
@Nonnull EntitySpec entitySpec,
@Nonnull SearchConfiguration configs,
@Nullable CustomSearchConfiguration customSearchConfiguration,
@ -139,7 +140,7 @@ public class SearchRequestHandler {
ImmutableList.of(entitySpec),
k ->
new SearchRequestHandler(
entityRegistry,
systemOperationContext,
entitySpec,
configs,
customSearchConfiguration,
@ -147,7 +148,7 @@ public class SearchRequestHandler {
}
public static SearchRequestHandler getBuilder(
@Nonnull EntityRegistry entityRegistry,
@Nonnull OperationContext systemOperationContext,
@Nonnull List<EntitySpec> entitySpecs,
@Nonnull SearchConfiguration configs,
@Nullable CustomSearchConfiguration customSearchConfiguration,
@ -156,7 +157,7 @@ public class SearchRequestHandler {
ImmutableList.copyOf(entitySpecs),
k ->
new SearchRequestHandler(
entityRegistry,
systemOperationContext,
entitySpecs,
configs,
customSearchConfiguration,
@ -185,6 +186,11 @@ public class SearchRequestHandler {
.collect(Collectors.toSet());
}
@Override
protected Collection<String> getValidQueryFieldNames() {
return searchableFieldTypes.keySet();
}
public BoolQueryBuilder getFilterQuery(
@Nonnull OperationContext opContext, @Nullable Filter filter) {
return getFilterQuery(opContext, filter, searchableFieldTypes, queryFilterRewriteChain);
@ -246,7 +252,7 @@ public class SearchRequestHandler {
if (Boolean.FALSE.equals(searchFlags.isSkipHighlighting())) {
if (CollectionUtils.isNotEmpty(searchFlags.getCustomHighlightingFields())) {
searchSourceBuilder.highlighter(
getValidatedHighlighter(searchFlags.getCustomHighlightingFields()));
getHighlights(opContext, searchFlags.getCustomHighlightingFields()));
} else {
searchSourceBuilder.highlighter(highlights);
}
@ -383,22 +389,10 @@ public class SearchRequestHandler {
return searchQueryBuilder.buildQuery(opContext, entitySpecs, query, fulltext);
}
@VisibleForTesting
public HighlightBuilder getHighlights() {
HighlightBuilder highlightBuilder =
new HighlightBuilder()
// Don't set tags to get the original field value
.preTags("")
.postTags("")
.numOfFragments(1);
// Check for each field name and any subfields
defaultQueryFieldNames.stream()
.flatMap(fieldName -> Stream.of(fieldName, fieldName + ".*"))
.distinct()
.forEach(highlightBuilder::field);
return highlightBuilder;
@Override
protected Stream<String> highlightFieldExpansion(
@Nonnull OperationContext opContext, @Nonnull String fieldName) {
return Stream.of(fieldName, fieldName + ".*");
}
@WithSpan
@ -597,18 +591,6 @@ public class SearchRequestHandler {
return searchSuggestions;
}
private HighlightBuilder getValidatedHighlighter(Collection<String> fieldsToHighlight) {
HighlightBuilder highlightBuilder = new HighlightBuilder();
highlightBuilder.preTags("");
highlightBuilder.postTags("");
fieldsToHighlight.stream()
.filter(defaultQueryFieldNames::contains)
.flatMap(fieldName -> Stream.of(fieldName, fieldName + ".*"))
.distinct()
.forEach(highlightBuilder::field);
return highlightBuilder;
}
/**
* Calculate the field types based on annotations if available, with fallback to ES mappings
*

40
metadata-io/src/test/java/com/linkedin/metadata/search/query/request/AutocompleteRequestHandlerTest.java
View File

@ -49,11 +49,12 @@ import org.opensearch.index.query.functionscore.FunctionScoreQueryBuilder;
import org.opensearch.index.query.functionscore.ScoreFunctionBuilders;
import org.opensearch.search.builder.SearchSourceBuilder;
import org.opensearch.search.fetch.subphase.highlight.HighlightBuilder;
import org.testng.annotations.BeforeClass;
import org.testng.annotations.Test;
public class AutocompleteRequestHandlerTest {
private static SearchConfiguration testQueryConfig;
private static AutocompleteRequestHandler handler;
private AutocompleteRequestHandler handler;
private OperationContext mockOpContext =
TestOperationContexts.systemContextNoSearchAuthorization(mock(EntityRegistry.class));
private OperationContext nonMockOpContext =
@ -83,9 +84,13 @@ public class AutocompleteRequestHandlerTest {
testQueryConfig.setExactMatch(exactMatchConfiguration);
testQueryConfig.setWordGram(wordGramConfiguration);
testQueryConfig.setPartial(partialConfiguration);
}
@BeforeClass
public void beforeTest() {
handler =
AutocompleteRequestHandler.getBuilder(
mockOpContext,
TestEntitySpecBuilder.getSpec(),
CustomSearchConfiguration.builder().build(),
QueryFilterRewriteChain.EMPTY,
@ -177,9 +182,9 @@ public class AutocompleteRequestHandlerTest {
@Test
public void testAutocompleteRequestWithField() {
// When field is null
// The field must be a valid field in the model. Pick from `keyPart1` or `urn`
SearchRequest autocompleteRequest =
handler.getSearchRequest(mockOpContext, "input", "field", null, 10);
handler.getSearchRequest(mockOpContext, "input", "keyPart1", null, 10);
SearchSourceBuilder sourceBuilder = autocompleteRequest.source();
assertEquals(sourceBuilder.size(), 10);
BoolQueryBuilder wrapper =
@ -189,19 +194,19 @@ public class AutocompleteRequestHandlerTest {
assertEquals(query.should().size(), 3);
MatchQueryBuilder matchQueryBuilder = (MatchQueryBuilder) query.should().get(0);
assertEquals("field.keyword", matchQueryBuilder.fieldName());
assertEquals("keyPart1.keyword", matchQueryBuilder.fieldName());
MultiMatchQueryBuilder autocompleteQuery = (MultiMatchQueryBuilder) query.should().get(2);
Map<String, Float> queryFields = autocompleteQuery.fields();
assertTrue(queryFields.containsKey("field.ngram"));
assertTrue(queryFields.containsKey("field.ngram._2gram"));
assertTrue(queryFields.containsKey("field.ngram._3gram"));
assertTrue(queryFields.containsKey("field.ngram._4gram"));
assertTrue(queryFields.containsKey("keyPart1.ngram"));
assertTrue(queryFields.containsKey("keyPart1.ngram._2gram"));
assertTrue(queryFields.containsKey("keyPart1.ngram._3gram"));
assertTrue(queryFields.containsKey("keyPart1.ngram._4gram"));
assertEquals(autocompleteQuery.type(), MultiMatchQueryBuilder.Type.BOOL_PREFIX);
MatchPhrasePrefixQueryBuilder prefixQuery =
(MatchPhrasePrefixQueryBuilder) query.should().get(1);
assertEquals("field.delimited", prefixQuery.fieldName());
assertEquals("keyPart1.delimited", prefixQuery.fieldName());
TermQueryBuilder removedFilter = (TermQueryBuilder) wrapper.mustNot().get(0);
assertEquals(removedFilter.fieldName(), "removed");
@ -209,11 +214,11 @@ public class AutocompleteRequestHandlerTest {
HighlightBuilder highlightBuilder = sourceBuilder.highlighter();
List<HighlightBuilder.Field> highlightedFields = highlightBuilder.fields();
assertEquals(highlightedFields.size(), 5);
assertEquals(highlightedFields.get(0).name(), "field");
assertEquals(highlightedFields.get(1).name(), "field.*");
assertEquals(highlightedFields.get(2).name(), "field.ngram");
assertEquals(highlightedFields.get(3).name(), "field.delimited");
assertEquals(highlightedFields.get(4).name(), "field.keyword");
assertEquals(highlightedFields.get(0).name(), "keyPart1");
assertEquals(highlightedFields.get(1).name(), "keyPart1.*");
assertEquals(highlightedFields.get(2).name(), "keyPart1.ngram");
assertEquals(highlightedFields.get(3).name(), "keyPart1.delimited");
assertEquals(highlightedFields.get(4).name(), "keyPart1.keyword");
}
@Test
@ -221,6 +226,7 @@ public class AutocompleteRequestHandlerTest {
// Exclude Default query
AutocompleteRequestHandler withoutDefaultQuery =
AutocompleteRequestHandler.getBuilder(
mockOpContext,
TestEntitySpecBuilder.getSpec(),
CustomSearchConfiguration.builder()
.autocompleteConfigurations(
@ -248,6 +254,7 @@ public class AutocompleteRequestHandlerTest {
// Include Default query
AutocompleteRequestHandler withDefaultQuery =
AutocompleteRequestHandler.getBuilder(
mockOpContext,
TestEntitySpecBuilder.getSpec(),
CustomSearchConfiguration.builder()
.autocompleteConfigurations(
@ -290,6 +297,7 @@ public class AutocompleteRequestHandlerTest {
// Pickup scoring functions from non-autocomplete
AutocompleteRequestHandler withInherit =
AutocompleteRequestHandler.getBuilder(
mockOpContext,
TestEntitySpecBuilder.getSpec(),
CustomSearchConfiguration.builder()
.queryConfigurations(List.of(TEST_QUERY_CONFIG))
@ -332,6 +340,7 @@ public class AutocompleteRequestHandlerTest {
// no search query customization
AutocompleteRequestHandler noQueryCustomization =
AutocompleteRequestHandler.getBuilder(
mockOpContext,
TestEntitySpecBuilder.getSpec(),
CustomSearchConfiguration.builder()
.autocompleteConfigurations(
@ -378,6 +387,7 @@ public class AutocompleteRequestHandlerTest {
// Scoring functions explicit autocomplete override
AutocompleteRequestHandler explicitNoInherit =
AutocompleteRequestHandler.getBuilder(
mockOpContext,
TestEntitySpecBuilder.getSpec(),
CustomSearchConfiguration.builder()
.queryConfigurations(List.of(TEST_QUERY_CONFIG)) // should be ignored
@ -432,6 +442,7 @@ public class AutocompleteRequestHandlerTest {
// inherit enabled)
AutocompleteRequestHandler explicit =
AutocompleteRequestHandler.getBuilder(
mockOpContext,
TestEntitySpecBuilder.getSpec(),
CustomSearchConfiguration.builder()
.queryConfigurations(List.of(TEST_QUERY_CONFIG)) // should be ignored
@ -605,6 +616,7 @@ public class AutocompleteRequestHandlerTest {
AutocompleteRequestHandler requestHandler =
AutocompleteRequestHandler.getBuilder(
mockOpContext,
entitySpec,
CustomSearchConfiguration.builder().build(),
QueryFilterRewriteChain.EMPTY,

36
metadata-io/src/test/java/com/linkedin/metadata/search/query/request/SearchRequestHandlerTest.java
View File

@ -119,17 +119,13 @@ public class SearchRequestHandlerTest extends AbstractTestNGSpringContextTests {
EntitySpec entitySpec = operationContext.getEntityRegistry().getEntitySpec("dataset");
SearchRequestHandler datasetHandler =
SearchRequestHandler.getBuilder(
operationContext.getEntityRegistry(),
entitySpec,
testQueryConfig,
null,
QueryFilterRewriteChain.EMPTY);
operationContext, entitySpec, testQueryConfig, null, QueryFilterRewriteChain.EMPTY);
/*
Ensure efficient query performance, we do not expect upstream/downstream/fineGrained lineage
*/
List<String> highlightFields =
datasetHandler.getHighlights().fields().stream()
datasetHandler.getDefaultHighlights(operationContext).fields().stream()
.map(HighlightBuilder.Field::name)
.collect(Collectors.toList());
assertTrue(
@ -144,7 +140,7 @@ public class SearchRequestHandlerTest extends AbstractTestNGSpringContextTests {
EntitySpec entitySpec = operationContext.getEntityRegistry().getEntitySpec("dataset");
SearchRequestHandler requestHandler =
SearchRequestHandler.getBuilder(
operationContext.getEntityRegistry(),
operationContext,
TestEntitySpecBuilder.getSpec(),
testQueryConfig,
null,
@ -174,7 +170,7 @@ public class SearchRequestHandlerTest extends AbstractTestNGSpringContextTests {
public void testSearchRequestHandlerHighlightingTurnedOff() {
SearchRequestHandler requestHandler =
SearchRequestHandler.getBuilder(
operationContext.getEntityRegistry(),
operationContext,
TestEntitySpecBuilder.getSpec(),
testQueryConfig,
null,
@ -220,7 +216,7 @@ public class SearchRequestHandlerTest extends AbstractTestNGSpringContextTests {
public void testSearchRequestHandler() {
SearchRequestHandler requestHandler =
SearchRequestHandler.getBuilder(
operationContext.getEntityRegistry(),
operationContext,
TestEntitySpecBuilder.getSpec(),
testQueryConfig,
null,
@ -294,7 +290,7 @@ public class SearchRequestHandlerTest extends AbstractTestNGSpringContextTests {
public void testAggregationsInSearch() {
SearchRequestHandler requestHandler =
SearchRequestHandler.getBuilder(
operationContext.getEntityRegistry(),
operationContext,
TestEntitySpecBuilder.getSpec(),
testQueryConfig,
null,
@ -368,7 +364,7 @@ public class SearchRequestHandlerTest extends AbstractTestNGSpringContextTests {
final SearchRequestHandler requestHandler =
SearchRequestHandler.getBuilder(
operationContext.getEntityRegistry(),
operationContext,
TestEntitySpecBuilder.getSpec(),
testQueryConfig,
null,
@ -793,11 +789,7 @@ public class SearchRequestHandlerTest extends AbstractTestNGSpringContextTests {
operationContext.getEntityRegistry().getEntitySpec(EntityTypeMapper.getName(entityType));
SearchRequestHandler handler =
SearchRequestHandler.getBuilder(
operationContext.getEntityRegistry(),
entitySpec,
testQueryConfig,
null,
QueryFilterRewriteChain.EMPTY);
operationContext, entitySpec, testQueryConfig, null, QueryFilterRewriteChain.EMPTY);
Set<String> unexpected = new HashSet<>(handler.getDefaultQueryFieldNames());
unexpected.removeAll(expectedEntityQueryByDefault);
@ -1027,11 +1019,7 @@ public class SearchRequestHandlerTest extends AbstractTestNGSpringContextTests {
final SearchRequestHandler requestHandler =
SearchRequestHandler.getBuilder(
operationContext.getEntityRegistry(),
entitySpec,
testQueryConfig,
null,
QueryFilterRewriteChain.EMPTY);
operationContext, entitySpec, testQueryConfig, null, QueryFilterRewriteChain.EMPTY);
return (BoolQueryBuilder)
requestHandler
@ -1059,11 +1047,7 @@ public class SearchRequestHandlerTest extends AbstractTestNGSpringContextTests {
final SearchRequestHandler requestHandler =
SearchRequestHandler.getBuilder(
operationContext.getEntityRegistry(),
entitySpec,
testQueryConfig,
null,
QueryFilterRewriteChain.EMPTY);
operationContext, entitySpec, testQueryConfig, null, QueryFilterRewriteChain.EMPTY);
return (BoolQueryBuilder)
requestHandler

11
metadata-service/servlet/src/main/java/com/datahub/gms/servlet/ConfigSearchExport.java
View File

@ -6,7 +6,6 @@ import static com.linkedin.metadata.search.elasticsearch.indexbuilder.SettingsBu
import com.datahub.gms.util.CSVWriter;
import com.linkedin.datahub.graphql.types.entitytype.EntityTypeMapper;
import com.linkedin.gms.factory.config.ConfigurationProvider;
import com.linkedin.metadata.aspect.AspectRetriever;
import com.linkedin.metadata.config.search.SearchConfiguration;
import com.linkedin.metadata.models.EntitySpec;
import com.linkedin.metadata.models.registry.EntityRegistry;
@ -40,10 +39,6 @@ public class ConfigSearchExport extends HttpServlet {
return (ConfigurationProvider) ctx.getBean("configurationProvider");
}
private static AspectRetriever getAspectRetriever(WebApplicationContext ctx) {
return (AspectRetriever) ctx.getBean("aspectRetriever");
}
private static OperationContext getOperationContext(WebApplicationContext ctx) {
return (OperationContext) ctx.getBean("systemOperationContext");
}
@ -53,9 +48,9 @@ public class ConfigSearchExport extends HttpServlet {
}
private void writeSearchCsv(WebApplicationContext ctx, PrintWriter pw) {
OperationContext systemOpContext = getOperationContext(ctx);
SearchConfiguration searchConfiguration = getConfigProvider(ctx).getElasticSearch().getSearch();
AspectRetriever aspectRetriever = getAspectRetriever(ctx);
EntityRegistry entityRegistry = aspectRetriever.getEntityRegistry();
EntityRegistry entityRegistry = systemOpContext.getEntityRegistry();
QueryFilterRewriteChain queryFilterRewriteChain = getQueryFilterRewriteChain(ctx);
CSVWriter writer = CSVWriter.builder().printWriter(pw).build();
@ -92,7 +87,7 @@ public class ConfigSearchExport extends HttpServlet {
EntitySpec entitySpec = entitySpecOpt.get();
SearchRequest searchRequest =
SearchRequestHandler.getBuilder(
entityRegistry,
systemOpContext,
entitySpec,
searchConfiguration,
null,