add empty or null check

Signed-off-by: Mingshi Liu <mingshl@amazon.com>
opensearch-project · Jan 23, 2025 · fcc44d5 · fcc44d5
1 parent 6b6a6de
commit fcc44d5
Show file tree

Hide file tree

Showing 2 changed files with 99 additions and 8 deletions.
diff --git a/server/src/main/java/org/opensearch/index/query/TemplateQueryBuilder.java b/server/src/main/java/org/opensearch/index/query/TemplateQueryBuilder.java
@@ -133,14 +133,16 @@ protected QueryBuilder doRewrite(QueryRewriteContext queryCoordinatorContext) th
         }
 
         // Convert Map<String, Object> to Map<String, String> with proper JSON escaping
-        Map<String, String> variablesMap = contextVariables.entrySet().stream().collect(Collectors.toMap(Map.Entry::getKey, entry -> {
-            try {
-                return JsonXContent.contentBuilder().value(entry.getValue()).toString();
-            } catch (IOException e) {
-                throw new RuntimeException("Error converting contextVariables to JSON string", e);
-            }
-        }));
-
+        Map<String, String> variablesMap = null;
+        if (contextVariables != null) {
+            variablesMap = contextVariables.entrySet().stream().collect(Collectors.toMap(Map.Entry::getKey, entry -> {
+                try {
+                    return JsonXContent.contentBuilder().value(entry.getValue()).toString();
+                } catch (IOException e) {
+                    throw new RuntimeException("Error converting contextVariables to JSON string", e);
+                }
+            }));
+        }
         String newQueryContent = replaceVariables(queryString, variablesMap);
 
         try {
@@ -159,6 +161,16 @@ protected QueryBuilder doRewrite(QueryRewriteContext queryCoordinatorContext) th
     }
 
     private String replaceVariables(String template, Map<String, String> variables) {
+        if (template == null || template.equals("null")) {
+            throw new IllegalArgumentException("Template string cannot be null. A valid template must be provided.");
+        }
+        if (template.isEmpty() || template.equals("{}")) {
+            throw new IllegalArgumentException("Template string cannot be empty. A valid template must be provided.");
+        }
+        if (variables == null || variables.isEmpty()) {
+            return template;
+        }
+
         StringBuilder result = new StringBuilder();
         int start = 0;
         while (true) {

diff --git a/server/src/test/java/org/opensearch/index/query/TemplateQueryBuilderTests.java b/server/src/test/java/org/opensearch/index/query/TemplateQueryBuilderTests.java
@@ -727,6 +727,85 @@ public void testDoRewriteWithMissingBracketVariable() throws IOException {
         assertTrue(exception.getMessage().contains("Unclosed variable in template"));
     }
 
+    /**
+     * Tests the replaceVariables method when the template is null.
+     * Verifies that an IllegalArgumentException is thrown with the appropriate error message.
+     */
+
+    public void testReplaceVariablesWithNullTemplate() {
+        TemplateQueryBuilder templateQueryBuilder = new TemplateQueryBuilder((Map<String, Object>) null);
+
+        QueryCoordinatorContext queryRewriteContext = mockQueryRewriteContext();
+        Map<String, Object> contextVariables = new HashMap<>();
+        contextVariables.put("response", "foo");
+        when(queryRewriteContext.getContextVariables()).thenReturn(contextVariables);
+
+        IllegalArgumentException exception = expectThrows(
+            IllegalArgumentException.class,
+            () -> templateQueryBuilder.doRewrite(queryRewriteContext)
+        );
+        assertEquals("Template string cannot be null. A valid template must be provided.", exception.getMessage());
+    }
+
+    /**
+     * Tests the replaceVariables method when the template is empty.
+     * Verifies that an IllegalArgumentException is thrown with the appropriate error message.
+     */
+
+    public void testReplaceVariablesWithEmptyTemplate() {
+        Map<String, Object> template = new HashMap<>();
+        TemplateQueryBuilder templateQueryBuilder = new TemplateQueryBuilder(template);
+
+        QueryCoordinatorContext queryRewriteContext = mockQueryRewriteContext();
+        Map<String, Object> contextVariables = new HashMap<>();
+        contextVariables.put("response", "foo");
+        when(queryRewriteContext.getContextVariables()).thenReturn(contextVariables);
+
+        IllegalArgumentException exception = expectThrows(
+            IllegalArgumentException.class,
+            () -> templateQueryBuilder.doRewrite(queryRewriteContext)
+        );
+        assertEquals("Template string cannot be empty. A valid template must be provided.", exception.getMessage());
+
+    }
+
+    /**
+     * Tests the replaceVariables method when the variables map is null.
+     * Verifies that the method returns the original template unchanged,
+     * since a null variables map is treated as no replacement.
+     */
+    public void testReplaceVariablesWithNullVariables() throws IOException {
+
+        Map<String, Object> template = new HashMap<>();
+        Map<String, Object> term = new HashMap<>();
+        Map<String, Object> message = new HashMap<>();
+
+        message.put("value", "foo");
+        term.put("message", message);
+        template.put("term", term);
+        TemplateQueryBuilder templateQueryBuilder = new TemplateQueryBuilder(template);
+        TermQueryBuilder termQueryBuilder = new TermQueryBuilder("message", "foo");
+
+        QueryCoordinatorContext queryRewriteContext = mockQueryRewriteContext();
+
+        when(queryRewriteContext.getContextVariables()).thenReturn(null);
+
+        TermQueryBuilder newQuery = (TermQueryBuilder) templateQueryBuilder.doRewrite(queryRewriteContext);
+
+        assertEquals(newQuery, termQueryBuilder);
+        assertEquals(
+            "{\n"
+                + "  \"term\" : {\n"
+                + "    \"message\" : {\n"
+                + "      \"value\" : \"foo\",\n"
+                + "      \"boost\" : 1.0\n"
+                + "    }\n"
+                + "  }\n"
+                + "}",
+            newQuery.toString()
+        );
+    }
+
     /**
      * Helper method to create a mock QueryCoordinatorContext for testing.
      */