address coverage

Signed-off-by: Lantao Jin <ltjin@amazon.com>
opensearch-project · Jul 4, 2024 · 873ecee · 873ecee
1 parent 09ca944
commit 873ecee
Show file tree

Hide file tree

Showing 4 changed files with 14 additions and 3 deletions.
diff --git a/legacy/src/main/java/org/opensearch/sql/legacy/plugin/RestSqlAction.java b/legacy/src/main/java/org/opensearch/sql/legacy/plugin/RestSqlAction.java
@@ -167,7 +167,7 @@ protected RestChannelConsumer prepareRequest(RestRequest request, NodeClient cli
               LOG.info(
                   "[{}] Request {} is not supported and falling back to old SQL engine",
                   QueryContext.getRequestId(),
-                  newSqlRequest.toAnonymousString(QueryDataAnonymizer::anonymizeData));
+                  newSqlRequest.toString(QueryDataAnonymizer::anonymizeData));
               LOG.info("Request Query: {}", QueryDataAnonymizer.anonymizeData(sqlRequest.getSql()));
               QueryAction queryAction = explainRequest(client, sqlRequest, format);
               executeSqlRequest(request, queryAction, client, restChannel);

diff --git a/legacy/src/test/java/org/opensearch/sql/legacy/unittest/utils/QueryDataAnonymizerTest.java b/legacy/src/test/java/org/opensearch/sql/legacy/unittest/utils/QueryDataAnonymizerTest.java
@@ -101,7 +101,7 @@ public void test_to_anonymous_string_of_SQLQueryRequest() {
             + "FROM accounts a JOIN employees e";
     SQLQueryRequest request =
         new SQLQueryRequest(null, query, "/_plugins/_sql", Map.of("pretty", "true"), null);
-    String actualQuery = request.toAnonymousString(QueryDataAnonymizer::anonymizeData);
+    String actualQuery = request.toString(QueryDataAnonymizer::anonymizeData);
     String expectedQuery =
         "SQLQueryRequest(query=( SELECT identifier, identifier, identifier, identifier, identifier"
             + " FROM table a JOIN table e ), path=/_plugins/_sql, format=jdbc,"

diff --git a/sql/src/main/java/org/opensearch/sql/sql/domain/SQLQueryRequest.java b/sql/src/main/java/org/opensearch/sql/sql/domain/SQLQueryRequest.java
@@ -150,7 +150,8 @@ private boolean shouldSanitize(Map<String, String> params) {
     return true;
   }
 
-  public String toAnonymousString(Function<String, String> anonymizer) {
+  /** A new toString() with anonymizer parameter to anonymize its query statement. */
+  public String toString(Function<String, String> anonymizer) {
     return "SQLQueryRequest("
         + "query="
         + anonymizer.apply(getQuery())

diff --git a/sql/src/test/java/org/opensearch/sql/sql/domain/SQLQueryRequestTest.java b/sql/src/test/java/org/opensearch/sql/sql/domain/SQLQueryRequestTest.java
@@ -222,6 +222,16 @@ public void should_support_raw_format() {
     assertTrue(csvRequest.isSupported());
   }
 
+  @Test
+  public void test_to_string_with_anonymizer() {
+    SQLQueryRequest request = SQLQueryRequestBuilder.request("SELECT 1").build();
+    String actual = request.toString(s -> "***");
+    String expected =
+        "SQLQueryRequest(query=***, path=_plugins/_sql, format=jdbc, params={}, sanitize=true,"
+            + " cursor=Optional.empty)";
+    assertEquals(expected, actual);
+  }
+
   /** SQL query request build helper to improve test data setup readability. */
   private static class SQLQueryRequestBuilder {
     private String jsonContent;