diff --git a/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/types/common/mappers/SearchFlagsInputMapper.java b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/types/common/mappers/SearchFlagsInputMapper.java index e6b75f9482..9f5025ccf3 100644 --- a/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/types/common/mappers/SearchFlagsInputMapper.java +++ b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/types/common/mappers/SearchFlagsInputMapper.java @@ -1,5 +1,6 @@ package com.linkedin.datahub.graphql.types.common.mappers; +import com.linkedin.data.template.StringArray; import com.linkedin.datahub.graphql.QueryContext; import com.linkedin.datahub.graphql.generated.SearchFlags; import com.linkedin.datahub.graphql.types.mappers.ModelMapper; @@ -64,6 +65,10 @@ public class SearchFlagsInputMapper .map(c -> GroupingCriterionInputMapper.map(context, c)) .collect(Collectors.toList())))); } + if (searchFlags.getCustomHighlightingFields() != null) { + result.setCustomHighlightingFields( + new StringArray(searchFlags.getCustomHighlightingFields())); + } return result; } } diff --git a/datahub-graphql-core/src/main/resources/search.graphql b/datahub-graphql-core/src/main/resources/search.graphql index 84e81e9096..ed214ebe66 100644 --- a/datahub-graphql-core/src/main/resources/search.graphql +++ b/datahub-graphql-core/src/main/resources/search.graphql @@ -162,6 +162,11 @@ input SearchFlags { Whether to include restricted entities """ includeRestricted: Boolean + + """ + fields to include for custom Highlighting + """ + customHighlightingFields: [String!] } """ 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 6e4210de6e..163bae6416 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 @@ -49,6 +49,7 @@ import java.util.stream.Stream; import javax.annotation.Nonnull; import javax.annotation.Nullable; import lombok.extern.slf4j.Slf4j; +import org.apache.commons.collections.CollectionUtils; import org.opensearch.action.search.SearchRequest; import org.opensearch.action.search.SearchResponse; import org.opensearch.common.unit.TimeValue; @@ -211,8 +212,14 @@ public class SearchRequestHandler { .forEach(searchSourceBuilder::aggregation); } if (Boolean.FALSE.equals(searchFlags.isSkipHighlighting())) { - searchSourceBuilder.highlighter(highlights); + if (CollectionUtils.isNotEmpty(searchFlags.getCustomHighlightingFields())) { + searchSourceBuilder.highlighter( + getValidatedHighlighter(searchFlags.getCustomHighlightingFields())); + } else { + searchSourceBuilder.highlighter(highlights); + } } + ESUtils.buildSortOrder(searchSourceBuilder, sortCriteria, entitySpecs); if (Boolean.TRUE.equals(searchFlags.isGetSuggestions())) { @@ -556,4 +563,16 @@ 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; + } } 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 1cd9a27446..2f7120e1f0 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 @@ -55,6 +55,14 @@ public class SearchRequestHandlerTest extends AbstractTestNGSpringContextTests { private OperationContext operationContext; public static SearchConfiguration testQueryConfig; + public static List validHighlightingFields = List.of("urn", "foreignKey"); + public static StringArray customHighlightFields = + new StringArray( + List.of( + validHighlightingFields.get(0), + validHighlightingFields.get(1), + "notExistingField", + "")); static { testQueryConfig = new SearchConfiguration(); @@ -102,6 +110,32 @@ public class SearchRequestHandlerTest extends AbstractTestNGSpringContextTests { "unexpected lineage fields in highlights: " + highlightFields); } + @Test + public void testCustomHighlights() { + EntitySpec entitySpec = operationContext.getEntityRegistry().getEntitySpec("dataset"); + SearchRequestHandler requestHandler = + SearchRequestHandler.getBuilder(TestEntitySpecBuilder.getSpec(), testQueryConfig, null); + SearchRequest searchRequest = + requestHandler.getSearchRequest( + operationContext.withSearchFlags( + flags -> + flags.setFulltext(false).setCustomHighlightingFields(customHighlightFields)), + "testQuery", + null, + null, + 0, + 10, + null); + SearchSourceBuilder sourceBuilder = searchRequest.source(); + assertNotNull(sourceBuilder.highlighter()); + assertEquals(4, sourceBuilder.highlighter().fields().size()); + assertTrue( + sourceBuilder.highlighter().fields().stream() + .map(HighlightBuilder.Field::name) + .toList() + .containsAll(validHighlightingFields)); + } + @Test public void testSearchRequestHandlerHighlightingTurnedOff() { SearchRequestHandler requestHandler = diff --git a/metadata-models/src/main/pegasus/com/linkedin/metadata/query/SearchFlags.pdl b/metadata-models/src/main/pegasus/com/linkedin/metadata/query/SearchFlags.pdl index 355a8bb7a5..0561a9c6f7 100644 --- a/metadata-models/src/main/pegasus/com/linkedin/metadata/query/SearchFlags.pdl +++ b/metadata-models/src/main/pegasus/com/linkedin/metadata/query/SearchFlags.pdl @@ -48,4 +48,9 @@ record SearchFlags { * include restricted entities in results (default is to filter) */ includeRestricted:optional boolean = false + + /** + * Include mentioned fields inside elastic highlighting query + */ + customHighlightingFields:optional array[string] } diff --git a/metadata-service/restli-api/src/main/snapshot/com.linkedin.entity.entities.snapshot.json b/metadata-service/restli-api/src/main/snapshot/com.linkedin.entity.entities.snapshot.json index 59894ed083..204793886b 100644 --- a/metadata-service/restli-api/src/main/snapshot/com.linkedin.entity.entities.snapshot.json +++ b/metadata-service/restli-api/src/main/snapshot/com.linkedin.entity.entities.snapshot.json @@ -6032,6 +6032,14 @@ "doc" : "include restricted entities in results (default is to filter)", "default" : false, "optional" : true + }, { + "name" : "customHighlightingFields", + "type" : { + "type" : "array", + "items" : "string" + }, + "doc" : "Include mentioned fields inside elastic highlighting query", + "optional" : true } ] }, { "type" : "enum",