2323import io .datahubproject .metadata .context .OperationContext ;
2424import java .net .URISyntaxException ;
2525import java .util .ArrayList ;
26+ import java .util .Collection ;
2627import java .util .Collections ;
2728import java .util .HashSet ;
2829import java .util .List ;
4546import org .opensearch .search .fetch .subphase .highlight .HighlightBuilder ;
4647
4748@ Slf4j
48- public class AutocompleteRequestHandler {
49+ public class AutocompleteRequestHandler extends BaseRequestHandler {
4950
50- private final List <Pair > _defaultAutocompleteFields ;
51+ private final List <Pair < String , String > > _defaultAutocompleteFields ;
5152 private final Map <String , Set <SearchableAnnotation .FieldType >> searchableFieldTypes ;
5253
5354 private static final Map <EntitySpec , AutocompleteRequestHandler >
@@ -58,8 +59,10 @@ public class AutocompleteRequestHandler {
5859 private final EntitySpec entitySpec ;
5960 private final QueryFilterRewriteChain queryFilterRewriteChain ;
6061 private final SearchConfiguration searchConfiguration ;
62+ @ Nonnull private final HighlightBuilder highlights ;
6163
6264 public AutocompleteRequestHandler (
65+ @ Nonnull OperationContext systemOperationContext ,
6366 @ Nonnull EntitySpec entitySpec ,
6467 @ Nullable CustomSearchConfiguration customSearchConfiguration ,
6568 @ Nonnull QueryFilterRewriteChain queryFilterRewriteChain ,
@@ -79,6 +82,7 @@ public AutocompleteRequestHandler(
7982 Double .toString (searchableAnnotation .getBoostScore ()))),
8083 Stream .of (Pair .of ("urn" , "1.0" )))
8184 .collect (Collectors .toList ());
85+ this .highlights = getDefaultHighlights (systemOperationContext );
8286 searchableFieldTypes =
8387 fieldSpecs .stream ()
8488 .collect (
@@ -98,6 +102,7 @@ public AutocompleteRequestHandler(
98102 }
99103
100104 public static AutocompleteRequestHandler getBuilder (
105+ @ Nonnull OperationContext systemOperationContext ,
101106 @ Nonnull EntitySpec entitySpec ,
102107 @ Nullable CustomSearchConfiguration customSearchConfiguration ,
103108 @ Nonnull QueryFilterRewriteChain queryFilterRewriteChain ,
@@ -106,6 +111,7 @@ public static AutocompleteRequestHandler getBuilder(
106111 entitySpec ,
107112 k ->
108113 new AutocompleteRequestHandler (
114+ systemOperationContext ,
109115 entitySpec ,
110116 customSearchConfiguration ,
111117 queryFilterRewriteChain ,
@@ -165,7 +171,8 @@ public SearchRequest getSearchRequest(
165171 ESUtils .buildSortOrder (searchSourceBuilder , null , List .of (entitySpec ));
166172
167173 // wire inner non-scored query
168- searchSourceBuilder .highlighter (getHighlights (field ));
174+ searchSourceBuilder .highlighter (
175+ field == null || field .isEmpty () ? highlights : getHighlights (opContext , List .of (field )));
169176 searchRequest .source (searchSourceBuilder );
170177 return searchRequest ;
171178 }
@@ -181,7 +188,7 @@ private BoolQueryBuilder getQuery(
181188 public BoolQueryBuilder getQuery (
182189 @ Nonnull ObjectMapper objectMapper ,
183190 @ Nullable AutocompleteConfiguration customAutocompleteConfig ,
184- List <Pair > autocompleteFields ,
191+ List <Pair < String , String > > autocompleteFields ,
185192 @ Nonnull String query ) {
186193
187194 BoolQueryBuilder finalQuery =
@@ -201,7 +208,7 @@ public BoolQueryBuilder getQuery(
201208
202209 private Optional <QueryBuilder > getAutocompleteQuery (
203210 @ Nullable AutocompleteConfiguration customConfig ,
204- List <Pair > autocompleteFields ,
211+ List <Pair < String , String > > autocompleteFields ,
205212 @ Nonnull String query ) {
206213 Optional <QueryBuilder > result = Optional .empty ();
207214
@@ -212,7 +219,8 @@ private Optional<QueryBuilder> getAutocompleteQuery(
212219 return result ;
213220 }
214221
215- private BoolQueryBuilder defaultQuery (List <Pair > autocompleteFields , @ Nonnull String query ) {
222+ private BoolQueryBuilder defaultQuery (
223+ List <Pair <String , String >> autocompleteFields , @ Nonnull String query ) {
216224 BoolQueryBuilder finalQuery = QueryBuilders .boolQuery ().minimumShouldMatch (1 );
217225
218226 // Search for exact matches with higher boost and ngram matches
@@ -248,38 +256,25 @@ private BoolQueryBuilder defaultQuery(List<Pair> autocompleteFields, @Nonnull St
248256 return finalQuery ;
249257 }
250258
251- // Get HighlightBuilder to highlight the matched field
252- private HighlightBuilder getHighlights (@ Nullable String field ) {
253- HighlightBuilder highlightBuilder =
254- new HighlightBuilder ()
255- // Don't set tags to get the original field value
256- .preTags ("" )
257- .postTags ("" )
258- .numOfFragments (1 );
259- // Check for each field name and any subfields
260- getAutocompleteFields (field )
261- .forEach (
262- pair -> {
263- final String fieldName = (String ) pair .getLeft ();
264- highlightBuilder
265- .field (fieldName )
266- .field (fieldName + ".*" )
267- .field (fieldName + ".ngram" )
268- .field (fieldName + ".delimited" );
269- if (!fieldName .equalsIgnoreCase ("urn" )) {
270- highlightBuilder .field (fieldName + ".keyword" );
271- }
272- });
273-
274- // set field match req false for ngram
275- highlightBuilder .fields ().stream ()
276- .filter (f -> f .name ().contains ("ngram" ))
277- .forEach (f -> f .requireFieldMatch (false ).noMatchSize (200 ));
278-
279- return highlightBuilder ;
259+ @ Override
260+ public Collection <String > getDefaultQueryFieldNames () {
261+ return _defaultAutocompleteFields .stream ().map (Pair ::getKey ).collect (Collectors .toList ());
280262 }
281263
282- private List <Pair > getAutocompleteFields (@ Nullable String field ) {
264+ @ Override
265+ protected Collection <String > getValidQueryFieldNames () {
266+ return searchableFieldTypes .keySet ();
267+ }
268+
269+ @ Override
270+ protected Stream <String > highlightFieldExpansion (
271+ @ Nonnull OperationContext opContext , @ Nonnull String fieldName ) {
272+ return Stream .concat (
273+ Stream .of (fieldName , fieldName + ".*" , fieldName + ".ngram" , fieldName + ".delimited" ),
274+ Stream .of (ESUtils .toKeywordField (fieldName , false , opContext .getAspectRetriever ())));
275+ }
276+
277+ private List <Pair <String , String >> getAutocompleteFields (@ Nullable String field ) {
283278 if (field != null && !field .isEmpty () && !field .equalsIgnoreCase ("urn" )) {
284279 return ImmutableList .of (Pair .of (field , "10.0" ));
285280 }
0 commit comments