From fd1129d3b25b03805e5b6798974405018ed02a04 Mon Sep 17 00:00:00 2001 From: Lantao Jin Date: Thu, 4 Jul 2024 19:31:14 +0800 Subject: [PATCH 1/3] Query anonymization should work in fallback log Signed-off-by: Lantao Jin --- .../sql/legacy/plugin/RestSqlAction.java | 2 +- .../utils/QueryDataAnonymizerTest.java | 18 ++++++++++++++++++ .../sql/sql/domain/SQLQueryRequest.java | 18 ++++++++++++++++++ 3 files changed, 37 insertions(+), 1 deletion(-) 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 index c47e5f05bd..11a6cdba7e 100644 --- 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); + newSqlRequest.toAnonymousString(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 index 073fec61e7..cbaf56bd0c 100644 --- 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 @@ -5,9 +5,11 @@ package org.opensearch.sql.legacy.unittest.utils; +import java.util.Map; import org.junit.Assert; import org.junit.Test; import org.opensearch.sql.legacy.utils.QueryDataAnonymizer; +import org.opensearch.sql.sql.domain.SQLQueryRequest; public class QueryDataAnonymizerTest { @@ -91,4 +93,20 @@ public void unionQueriesShouldAnonymizeSensitiveData() { + "UNION SELECT identifier, identifier FROM table )"; Assert.assertEquals(expectedQuery, QueryDataAnonymizer.anonymizeData(query)); } + + @Test + public void test_to_anonymous_string_of_SQLQueryRequest() { + String query = + "SELECT a.account_number, a.firstname, a.lastname, e.id, e.name " + + "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 expectedQuery = + "SQLQueryRequest(query=( SELECT identifier, identifier, identifier, identifier, identifier" + + " FROM table a JOIN table e ), path=/_plugins/_sql, format=jdbc," + + " params={pretty=true}, sanitize=true, cursor=Optional.empty)"; + Assert.assertEquals(expectedQuery, actualQuery); + } } 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 index 4e902cb67d..ac028b9856 100644 --- a/sql/src/main/java/org/opensearch/sql/sql/domain/SQLQueryRequest.java +++ b/sql/src/main/java/org/opensearch/sql/sql/domain/SQLQueryRequest.java @@ -10,6 +10,7 @@ import java.util.Map; import java.util.Optional; import java.util.Set; +import java.util.function.Function; import java.util.stream.Stream; import lombok.EqualsAndHashCode; import lombok.Getter; @@ -148,4 +149,21 @@ private boolean shouldSanitize(Map params) { } return true; } + + public String toAnonymousString(Function anonymizer) { + return "SQLQueryRequest(" + + "query=" + + anonymizer.apply(getQuery()) + + ", path=" + + path + + ", format=" + + format + + ", params=" + + params + + ", sanitize=" + + sanitize + + ", cursor=" + + getCursor() + + ')'; + } } From 16d681aa6b13fabd58c993249996100faeec0df6 Mon Sep 17 00:00:00 2001 From: Lantao Jin Date: Thu, 4 Jul 2024 19:46:44 +0800 Subject: [PATCH 2/3] fix typo Signed-off-by: Lantao Jin --- .../sql/legacy/unittest/utils/QueryDataAnonymizerTest.java | 1 - 1 file changed, 1 deletion(-) 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 index cbaf56bd0c..6f96bc39a3 100644 --- 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,6 @@ 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 expectedQuery = "SQLQueryRequest(query=( SELECT identifier, identifier, identifier, identifier, identifier" From 4572e25ab9f8532c4e00df81aebe3f23c6938da3 Mon Sep 17 00:00:00 2001 From: Lantao Jin Date: Thu, 4 Jul 2024 22:07:04 +0800 Subject: [PATCH 3/3] address coverage Signed-off-by: Lantao Jin --- .../opensearch/sql/legacy/plugin/RestSqlAction.java | 2 +- .../legacy/unittest/utils/QueryDataAnonymizerTest.java | 2 +- .../org/opensearch/sql/sql/domain/SQLQueryRequest.java | 3 ++- .../opensearch/sql/sql/domain/SQLQueryRequestTest.java | 10 ++++++++++ 4 files changed, 14 insertions(+), 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 index 11a6cdba7e..3a47af2700 100644 --- 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 index 6f96bc39a3..be17d828a1 100644 --- 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 index ac028b9856..1169afb121 100644 --- 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 params) { return true; } - public String toAnonymousString(Function anonymizer) { + /** A new toString() with anonymizer parameter to anonymize its query statement. */ + public String toString(Function 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 index 2b64b13b35..fc1434301e 100644 --- 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;