Add check for duplicate sqlstate inside condition value list in signl…

…e handler
miland-db · Aug 8, 2024 · 17282c4 · 17282c4
1 parent 29b4f3d
commit 17282c4
Show file tree

Hide file tree

Showing 5 changed files with 41 additions and 6 deletions.
diff --git a/common/utils/src/main/resources/error/error-conditions.json b/common/utils/src/main/resources/error/error-conditions.json
@@ -1127,6 +1127,12 @@
     },
     "sqlState" : "4274K"
   },
+  "DUPLICATE_SQL_STATE_FOR_SAME_HANDLER" : {
+    "message" : [
+      "Found duplicate SQL state <sqlState> for the same handler. Please, remove one of them."
+    ],
+    "sqlState" : "42710"
+  },
   "EMITTING_ROWS_OLDER_THAN_WATERMARK_NOT_ALLOWED" : {
     "message" : [
       "Previous node emitted a row with eventTime=<emittedRowEventTime> which is older than current_watermark_value=<currentWatermark>",

diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/parser/AstBuilder.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/parser/AstBuilder.scala
@@ -245,8 +245,14 @@ class AstBuilder extends DataTypeAstBuilder
     Option(ctx.SQLEXCEPTION()).map(_ => Seq("SQLEXCEPTION")).getOrElse {
       Option(ctx.NOT()).map(_ => Seq("NOT FOUND")).getOrElse {
         val buff = ListBuffer[String]()
+        val seen = scala.collection.mutable.Set[String]()
         ctx.conditionValues.forEach { conditionValue =>
-          buff += visit(conditionValue).asInstanceOf[String]
+          val elem = visit(conditionValue).asInstanceOf[String]
+          if (seen(elem)) {
+            throw SqlScriptingErrors.duplicateSqlStateForSameHandler(CurrentOrigin.get, elem)
+          }
+          buff += elem
+          seen += elem
         }
         buff.toSeq
       }

diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/errors/SqlScriptingErrors.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/errors/SqlScriptingErrors.scala
@@ -59,6 +59,14 @@ private[sql] object SqlScriptingErrors {
       messageParameters = Map("sqlState" -> sqlState))
   }
 
+  def duplicateSqlStateForSameHandler(origin: Origin, sqlState: String): Throwable = {
+    new SqlScriptingException(
+      origin = origin,
+          errorClass = "DUPLICATE_SQL_STATE_FOR_SAME_HANDLER",
+      cause = null,
+      messageParameters = Map("sqlState" -> sqlState))
+  }
+
   def variableDeclarationNotAllowedInScope(
       origin: Origin,
       varName: String,

diff --git a/...atalyst/src/test/scala/org/apache/spark/sql/catalyst/parser/SqlScriptingParserSuite.scala b/...atalyst/src/test/scala/org/apache/spark/sql/catalyst/parser/SqlScriptingParserSuite.scala
@@ -396,6 +396,20 @@ class SqlScriptingParserSuite extends SparkFunSuite with SQLHelper {
     assert(tree.handlers.head.isInstanceOf[ErrorHandler])
   }
 
+  test("declare handler duplicate sqlState") {
+    val sqlScriptText =
+      """
+        |BEGIN
+        |  DECLARE CONTINUE HANDLER FOR test, test BEGIN SELECT 1; END;
+        |END""".stripMargin
+    checkError(
+      exception = intercept[SqlScriptingException] {
+        parseScript(sqlScriptText)
+      },
+      errorClass = "DUPLICATE_SQL_STATE_FOR_SAME_HANDLER",
+      parameters = Map("sqlState" -> "test"))
+  }
+
   test("SQL Scripting not enabled") {
     withSQLConf(SQLConf.SQL_SCRIPTING_ENABLED.key -> "false") {
       val sqlScriptText =

diff --git a/sql/core/src/main/scala/org/apache/spark/sql/scripting/SqlScriptingExecutionNode.scala b/sql/core/src/main/scala/org/apache/spark/sql/scripting/SqlScriptingExecutionNode.scala
@@ -199,11 +199,12 @@ class CompoundBodyExec(
   extends NonLeafStatementExec {
 
   private def getHandler(condition: String): Option[ErrorHandlerExec] = {
-    var ret = conditionHandlerMap.get(condition)
-    if (ret.isEmpty) {
-      ret = conditionHandlerMap.get("UNKNOWN")
-    }
-    ret
+    conditionHandlerMap.get(condition)
+      .orElse(conditionHandlerMap.get("NOT FOUND") match {
+        case Some(handler) if condition.startsWith("02") => Some(handler)
+        case _ => None
+      })
+      .orElse(conditionHandlerMap.get("UNKNOWN"))
   }
 
   /**