diff --git a/metadata-io/src/main/java/com/linkedin/metadata/search/elasticsearch/query/ESBrowseDAO.java b/metadata-io/src/main/java/com/linkedin/metadata/search/elasticsearch/query/ESBrowseDAO.java index 3a84d1cb2d..1c20d5a6ef 100644 --- a/metadata-io/src/main/java/com/linkedin/metadata/search/elasticsearch/query/ESBrowseDAO.java +++ b/metadata-io/src/main/java/com/linkedin/metadata/search/elasticsearch/query/ESBrowseDAO.java @@ -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, diff --git a/metadata-io/src/main/java/com/linkedin/metadata/search/elasticsearch/query/ESSearchDAO.java b/metadata-io/src/main/java/com/linkedin/metadata/search/elasticsearch/query/ESSearchDAO.java index a720e0bf81..ae483ad771 100644 --- a/metadata-io/src/main/java/com/linkedin/metadata/search/elasticsearch/query/ESSearchDAO.java +++ b/metadata-io/src/main/java/com/linkedin/metadata/search/elasticsearch/query/ESSearchDAO.java @@ -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, diff --git a/metadata-io/src/main/java/com/linkedin/metadata/search/elasticsearch/query/request/AutocompleteRequestHandler.java b/metadata-io/src/main/java/com/linkedin/metadata/search/elasticsearch/query/request/AutocompleteRequestHandler.java index 45359285b4..d755d5c6f7 100644 --- a/metadata-io/src/main/java/com/linkedin/metadata/search/elasticsearch/query/request/AutocompleteRequestHandler.java +++ b/metadata-io/src/main/java/com/linkedin/metadata/search/elasticsearch/query/request/AutocompleteRequestHandler.java @@ -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 _defaultAutocompleteFields; + private final List> _defaultAutocompleteFields; private final Map> searchableFieldTypes; private static final Map @@ -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 autocompleteFields, + List> autocompleteFields, @Nonnull String query) { BoolQueryBuilder finalQuery = @@ -201,7 +208,7 @@ public class AutocompleteRequestHandler { private Optional getAutocompleteQuery( @Nullable AutocompleteConfiguration customConfig, - List autocompleteFields, + List> autocompleteFields, @Nonnull String query) { Optional result = Optional.empty(); @@ -212,7 +219,8 @@ public class AutocompleteRequestHandler { return result; } - private BoolQueryBuilder defaultQuery(List autocompleteFields, @Nonnull String query) { + private BoolQueryBuilder defaultQuery( + List> 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 getDefaultQueryFieldNames() { + return _defaultAutocompleteFields.stream().map(Pair::getKey).collect(Collectors.toList()); } - private List getAutocompleteFields(@Nullable String field) { + @Override + protected Collection getValidQueryFieldNames() { + return searchableFieldTypes.keySet(); + } + + @Override + protected Stream 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> getAutocompleteFields(@Nullable String field) { if (field != null && !field.isEmpty() && !field.equalsIgnoreCase("urn")) { return ImmutableList.of(Pair.of(field, "10.0")); } diff --git a/metadata-io/src/main/java/com/linkedin/metadata/search/elasticsearch/query/request/BaseRequestHandler.java b/metadata-io/src/main/java/com/linkedin/metadata/search/elasticsearch/query/request/BaseRequestHandler.java new file mode 100644 index 0000000000..ec43122107 --- /dev/null +++ b/metadata-io/src/main/java/com/linkedin/metadata/search/elasticsearch/query/request/BaseRequestHandler.java @@ -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 getDefaultQueryFieldNames(); + + protected abstract Collection getValidQueryFieldNames(); + + protected abstract Stream 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 fieldsToHighlight) { + HighlightBuilder highlightBuilder = + new HighlightBuilder() + // Don't set tags to get the original field value + .preTags("") + .postTags("") + .numOfFragments(1); + + final Stream 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; + } +} diff --git a/metadata-io/src/main/java/com/linkedin/metadata/search/elasticsearch/query/request/SearchRequestHandler.java b/metadata-io/src/main/java/com/linkedin/metadata/search/elasticsearch/query/request/SearchRequestHandler.java index 0ecf1a932e..90ab6cd597 100644 --- a/metadata-io/src/main/java/com/linkedin/metadata/search/elasticsearch/query/request/SearchRequestHandler.java +++ b/metadata-io/src/main/java/com/linkedin/metadata/search/elasticsearch/query/request/SearchRequestHandler.java @@ -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, 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 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 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 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 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 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 * diff --git a/metadata-io/src/test/java/com/linkedin/metadata/search/query/request/AutocompleteRequestHandlerTest.java b/metadata-io/src/test/java/com/linkedin/metadata/search/query/request/AutocompleteRequestHandlerTest.java index e91dc1e6a8..4055736926 100644 --- a/metadata-io/src/test/java/com/linkedin/metadata/search/query/request/AutocompleteRequestHandlerTest.java +++ b/metadata-io/src/test/java/com/linkedin/metadata/search/query/request/AutocompleteRequestHandlerTest.java @@ -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 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 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, diff --git a/metadata-io/src/test/java/com/linkedin/metadata/search/query/request/SearchRequestHandlerTest.java b/metadata-io/src/test/java/com/linkedin/metadata/search/query/request/SearchRequestHandlerTest.java index 62b9572d3c..f1e56e6350 100644 --- a/metadata-io/src/test/java/com/linkedin/metadata/search/query/request/SearchRequestHandlerTest.java +++ b/metadata-io/src/test/java/com/linkedin/metadata/search/query/request/SearchRequestHandlerTest.java @@ -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 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 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 diff --git a/metadata-service/servlet/src/main/java/com/datahub/gms/servlet/ConfigSearchExport.java b/metadata-service/servlet/src/main/java/com/datahub/gms/servlet/ConfigSearchExport.java index f80581bf6f..720845e7ac 100644 --- a/metadata-service/servlet/src/main/java/com/datahub/gms/servlet/ConfigSearchExport.java +++ b/metadata-service/servlet/src/main/java/com/datahub/gms/servlet/ConfigSearchExport.java @@ -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,