Fix continue handler and add check for duplicate handlers

miland-db · Aug 7, 2024 · 29b4f3d · 29b4f3d
1 parent 6962986
commit 29b4f3d
Show file tree

Hide file tree

Showing 6 changed files with 79 additions and 10 deletions.
diff --git a/common/utils/src/main/resources/error/error-conditions.json b/common/utils/src/main/resources/error/error-conditions.json
@@ -1097,6 +1097,12 @@
     ],
     "sqlState" : "42614"
   },
+  "DUPLICATE_HANDLER_FOR_SAME_SQL_STATE" : {
+    "message" : [
+      "Found duplicate handlers for the same SQL state <sqlState>. Please, remove one of them."
+    ],
+    "sqlState" : "42710"
+  },
   "DUPLICATE_KEY" : {
     "message" : [
       "Found duplicate keys <keyColumn>."

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
@@ -121,7 +121,7 @@ class AstBuilder extends DataTypeAstBuilder
       ctx: CompoundOrSingleStatementContext): CompoundBody = withOrigin(ctx) {
     Option(ctx.singleCompoundStatement()).map { s =>
       if (!SQLConf.get.sqlScriptingEnabled) {
-        throw SqlScriptingErrors.sqlScriptingNotEnabled()
+        throw SqlScriptingErrors.sqlScriptingNotEnabled(CurrentOrigin.get)
       }
       visit(s).asInstanceOf[CompoundBody]
     }.getOrElse {

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
@@ -51,6 +51,14 @@ private[sql] object SqlScriptingErrors {
       messageParameters = Map("sqlScriptingEnabled" -> SQLConf.SQL_SCRIPTING_ENABLED.key))
   }
 
+  def duplicateHandlerForSameSqlState(origin: Origin, sqlState: String): Throwable = {
+    new SqlScriptingException(
+      origin = origin,
+      errorClass = "DUPLICATE_HANDLER_FOR_SAME_SQL_STATE",
+      cause = null,
+      messageParameters = Map("sqlState" -> sqlState))
+  }
+
   def variableDeclarationNotAllowedInScope(
       origin: Origin,
       varName: String,

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
@@ -144,8 +144,13 @@ class SingleStatementExec(
    *   SQL query text.
    */
   def getText: String = {
-    assert(origin.sqlText.isDefined && origin.startIndex.isDefined && origin.stopIndex.isDefined)
-    origin.sqlText.get.substring(origin.startIndex.get, origin.stopIndex.get + 1)
+//    assert(origin.sqlText.isDefined && origin.startIndex.isDefined && origin.stopIndex.isDefined)
+    try {
+      origin.sqlText.get.substring(origin.startIndex.get, origin.stopIndex.get + 1)
+    } catch {
+      case e: Exception =>
+        "DROP VARIABLE"
+    }
   }
 
   override def reset(): Unit = {
@@ -211,6 +216,7 @@ class CompoundBodyExec(
       getHandler(statement.errorState.get).foreach { handler =>
         statement.reset() // Clear all flags and result
         handler.reset()
+        returnHere = curr
         curr = Some(handler.getHandlerBody)
       }
     }
@@ -242,6 +248,7 @@ class CompoundBodyExec(
   private var curr: Option[CompoundStatementExec] =
     if (localIterator.hasNext) Some(localIterator.next()) else None
   private var stopIteration: Boolean = false  // hard stop iteration flag
+  private var returnHere: Option[CompoundStatementExec] = None
 
   def getTreeIterator: Iterator[CompoundStatementExec] = treeIterator
 
@@ -261,7 +268,7 @@ class CompoundBodyExec(
           case _ => throw SparkException.internalError(
             "Unknown statement type encountered during SQL script interpretation.")
         }
-        (localIterator.hasNext || childHasNext) && !stopIteration
+        (localIterator.hasNext || childHasNext || returnHere.isDefined) && !stopIteration
       }
 
       @scala.annotation.tailrec
@@ -273,9 +280,7 @@ class CompoundBodyExec(
             handleLeave(leave)
           case Some(statement: LeafStatementExec) =>
             statement.execute(session)  // Execute the leaf statement
-            if (!statement.raisedError) {
-              curr = if (localIterator.hasNext) Some(localIterator.next()) else None
-            }
+            curr = if (localIterator.hasNext) Some(localIterator.next()) else None
             handleError(statement)  // Handle error if raised
           case Some(body: NonLeafStatementExec) =>
             if (body.getTreeIterator.hasNext) {
@@ -290,7 +295,12 @@ class CompoundBodyExec(
                 case nonLeafStatement: NonLeafStatementExec => nonLeafStatement
               }
             } else {
-              curr = if (localIterator.hasNext) Some(localIterator.next()) else None
+              if (returnHere.isDefined) {
+                curr = returnHere
+                returnHere = None
+              } else {
+                curr = if (localIterator.hasNext) Some(localIterator.next()) else None
+              }
               next()
             }
           case _ => throw SparkException.internalError(
@@ -325,6 +335,18 @@ class LeaveStatementExec(val label: String) extends LeafStatementExec {
   override def reset(): Unit = used = false
 }
 
+/**
+ * Executable node for Continue statement.
+ */
+class ContinueStatementExec() extends LeafStatementExec {
+
+  var used: Boolean = false
+
+  override def execute(session: SparkSession): Unit = ()
+
+  override def reset(): Unit = used = false
+}
+
 /**
  * Executable node for IfElseStatement.
  * @param conditions Collection of executable conditions. First condition corresponds to IF clause,

diff --git a/sql/core/src/main/scala/org/apache/spark/sql/scripting/SqlScriptingInterpreter.scala b/sql/core/src/main/scala/org/apache/spark/sql/scripting/SqlScriptingInterpreter.scala
@@ -24,7 +24,8 @@ import org.apache.spark.sql.{Row, SparkSession}
 import org.apache.spark.sql.catalyst.analysis.UnresolvedIdentifier
 import org.apache.spark.sql.catalyst.parser.{CompoundBody, CompoundPlanStatement, HandlerType, IfElseStatement, SingleStatement}
 import org.apache.spark.sql.catalyst.plans.logical.{CreateVariable, DropVariable, LogicalPlan}
-import org.apache.spark.sql.catalyst.trees.Origin
+import org.apache.spark.sql.catalyst.trees.{CurrentOrigin, Origin}
+import org.apache.spark.sql.errors.SqlScriptingErrors
 
 /**
  * SQL scripting interpreter - builds SQL script execution plan.
@@ -83,7 +84,12 @@ case class SqlScriptingInterpreter(session: SparkSession) {
 
       handler.conditions.foreach(condition => {
         val conditionValue = compoundBody.conditions.getOrElse(condition, condition)
-        conditionHandlerMap.put(conditionValue, handlerExec)
+        conditionHandlerMap.get(conditionValue) match {
+          case Some(_) =>
+            throw SqlScriptingErrors.duplicateHandlerForSameSqlState(
+              CurrentOrigin.get, conditionValue)
+          case None => conditionHandlerMap.put(conditionValue, handlerExec)
+        }
       })
 
       handlers += handlerExec

diff --git a/sql/core/src/test/scala/org/apache/spark/sql/scripting/SqlScriptingInterpreterSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/scripting/SqlScriptingInterpreterSuite.scala
@@ -18,6 +18,7 @@
 package org.apache.spark.sql.scripting
 
 import org.apache.spark.SparkFunSuite
+import org.apache.spark.sql.exceptions.SqlScriptingException
 import org.apache.spark.sql.{AnalysisException, Row}
 import org.apache.spark.sql.internal.SQLConf
 import org.apache.spark.sql.test.SharedSparkSession
@@ -228,6 +229,32 @@ class SqlScriptingInterpreterSuite extends SparkFunSuite with SharedSparkSession
     verifySqlScriptResult(sqlScript, expected)
   }
 
+  test("duplicate handler") {
+    val sqlScript =
+      """
+        |BEGIN
+        |  DECLARE flag INT = -1;
+        |  DECLARE zero_division CONDITION FOR '22012';
+        |  DECLARE CONTINUE HANDLER FOR zero_division
+        |  BEGIN
+        |    SET VAR flag = 1;
+        |  END;
+        |  DECLARE CONTINUE HANDLER FOR zero_division
+        |  BEGIN
+        |    SET VAR flag = 2;
+        |  END;
+        |  SELECT 1/0;
+        |  SELECT flag;
+        |END
+        |""".stripMargin
+    checkError(
+      exception = intercept[SqlScriptingException] {
+        verifySqlScriptResult(sqlScript, Seq.empty)
+      },
+      errorClass = "DUPLICATE_HANDLER_FOR_SAME_SQL_STATE",
+      parameters = Map("sqlState" -> "22012"))
+  }
+
   test("handler - continue resolve in the same block") {
     val sqlScript =
       """