Fix

apache · Jul 16, 2020 · 7f8907e · 7f8907e
1 parent c28a6fa
commit 7f8907e
Show file tree

Hide file tree

Showing 2 changed files with 37 additions and 10 deletions.
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/SparkSqlParser.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/SparkSqlParser.scala
@@ -65,16 +65,19 @@ class SparkSqlAstBuilder(conf: SQLConf) extends AstBuilder(conf) {
    */
   override def visitSetConfiguration(ctx: SetConfigurationContext): LogicalPlan = withOrigin(ctx) {
     // Construct the command.
-    val raw = remainder(ctx.SET.getSymbol)
-    val keyValueSeparatorIndex = raw.indexOf('=')
-    if (keyValueSeparatorIndex >= 0) {
-      val key = raw.substring(0, keyValueSeparatorIndex).trim
-      val value = raw.substring(keyValueSeparatorIndex + 1).trim
-      SetCommand(Some(key -> Option(value)))
-    } else if (raw.nonEmpty) {
-      SetCommand(Some(raw.trim -> None))
-    } else {
-      SetCommand(None)
+    val keyValueDef = """\s*([a-zA-Z.]+|`[a-zA-Z.\s]+`)\s*=(.*)""".r
+    val keyOnlyDef = """\s*([a-zA-Z.]+|`[a-zA-Z.\s]+`)\s*""".r
+    remainder(ctx.SET.getSymbol) match {
+      case keyValueDef(key, value) =>
+        SetCommand(Some(key.replaceAll("`", "") -> Option(value)))
+      case keyOnlyDef(key) =>
+        SetCommand(Some(key.replaceAll("`", "") -> None))
+      case s if s.trim.isEmpty =>
+        SetCommand(None)
+      case _ =>
+        throw new ParseException("Expected format is 'SET key=value' or 'SET key'. " +
+          "If you want to include spaces in key, please use backquotes, " +
+          "e.g., 'SET `ke y`=value'.", ctx)
     }
   }
 

diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/SparkSqlParserSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/SparkSqlParserSuite.scala
@@ -61,6 +61,30 @@ class SparkSqlParserSuite extends AnalysisTest {
   private def intercept(sqlCommand: String, messages: String*): Unit =
     interceptParseException(parser.parsePlan)(sqlCommand, messages: _*)
 
+  test("Report Error for invalid usage of SET command") {
+    assertEqual("SET", SetCommand(None))
+    assertEqual("SET   ", SetCommand(None))
+    assertEqual("SET spark.sql.key", SetCommand(Some("spark.sql.key" -> None)))
+    assertEqual("SET   spark.sql.key ", SetCommand(Some("spark.sql.key" -> None)))
+    assertEqual("SET spark.sql.key=value", SetCommand(Some("spark.sql.key" -> Some("value"))))
+    assertEqual("SET spark.sql.key =  value", SetCommand(Some("spark.sql.key" -> Some("  value"))))
+    assertEqual("SET spark.sql.key= v al u   e ",
+      SetCommand(Some("spark.sql.key" -> Some(" v al u   e "))))
+    assertEqual("SET spark.sql.key= value1= value2",
+      SetCommand(Some("spark.sql.key" -> Some(" value1= value2"))))
+    assertEqual("SET `spark.sql.    key`=value",
+      SetCommand(Some("spark.sql.    key" -> Some("value"))))
+
+    val expectedErrMsg = "Expected format is 'SET key=value' or 'SET key'. " +
+      "If you want to include spaces in key, please use backquotes, " +
+      "e.g., 'SET `ke y`=value'."
+    intercept("SET spark.sql.key value", expectedErrMsg)
+    intercept("SET spark.sql.key   'value'", expectedErrMsg)
+    intercept("""SET    spark.sql.key "value" """, expectedErrMsg)
+    intercept("SET spark.sql.key value1 value2", expectedErrMsg)
+    intercept("SET spark.sql.    key=value", expectedErrMsg)
+  }
+
   test("refresh resource") {
     assertEqual("REFRESH prefix_path", RefreshResource("prefix_path"))
     assertEqual("REFRESH /", RefreshResource("/"))