From 405068b3b147b96f4db9c510a699440c27ac8e6a Mon Sep 17 00:00:00 2001
From: Louis Chu <clingzhi@amazon.com>
Date: Tue, 19 Mar 2024 16:27:23 -0700
Subject: [PATCH] Wrap the query with double quotes (#2565)

Signed-off-by: Louis Chu <clingzhi@amazon.com>
---
 .../sql/spark/asyncquery/model/SparkSubmitParameters.java      | 3 ++-
 .../sql/spark/dispatcher/SparkQueryDispatcherTest.java         | 1 +
 2 files changed, 3 insertions(+), 1 deletion(-)

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
index e3fe931a9e..e6d1dcd8c8 100644
--- 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
@@ -86,7 +86,8 @@ public Builder clusterName(String clusterName) {
     }
 
     public Builder query(String query) {
-      config.put(FLINT_JOB_QUERY, query);
+      String wrappedQuery = "\"" + query + "\""; // Wrap the query with double quotes
+      config.put(FLINT_JOB_QUERY, wrappedQuery);
       return this;
     }
 
diff --git a/spark/src/test/java/org/opensearch/sql/spark/dispatcher/SparkQueryDispatcherTest.java b/spark/src/test/java/org/opensearch/sql/spark/dispatcher/SparkQueryDispatcherTest.java
index d1d5033ee0..fc4bfb1923 100644
--- a/spark/src/test/java/org/opensearch/sql/spark/dispatcher/SparkQueryDispatcherTest.java
+++ b/spark/src/test/java/org/opensearch/sql/spark/dispatcher/SparkQueryDispatcherTest.java
@@ -947,6 +947,7 @@ private String constructExpectedSparkSubmitParameterString(
       authParamConfigBuilder.append(authParams.get(key));
       authParamConfigBuilder.append(" ");
     }
+    query = "\"" + query + "\"";
     return " --class org.apache.spark.sql.FlintJob  --conf"
                + " spark.hadoop.fs.s3.customAWSCredentialsProvider=com.amazonaws.emr.AssumeRoleAWSCredentialsProvider"
                + "  --conf"