Bug Fix: Escape query in spark submit parameter (#2587)

* escape query in spark submit parameter Signed-off-by: Sean Kao <seankao@amazon.com> * spotless Signed-off-by: Sean Kao <seankao@amazon.com> * test case for special character Signed-off-by: Sean Kao <seankao@amazon.com> --------- Signed-off-by: Sean Kao <seankao@amazon.com>
opensearch-project · Mar 25, 2024 · 2bcf0b8 · 2bcf0b8
1 parent 85dae6f
commit 2bcf0b8
Show file tree

Hide file tree

Showing 5 changed files with 30 additions and 4 deletions.
diff --git a/common/build.gradle b/common/build.gradle
@@ -37,6 +37,7 @@ dependencies {
     api group: 'com.google.guava', name: 'guava', version: '32.0.1-jre'
     api group: 'org.apache.logging.log4j', name: 'log4j-core', version:"${versions.log4j}"
     api group: 'org.apache.commons', name: 'commons-lang3', version: '3.12.0'
+    api group: 'org.apache.commons', name: 'commons-text', version: '1.10.0'
     api group: 'com.squareup.okhttp3', name: 'okhttp', version: '4.12.0'
     implementation 'com.github.babbel:okhttp-aws-signer:1.0.2'
     api group: 'com.amazonaws', name: 'aws-java-sdk-core', version: "${aws_java_sdk_version}"

diff --git a/core/build.gradle b/core/build.gradle
@@ -48,6 +48,7 @@ pitest {
 dependencies {
     api group: 'com.google.guava', name: 'guava', version: '32.0.1-jre'
     api group: 'org.apache.commons', name: 'commons-lang3', version: '3.12.0'
+    api group: 'org.apache.commons', name: 'commons-text', version: '1.10.0'
     api group: 'com.facebook.presto', name: 'presto-matching', version: '0.240'
     api group: 'org.apache.commons', name: 'commons-math3', version: '3.6.1'
     api "com.fasterxml.jackson.core:jackson-core:${versions.jackson}"

diff --git a/legacy/build.gradle b/legacy/build.gradle
@@ -110,6 +110,7 @@ dependencies {
     implementation group: 'com.google.guava', name: 'guava', version: '32.0.1-jre'
     implementation group: 'org.json', name: 'json', version:'20231013'
     implementation group: 'org.apache.commons', name: 'commons-lang3', version: '3.12.0'
+    implementation group: 'org.apache.commons', name: 'commons-text', version: '1.10.0'
     implementation group: 'org.opensearch', name: 'opensearch', version: "${opensearch_version}"
     // add geo module as dependency. https://github.com/opensearch-project/OpenSearch/pull/4180/.
     implementation group: 'org.opensearch.plugin', name: 'geo', version: "${opensearch_version}"

diff --git a/spark/src/main/java/org/opensearch/sql/spark/asyncquery/model/SparkSubmitParameters.java b/spark/src/main/java/org/opensearch/sql/spark/asyncquery/model/SparkSubmitParameters.java
@@ -21,6 +21,7 @@
 import java.util.function.Supplier;
 import lombok.AllArgsConstructor;
 import lombok.RequiredArgsConstructor;
+import org.apache.commons.text.StringEscapeUtils;
 import org.opensearch.sql.datasource.model.DataSourceMetadata;
 import org.opensearch.sql.datasource.model.DataSourceType;
 import org.opensearch.sql.datasources.auth.AuthenticationType;
@@ -85,8 +86,13 @@ public Builder clusterName(String clusterName) {
       return this;
     }
 
+    /**
+     * For query in spark submit parameters to be parsed correctly, escape the characters in the
+     * query, then wrap the query with double quotes.
+     */
     public Builder query(String query) {
-      String wrappedQuery = "\"" + query + "\""; // Wrap the query with double quotes
+      String escapedQuery = StringEscapeUtils.escapeJava(query);
+      String wrappedQuery = "\"" + escapedQuery + "\"";
       config.put(FLINT_JOB_QUERY, wrappedQuery);
       return this;
     }

diff --git a/spark/src/test/java/org/opensearch/sql/spark/asyncquery/model/SparkSubmitParametersTest.java b/spark/src/test/java/org/opensearch/sql/spark/asyncquery/model/SparkSubmitParametersTest.java
@@ -30,8 +30,25 @@ public void testBuildWithExtraParameters() {
 
   @Test
   public void testBuildQueryString() {
-    String query = "SHOW tables LIKE \"%\";";
-    String params = SparkSubmitParameters.Builder.builder().query(query).build().toString();
-    assertTrue(params.contains(query));
+    String rawQuery = "SHOW tables LIKE \"%\";";
+    String expectedQueryInParams = "\"SHOW tables LIKE \\\"%\\\";\"";
+    String params = SparkSubmitParameters.Builder.builder().query(rawQuery).build().toString();
+    assertTrue(params.contains(expectedQueryInParams));
+  }
+
+  @Test
+  public void testBuildQueryStringNestedQuote() {
+    String rawQuery = "SELECT '\"1\"'";
+    String expectedQueryInParams = "\"SELECT '\\\"1\\\"'\"";
+    String params = SparkSubmitParameters.Builder.builder().query(rawQuery).build().toString();
+    assertTrue(params.contains(expectedQueryInParams));
+  }
+
+  @Test
+  public void testBuildQueryStringSpecialCharacter() {
+    String rawQuery = "SELECT '{\"test ,:+\\\"inner\\\"/\\|?#><\"}'";
+    String expectedQueryInParams = "SELECT '{\\\"test ,:+\\\\\\\"inner\\\\\\\"/\\\\|?#><\\\"}'";
+    String params = SparkSubmitParameters.Builder.builder().query(rawQuery).build().toString();
+    assertTrue(params.contains(expectedQueryInParams));
   }
 }