fix(customSearch): apply query string interpolation to function score

datahub-project · Sep 17, 2024 · 5ba731a · 5ba731a
1 parent da958b7
commit 5ba731a
Show file tree

Hide file tree

Showing 4 changed files with 30 additions and 13 deletions.
diff --git a/.../com/linkedin/metadata/search/elasticsearch/query/request/AutocompleteRequestHandler.java b/.../com/linkedin/metadata/search/elasticsearch/query/request/AutocompleteRequestHandler.java
@@ -140,10 +140,15 @@ public SearchRequest getSearchRequest(
                         opContext.getObjectMapper(),
                         cac,
                         queryWithDefaultFilters,
-                        customQueryConfig))
+                        customQueryConfig,
+                        input))
             .orElse(
                 SearchQueryBuilder.buildScoreFunctions(
-                    opContext, customQueryConfig, List.of(entitySpec), queryWithDefaultFilters));
+                    opContext,
+                    customQueryConfig,
+                    List.of(entitySpec),
+                    input,
+                    queryWithDefaultFilters));
     searchSourceBuilder.query(functionScoreQueryBuilder);
 
     ESUtils.buildSortOrder(searchSourceBuilder, null, List.of(entitySpec));

diff --git a/...java/com/linkedin/metadata/search/elasticsearch/query/request/CustomizedQueryHandler.java b/...java/com/linkedin/metadata/search/elasticsearch/query/request/CustomizedQueryHandler.java
@@ -119,16 +119,18 @@ private static BoolQueryBuilder toBoolQueryBuilder(
   public static FunctionScoreQueryBuilder functionScoreQueryBuilder(
       @Nonnull ObjectMapper objectMapper,
       @Nonnull QueryConfiguration customQueryConfiguration,
-      QueryBuilder queryBuilder) {
+      QueryBuilder queryBuilder,
+      String query) {
     return toFunctionScoreQueryBuilder(
-        objectMapper, queryBuilder, customQueryConfiguration.getFunctionScore());
+        objectMapper, queryBuilder, customQueryConfiguration.getFunctionScore(), query);
   }
 
   public static Optional<FunctionScoreQueryBuilder> functionScoreQueryBuilder(
       @Nonnull ObjectMapper objectMapper,
       @Nonnull AutocompleteConfiguration customAutocompleteConfiguration,
       QueryBuilder queryBuilder,
-      @Nullable QueryConfiguration customQueryConfiguration) {
+      @Nullable QueryConfiguration customQueryConfiguration,
+      String query) {
 
     Optional<FunctionScoreQueryBuilder> result = Optional.empty();
 
@@ -143,14 +145,17 @@ public static Optional<FunctionScoreQueryBuilder> functionScoreQueryBuilder(
       result =
           Optional.of(
               toFunctionScoreQueryBuilder(
-                  objectMapper, queryBuilder, customQueryConfiguration.getFunctionScore()));
+                  objectMapper, queryBuilder, customQueryConfiguration.getFunctionScore(), query));
     } else if (customAutocompleteConfiguration.getFunctionScore() != null
         && !customAutocompleteConfiguration.getFunctionScore().isEmpty()) {
       log.debug("Applying custom autocomplete function scores.");
       result =
           Optional.of(
               toFunctionScoreQueryBuilder(
-                  objectMapper, queryBuilder, customAutocompleteConfiguration.getFunctionScore()));
+                  objectMapper,
+                  queryBuilder,
+                  customAutocompleteConfiguration.getFunctionScore(),
+                  query));
     }
 
     return result;
@@ -159,7 +164,8 @@ public static Optional<FunctionScoreQueryBuilder> functionScoreQueryBuilder(
   private static FunctionScoreQueryBuilder toFunctionScoreQueryBuilder(
       @Nonnull ObjectMapper objectMapper,
       @Nonnull QueryBuilder queryBuilder,
-      @Nonnull Map<String, Object> params) {
+      @Nonnull Map<String, Object> params,
+      String query) {
     try {
       HashMap<String, Object> body = new HashMap<>(params);
       if (!body.isEmpty()) {
@@ -168,7 +174,12 @@ private static FunctionScoreQueryBuilder toFunctionScoreQueryBuilder(
 
       body.put("query", objectMapper.readValue(queryBuilder.toString(), Map.class));
 
-      String jsonFragment = objectMapper.writeValueAsString(Map.of("function_score", body));
+      String jsonFragment =
+          objectMapper
+              .writeValueAsString(Map.of("function_score", body))
+              .replace("\"{{query_string}}\"", objectMapper.writeValueAsString(query))
+              .replace(
+                  "\"{{unquoted_query_string}}\"", objectMapper.writeValueAsString(unquote(query)));
       XContentParser parser =
           XContentType.JSON
               .xContent()

diff --git a/...ain/java/com/linkedin/metadata/search/elasticsearch/query/request/SearchQueryBuilder.java b/...ain/java/com/linkedin/metadata/search/elasticsearch/query/request/SearchQueryBuilder.java
@@ -78,7 +78,7 @@ public QueryBuilder buildQuery(
 
     final QueryBuilder queryBuilder =
         buildInternalQuery(opContext, customQueryConfig, entitySpecs, query, fulltext);
-    return buildScoreFunctions(opContext, customQueryConfig, entitySpecs, queryBuilder);
+    return buildScoreFunctions(opContext, customQueryConfig, entitySpecs, query, queryBuilder);
   }
 
   /**
@@ -485,12 +485,13 @@ static FunctionScoreQueryBuilder buildScoreFunctions(
       @Nonnull OperationContext opContext,
       @Nullable QueryConfiguration customQueryConfig,
       @Nonnull List<EntitySpec> entitySpecs,
+      String query,
       @Nonnull QueryBuilder queryBuilder) {
 
     if (customQueryConfig != null) {
       // Prefer configuration function scoring over annotation scoring
       return CustomizedQueryHandler.functionScoreQueryBuilder(
-          opContext.getObjectMapper(), customQueryConfig, queryBuilder);
+          opContext.getObjectMapper(), customQueryConfig, queryBuilder, query);
     } else {
       return QueryBuilders.functionScoreQuery(
               queryBuilder, buildAnnotationScoreFunctions(entitySpecs))

diff --git a/.../src/test/java/com/linkedin/metadata/search/query/request/CustomizedQueryHandlerTest.java b/.../src/test/java/com/linkedin/metadata/search/query/request/CustomizedQueryHandlerTest.java
@@ -179,7 +179,7 @@ public void functionScoreQueryBuilderTest() {
      */
     FunctionScoreQueryBuilder selectStarTest =
         CustomizedQueryHandler.functionScoreQueryBuilder(
-            new ObjectMapper(), test.lookupQueryConfig("*").get(), inputQuery);
+            new ObjectMapper(), test.lookupQueryConfig("*").get(), inputQuery, "*");
 
     FunctionScoreQueryBuilder.FilterFunctionBuilder[] expectedSelectStarScoreFunctions = {
       new FunctionScoreQueryBuilder.FilterFunctionBuilder(
@@ -203,7 +203,7 @@ public void functionScoreQueryBuilderTest() {
      */
     FunctionScoreQueryBuilder defaultTest =
         CustomizedQueryHandler.functionScoreQueryBuilder(
-            new ObjectMapper(), test.lookupQueryConfig("foobar").get(), inputQuery);
+            new ObjectMapper(), test.lookupQueryConfig("foobar").get(), inputQuery, "foobar");
 
     FunctionScoreQueryBuilder.FilterFunctionBuilder[] expectedDefaultScoreFunctions = {
       new FunctionScoreQueryBuilder.FilterFunctionBuilder(