fix(browsev2): align browse and aggregate queries (#9790)

Author: RyanHolstien

datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/chart/BrowseV2Resolver.java

@@ -18,6 +18,7 @@
 import com.linkedin.datahub.graphql.types.entitytype.EntityTypeMapper;
 import com.linkedin.entity.client.EntityClient;
 import com.linkedin.metadata.browse.BrowseResultV2;
+import com.linkedin.metadata.query.SearchFlags;
 import com.linkedin.metadata.query.filter.Filter;
 import com.linkedin.metadata.service.FormService;
 import com.linkedin.metadata.service.ViewService;
@@ -52,6 +53,7 @@ public CompletableFuture<BrowseResultsV2> get(DataFetchingEnvironment environmen
     final int start = input.getStart() != null ? input.getStart() : DEFAULT_START;
     final int count = input.getCount() != null ? input.getCount() : DEFAULT_COUNT;
     final String query = input.getQuery() != null ? input.getQuery() : "*";
+    final SearchFlags searchFlags = mapInputFlags(input.getSearchFlags());
     // escape forward slash since it is a reserved character in Elasticsearch
     final String sanitizedQuery = ResolverUtils.escapeForwardSlash(query);
 
@@ -83,7 +85,8 @@ public CompletableFuture<BrowseResultsV2> get(DataFetchingEnvironment environmen
                     sanitizedQuery,
                     start,
                     count,
-                    context.getAuthentication());
+                    context.getAuthentication(),
+                    searchFlags);
             return mapBrowseResults(browseResults);
           } catch (Exception e) {
             throw new RuntimeException("Failed to execute browse V2", e);

datahub-graphql-core/src/main/resources/search.graphql

@@ -1230,6 +1230,11 @@ input BrowseV2Input {
   The search query string
   """
   query: String
+
+  """
+  Flags controlling search options
+  """
+  searchFlags: SearchFlags
 }
 
 """

datahub-graphql-core/src/test/java/com/linkedin/datahub/graphql/resolvers/browse/BrowseV2ResolverTest.java

@@ -21,6 +21,7 @@
 import com.linkedin.metadata.browse.BrowseResultGroupV2Array;
 import com.linkedin.metadata.browse.BrowseResultMetadata;
 import com.linkedin.metadata.browse.BrowseResultV2;
+import com.linkedin.metadata.query.SearchFlags;
 import com.linkedin.metadata.query.filter.ConjunctiveCriterion;
 import com.linkedin.metadata.query.filter.ConjunctiveCriterionArray;
 import com.linkedin.metadata.query.filter.Criterion;
@@ -262,7 +263,8 @@ private static EntityClient initMockEntityClient(
                 Mockito.eq(query),
                 Mockito.eq(start),
                 Mockito.eq(limit),
-                Mockito.any(Authentication.class)))
+                Mockito.any(Authentication.class),
+                Mockito.nullable(SearchFlags.class)))
         .thenReturn(result);
     return client;
   }

metadata-io/src/main/java/com/linkedin/metadata/client/JavaEntityClient.java

@@ -229,9 +229,11 @@ public BrowseResultV2 browseV2(
       @Nonnull String input,
       int start,
       int count,
-      @Nonnull Authentication authentication) {
+      @Nonnull Authentication authentication,
+      @Nullable SearchFlags searchFlags) {
     // TODO: cache browseV2 results
-    return _entitySearchService.browseV2(entityName, path, filter, input, start, count);
+    return _entitySearchService.browseV2(
+        entityName, path, filter, input, start, count, searchFlags);
   }
 
   /**
@@ -253,9 +255,11 @@ public BrowseResultV2 browseV2(
       @Nonnull String input,
       int start,
       int count,
-      @Nonnull Authentication authentication) {
+      @Nonnull Authentication authentication,
+      @Nullable SearchFlags searchFlags) {
     // TODO: cache browseV2 results
-    return _entitySearchService.browseV2(entityNames, path, filter, input, start, count);
+    return _entitySearchService.browseV2(
+        entityNames, path, filter, input, start, count, searchFlags);
   }
 
   @SneakyThrows

metadata-io/src/main/java/com/linkedin/metadata/search/elasticsearch/ElasticSearchService.java

@@ -215,8 +215,9 @@ public BrowseResultV2 browseV2(
       @Nullable Filter filter,
       @Nonnull String input,
       int start,
-      int count) {
-    return esBrowseDAO.browseV2(entityName, path, filter, input, start, count);
+      int count,
+      @Nullable SearchFlags searchFlags) {
+    return esBrowseDAO.browseV2(entityName, path, filter, input, start, count, searchFlags);
   }
 
   @Nonnull
@@ -227,8 +228,9 @@ public BrowseResultV2 browseV2(
       @Nullable Filter filter,
       @Nonnull String input,
       int start,
-      int count) {
-    return esBrowseDAO.browseV2(entityNames, path, filter, input, start, count);
+      int count,
+      @Nullable SearchFlags searchFlags) {
+    return esBrowseDAO.browseV2(entityNames, path, filter, input, start, count, searchFlags);
   }
 
   @Nonnull

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

@@ -21,6 +21,7 @@
 import com.linkedin.metadata.models.EntitySpec;
 import com.linkedin.metadata.models.annotation.SearchableAnnotation;
 import com.linkedin.metadata.models.registry.EntityRegistry;
+import com.linkedin.metadata.query.SearchFlags;
 import com.linkedin.metadata.query.filter.Filter;
 import com.linkedin.metadata.search.elasticsearch.query.request.SearchRequestHandler;
 import com.linkedin.metadata.search.utils.ESUtils;
@@ -34,6 +35,7 @@
 import java.util.Collections;
 import java.util.List;
 import java.util.Map;
+import java.util.Optional;
 import java.util.Set;
 import java.util.stream.Collectors;
 import javax.annotation.Nonnull;
@@ -399,14 +401,15 @@ public BrowseResultV2 browseV2(
       @Nullable Filter filter,
       @Nonnull String input,
       int start,
-      int count) {
+      int count,
+      @Nullable SearchFlags searchFlags) {
     try {
       final SearchResponse groupsResponse;
       try (Timer.Context ignored = MetricUtils.timer(this.getClass(), "esGroupSearch").time()) {
         final String finalInput = input.isEmpty() ? "*" : input;
         groupsResponse =
             client.search(
-                constructGroupsSearchRequestV2(entityName, path, filter, finalInput),
+                constructGroupsSearchRequestV2(entityName, path, filter, finalInput, searchFlags),
                 RequestOptions.DEFAULT);
       }
 
@@ -435,7 +438,8 @@ public BrowseResultV2 browseV2(
       @Nullable Filter filter,
       @Nonnull String input,
       int start,
-      int count) {
+      int count,
+      @Nullable SearchFlags searchFlags) {
     try {
       final SearchResponse groupsResponse;
 
@@ -444,7 +448,7 @@ public BrowseResultV2 browseV2(
         groupsResponse =
             client.search(
                 constructGroupsSearchRequestBrowseAcrossEntities(
-                    entities, path, filter, finalInput),
+                    entities, path, filter, finalInput, searchFlags),
                 RequestOptions.DEFAULT);
       }
 
@@ -472,7 +476,8 @@ private SearchRequest constructGroupsSearchRequestV2(
       @Nonnull String entityName,
       @Nonnull String path,
       @Nullable Filter filter,
-      @Nonnull String input) {
+      @Nonnull String input,
+      @Nullable SearchFlags searchFlags) {
     final String indexName = indexConvention.getIndexName(entityRegistry.getEntitySpec(entityName));
     final SearchRequest searchRequest = new SearchRequest(indexName);
     final SearchSourceBuilder searchSourceBuilder = new SearchSourceBuilder();
@@ -482,7 +487,8 @@ private SearchRequest constructGroupsSearchRequestV2(
             entityName,
             path,
             SearchUtil.transformFilterForEntities(filter, indexConvention),
-            input));
+            input,
+            searchFlags));
     searchSourceBuilder.aggregation(buildAggregationsV2(path));
     searchRequest.source(searchSourceBuilder);
     return searchRequest;
@@ -493,7 +499,8 @@ private SearchRequest constructGroupsSearchRequestBrowseAcrossEntities(
       @Nonnull List<String> entities,
       @Nonnull String path,
       @Nullable Filter filter,
-      @Nonnull String input) {
+      @Nonnull String input,
+      @Nullable SearchFlags searchFlags) {
 
     List<EntitySpec> entitySpecs =
         entities.stream().map(entityRegistry::getEntitySpec).collect(Collectors.toList());
@@ -509,7 +516,8 @@ private SearchRequest constructGroupsSearchRequestBrowseAcrossEntities(
             entitySpecs,
             path,
             SearchUtil.transformFilterForEntities(filter, indexConvention),
-            input));
+            input,
+            searchFlags));
     searchSourceBuilder.aggregation(buildAggregationsV2(path));
     searchRequest.source(searchSourceBuilder);
     return searchRequest;
@@ -537,15 +545,18 @@ private QueryBuilder buildQueryStringV2(
       @Nonnull String entityName,
       @Nonnull String path,
       @Nullable Filter filter,
-      @Nonnull String input) {
+      @Nonnull String input,
+      @Nullable SearchFlags searchFlags) {
+    SearchFlags finalSearchFlags =
+        Optional.ofNullable(searchFlags).orElse(new SearchFlags().setFulltext(true));
     final int browseDepthVal = getPathDepthV2(path);
 
     final BoolQueryBuilder queryBuilder = QueryBuilders.boolQuery();
 
     EntitySpec entitySpec = entityRegistry.getEntitySpec(entityName);
     QueryBuilder query =
         SearchRequestHandler.getBuilder(entitySpec, searchConfiguration, customSearchConfiguration)
-            .getQuery(input, false);
+            .getQuery(input, Boolean.TRUE.equals(finalSearchFlags.isFulltext()));
     queryBuilder.must(query);
 
     filterSoftDeletedByDefault(filter, queryBuilder);
@@ -567,14 +578,17 @@ private QueryBuilder buildQueryStringBrowseAcrossEntities(
       @Nonnull List<EntitySpec> entitySpecs,
       @Nonnull String path,
       @Nullable Filter filter,
-      @Nonnull String input) {
+      @Nonnull String input,
+      @Nullable SearchFlags searchFlags) {
+    SearchFlags finalSearchFlags =
+        Optional.ofNullable(searchFlags).orElse(new SearchFlags().setFulltext(true));
     final int browseDepthVal = getPathDepthV2(path);
 
     final BoolQueryBuilder queryBuilder = QueryBuilders.boolQuery();
 
     QueryBuilder query =
         SearchRequestHandler.getBuilder(entitySpecs, searchConfiguration, customSearchConfiguration)
-            .getQuery(input, false);
+            .getQuery(input, Boolean.TRUE.equals(finalSearchFlags.isFulltext()));
     queryBuilder.must(query);
 
     if (!path.isEmpty()) {

metadata-io/src/main/java/com/linkedin/metadata/search/elasticsearch/query/request/SearchQueryBuilder.java

@@ -135,14 +135,10 @@ private QueryBuilder buildInternalQuery(
           query.startsWith(STRUCTURED_QUERY_PREFIX)
               ? query.substring(STRUCTURED_QUERY_PREFIX.length())
               : query;
-
-      QueryStringQueryBuilder queryBuilder = QueryBuilders.queryStringQuery(withoutQueryPrefix);
-      queryBuilder.defaultOperator(Operator.AND);
-      getStandardFields(entitySpecs)
-          .forEach(entitySpec -> queryBuilder.field(entitySpec.fieldName(), entitySpec.boost()));
-      finalQuery.should(queryBuilder);
+      getStructuredQuery(customQueryConfig, entitySpecs, withoutQueryPrefix)
+          .ifPresent(finalQuery::should);
       if (exactMatchConfiguration.isEnableStructured()) {
-        getPrefixAndExactMatchQuery(null, entitySpecs, withoutQueryPrefix)
+        getPrefixAndExactMatchQuery(customQueryConfig, entitySpecs, withoutQueryPrefix)
             .ifPresent(finalQuery::should);
       }
     }
@@ -415,6 +411,29 @@ private Optional<QueryBuilder> getPrefixAndExactMatchQuery(
     return finalQuery.should().size() > 0 ? Optional.of(finalQuery) : Optional.empty();
   }
 
+  private Optional<QueryBuilder> getStructuredQuery(
+      @Nullable QueryConfiguration customQueryConfig,
+      List<EntitySpec> entitySpecs,
+      String sanitizedQuery) {
+    Optional<QueryBuilder> result = Optional.empty();
+
+    final boolean executeStructuredQuery;
+    if (customQueryConfig != null) {
+      executeStructuredQuery = customQueryConfig.isStructuredQuery();
+    } else {
+      executeStructuredQuery = !(isQuoted(sanitizedQuery) && exactMatchConfiguration.isExclusive());
+    }
+
+    if (executeStructuredQuery) {
+      QueryStringQueryBuilder queryBuilder = QueryBuilders.queryStringQuery(sanitizedQuery);
+      queryBuilder.defaultOperator(Operator.AND);
+      getStandardFields(entitySpecs)
+          .forEach(entitySpec -> queryBuilder.field(entitySpec.fieldName(), entitySpec.boost()));
+      result = Optional.of(queryBuilder);
+    }
+    return result;
+  }
+
   private FunctionScoreQueryBuilder buildScoreFunctions(
       @Nullable QueryConfiguration customQueryConfig,
       @Nonnull List<EntitySpec> entitySpecs,

metadata-service/configuration/src/main/java/com/linkedin/metadata/config/search/custom/QueryConfiguration.java

@@ -19,6 +19,13 @@ public class QueryConfiguration {
 
   private String queryRegex;
   @Builder.Default private boolean simpleQuery = true;
+
+  /**
+   * Used to determine if standard structured query logic should be applied when relevant, i.e.
+   * fullText flag is false. Will not be added in cases where simpleQuery would be the standard.
+   */
+  @Builder.Default private boolean structuredQuery = true;
+
   @Builder.Default private boolean exactMatchQuery = true;
   @Builder.Default private boolean prefixMatchQuery = true;
   private BoolQueryConfiguration boolQuery;

metadata-service/restli-client/src/main/java/com/linkedin/entity/client/EntityClient.java

@@ -154,7 +154,8 @@ public BrowseResultV2 browseV2(
       @Nonnull String input,
       int start,
       int count,
-      @Nonnull Authentication authentication)
+      @Nonnull Authentication authentication,
+      @Nullable SearchFlags searchFlags)
       throws RemoteInvocationException;
 
   /**
@@ -176,7 +177,8 @@ public BrowseResultV2 browseV2(
       @Nonnull String input,
       int start,
       int count,
-      @Nonnull Authentication authentication)
+      @Nonnull Authentication authentication,
+      @Nullable SearchFlags searchFlags)
       throws RemoteInvocationException;
 
   @Deprecated

metadata-service/restli-client/src/main/java/com/linkedin/entity/client/RestliEntityClient.java

@@ -378,7 +378,8 @@ public BrowseResultV2 browseV2(
       @Nonnull String input,
       int start,
       int count,
-      @Nonnull Authentication authentication) {
+      @Nonnull Authentication authentication,
+      @Nullable SearchFlags searchFlags) {
     throw new NotImplementedException("BrowseV2 is not implemented in Restli yet");
   }
 
@@ -391,7 +392,8 @@ public BrowseResultV2 browseV2(
       @Nonnull String input,
       int start,
       int count,
-      @Nonnull Authentication authentication)
+      @Nonnull Authentication authentication,
+      @Nullable SearchFlags searchFlags)
       throws RemoteInvocationException {
     throw new NotImplementedException("BrowseV2 is not implemented in Restli yet");
   }

metadata-service/services/src/main/java/com/linkedin/metadata/search/EntitySearchService.java

@@ -197,6 +197,7 @@ BrowseResult browse(
    * @param input search query
    * @param start start offset of first group
    * @param count max number of results requested
+   * @param searchFlags configuration options for search
    */
   @Nonnull
   public BrowseResultV2 browseV2(
@@ -205,7 +206,8 @@ public BrowseResultV2 browseV2(
       @Nullable Filter filter,
       @Nonnull String input,
       int start,
-      int count);
+      int count,
+      @Nullable SearchFlags searchFlags);
 
   /**
    * Gets browse snapshot of a given path
@@ -216,6 +218,7 @@ public BrowseResultV2 browseV2(
    * @param input search query
    * @param start start offset of first group
    * @param count max number of results requested
+   * @param searchFlags configuration options for search
    */
   @Nonnull
   public BrowseResultV2 browseV2(
@@ -224,7 +227,8 @@ public BrowseResultV2 browseV2(
       @Nullable Filter filter,
       @Nonnull String input,
       int start,
-      int count);
+      int count,
+      @Nullable SearchFlags searchFlags);
 
   /**
    * Gets a list of paths for a given urn.