[WIP][SPARK-48353][SQL] Introduction of Error Handling mechanism in SQL Scripting

common/utils/src/main/resources/error/error-conditions.json

@@ -1128,6 +1128,18 @@
     ],
     "sqlState" : "42614"
   },
+  "DUPLICATE_CONDITION_NAME_FOR_DIFFERENT_SQL_STATE" : {
+    "message" : [
+      "Found duplicate condition name <conditionName> for different SQL states. Please, remove one of them."
+    ],
+    "sqlState" : "42710"
+  },
+  "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>."
@@ -1152,6 +1164,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>",
@@ -4824,6 +4842,11 @@
           "<variableName> is a VARIABLE and cannot be updated using the SET statement. Use SET VARIABLE <variableName> = ... instead."
         ]
       },
+      "SQL_SCRIPTING_NOT_ENABLED" : {
+        "message" : [
+          "SQL scripting is under development and not all features are supported. To enable existing features set <sqlScriptingEnabled> to `true`."
+        ]
+      },
       "STATE_STORE_MULTIPLE_COLUMN_FAMILIES" : {
         "message" : [
           "Creating multiple column families with <stateStoreProvider> is not supported."

docs/sql-ref-ansi-compliance.md

@@ -454,6 +454,7 @@ Below is a list of all the keywords in Spark SQL.
 |CONCATENATE|non-reserved|non-reserved|non-reserved|
 |CONSTRAINT|reserved|non-reserved|reserved|
 |CONTAINS|non-reserved|non-reserved|non-reserved|
+|CONTINUE|non-reserved|non-reserved|non-reserved|
 |COST|non-reserved|non-reserved|non-reserved|
 |CREATE|reserved|non-reserved|reserved|
 |CROSS|reserved|strict-non-reserved|reserved|
@@ -530,6 +531,7 @@ Below is a list of all the keywords in Spark SQL.
 |GRANT|reserved|non-reserved|reserved|
 |GROUP|reserved|non-reserved|reserved|
 |GROUPING|non-reserved|non-reserved|reserved|
+|HANDLER|non-reserved|non-reserved|non-reserved|
 |HAVING|reserved|non-reserved|reserved|
 |HOUR|non-reserved|non-reserved|non-reserved|
 |HOURS|non-reserved|non-reserved|non-reserved|

python/pyspark/errors/error-conditions.json

@@ -913,6 +913,11 @@
       "Slice with step is not supported."
     ]
   },
+  "SQL_SCRIPTING_NOT_ENABLED": {
+    "message": [
+      "SQL Scripting is under development and not all features are supported. To enable existing features set <sqlScriptingEnabled> to `true`."
+    ]
+  },
   "STATE_NOT_EXISTS": {
     "message": [
       "State is either not defined or has already been removed."

sql/api/src/main/antlr4/org/apache/spark/sql/catalyst/parser/SqlBaseLexer.g4

@@ -172,8 +172,10 @@ COMPACTIONS: 'COMPACTIONS';
 COMPENSATION: 'COMPENSATION';
 COMPUTE: 'COMPUTE';
 CONCATENATE: 'CONCATENATE';
+CONDITION: 'CONDITION';
 CONSTRAINT: 'CONSTRAINT';
 CONTAINS: 'CONTAINS';
+CONTINUE: 'CONTINUE';
 COST: 'COST';
 CREATE: 'CREATE';
 CROSS: 'CROSS';
@@ -223,6 +225,7 @@ EXCEPT: 'EXCEPT';
 EXCHANGE: 'EXCHANGE';
 EXCLUDE: 'EXCLUDE';
 EXISTS: 'EXISTS';
+EXIT: 'EXIT';
 EXPLAIN: 'EXPLAIN';
 EXPORT: 'EXPORT';
 EXTENDED: 'EXTENDED';
@@ -240,6 +243,7 @@ FOR: 'FOR';
 FOREIGN: 'FOREIGN';
 FORMAT: 'FORMAT';
 FORMATTED: 'FORMATTED';
+FOUND: 'FOUND';
 FROM: 'FROM';
 FULL: 'FULL';
 FUNCTION: 'FUNCTION';
@@ -249,6 +253,7 @@ GLOBAL: 'GLOBAL';
 GRANT: 'GRANT';
 GROUP: 'GROUP';
 GROUPING: 'GROUPING';
+HANDLER: 'HANDLER';
 HAVING: 'HAVING';
 BINARY_HEX: 'X';
 HOUR: 'HOUR';
@@ -402,6 +407,7 @@ SORTED: 'SORTED';
 SOURCE: 'SOURCE';
 SPECIFIC: 'SPECIFIC';
 SQL: 'SQL';
+SQLEXCEPTION: 'SQLEXCEPTION';
 START: 'START';
 STATISTICS: 'STATISTICS';
 STORED: 'STORED';
@@ -570,6 +576,10 @@ IDENTIFIER
     | UNICODE_LETTER+ '://' (UNICODE_LETTER | DIGIT | '_' | '/' | '-' | '.' | '?' | '=' | '&' | '#' | '%')+
     ;
 
+SQLSTATE
+    : '\'' (LETTER | DIGIT)(LETTER | DIGIT)(LETTER | DIGIT)(LETTER | DIGIT)(LETTER | DIGIT) '\''
+    ;
+
 BACKQUOTED_IDENTIFIER
     : '`' ( ~'`' | '``' )* '`'
     ;

sql/api/src/main/antlr4/org/apache/spark/sql/catalyst/parser/SqlBaseParser.g4

@@ -63,6 +63,8 @@ compoundStatement
     : statement
     | setStatementWithOptionalVarKeyword
     | beginEndCompoundBlock
+    | declareCondition
+    | declareHandler
     | ifElseStatement
     ;
 
@@ -82,6 +84,23 @@ singleStatement
     : (statement|setResetStatement) SEMICOLON* EOF
     ;
 
+conditionValue
+    : stringLit
+    | multipartIdentifier
+    ;
+
+conditionValueList
+    : ((conditionValues+=conditionValue (COMMA conditionValues+=conditionValue)*) | SQLEXCEPTION | NOT FOUND)
+    ;
+
+declareCondition
+    : DECLARE multipartIdentifier CONDITION (FOR stringLit)?
+    ;
+
+declareHandler
+    : DECLARE (CONTINUE | EXIT) HANDLER FOR conditionValueList (BEGIN compoundBody END | statement)
+    ;
+
 beginLabel
     : multipartIdentifier COLON
     ;
@@ -1498,6 +1517,7 @@ ansiNonReserved
     | COMPUTE
     | CONCATENATE
     | CONTAINS
+    | CONTINUE
     | COST
     | CUBE
     | CURRENT
@@ -1536,6 +1556,7 @@ ansiNonReserved
     | EXCHANGE
     | EXCLUDE
     | EXISTS
+    | EXIT
     | EXPLAIN
     | EXPORT
     | EXTENDED
@@ -1823,6 +1844,7 @@ nonReserved
     | CONCATENATE
     | CONSTRAINT
     | CONTAINS
+    | CONTINUE
     | COST
     | CREATE
     | CUBE
@@ -1871,6 +1893,7 @@ nonReserved
     | EXCLUDE
     | EXECUTE
     | EXISTS
+    | EXIT
     | EXPLAIN
     | EXPORT
     | EXTENDED

sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/parser/AstBuilder.scala

@@ -20,6 +20,7 @@ package org.apache.spark.sql.catalyst.parser
 import java.util.Locale
 import java.util.concurrent.TimeUnit
 
+import scala.collection.mutable
 import scala.collection.mutable.{ArrayBuffer, ListBuffer, Set}
 import scala.jdk.CollectionConverters._
 import scala.util.{Left, Right}
@@ -119,11 +120,13 @@ class AstBuilder extends DataTypeAstBuilder
   override def visitCompoundOrSingleStatement(
       ctx: CompoundOrSingleStatementContext): CompoundBody = withOrigin(ctx) {
     Option(ctx.singleCompoundStatement()).map { s =>
+      if (!SQLConf.get.sqlScriptingEnabled) {
+        throw SqlScriptingErrors.sqlScriptingNotEnabled(CurrentOrigin.get)
+      }
       visit(s).asInstanceOf[CompoundBody]
     }.getOrElse {
       val logicalPlan = visitSingleStatement(ctx.singleStatement())
-      CompoundBody(Seq(SingleStatement(parsedPlan = logicalPlan)),
-        Some(java.util.UUID.randomUUID.toString.toLowerCase(Locale.ROOT)))
+      CompoundBody(Seq(SingleStatement(parsedPlan = logicalPlan)))
     }
   }
 
@@ -136,8 +139,24 @@ class AstBuilder extends DataTypeAstBuilder
       label: Option[String],
       allowVarDeclare: Boolean): CompoundBody = {
     val buff = ListBuffer[CompoundPlanStatement]()
+    val handlers = ListBuffer[ErrorHandler]()
+    val conditions = mutable.HashMap[String, String]()
+    val sqlStates = mutable.Set[String]()
+
     ctx.compoundStatements.forEach(compoundStatement => {
-      buff += visit(compoundStatement).asInstanceOf[CompoundPlanStatement]
+      val stmt = visit(compoundStatement).asInstanceOf[CompoundPlanStatement]
+
+      stmt match {
+        case handler: ErrorHandler => handlers += handler
+        case condition: ErrorCondition =>
+          if (conditions.contains(condition.conditionName)) {
+            throw SqlScriptingErrors.duplicateConditionNameForDifferentSqlState(
+              CurrentOrigin.get, condition.conditionName)
+          }
+          conditions += condition.conditionName -> condition.value
+          sqlStates += condition.value
+        case s => buff += s
+      }
     })
 
     val compoundStatements = buff.toList
@@ -171,7 +190,7 @@ class AstBuilder extends DataTypeAstBuilder
       case _ =>
     }
 
-    CompoundBody(buff.toSeq, label)
+    CompoundBody(buff.toSeq, label, handlers.toSeq, conditions)
   }
 
   override def visitBeginEndCompoundBlock(ctx: BeginEndCompoundBlockContext): CompoundBody = {
@@ -216,6 +235,27 @@ class AstBuilder extends DataTypeAstBuilder
         }
     }
 
+  override def visitConditionValue(ctx: ConditionValueContext): String = {
+    Option(ctx.multipartIdentifier()).map(_.getText)
+      .getOrElse(ctx.stringLit().getText).replace("'", "")
+  }
+
+  override def visitConditionValueList(ctx: ConditionValueListContext): Seq[String] = {
+    Option(ctx.SQLEXCEPTION()).map(_ => Seq("SQLEXCEPTION")).getOrElse {
+      Option(ctx.NOT()).map(_ => Seq("NOT FOUND")).getOrElse {
+        val buff = scala.collection.mutable.Set[String]()
+        ctx.conditionValues.forEach { conditionValue =>
+          val elem = visit(conditionValue).asInstanceOf[String]
+          if (buff(elem)) {
+            throw SqlScriptingErrors.duplicateSqlStateForSameHandler(CurrentOrigin.get, elem)
+          }
+          buff += elem
+        }
+        buff.toSeq
+      }
+    }
+  }
+
   override def visitIfElseStatement(ctx: IfElseStatementContext): IfElseStatement = {
     IfElseStatement(
       conditions = ctx.booleanExpression().asScala.toList.map(boolExpr => withOrigin(boolExpr) {
@@ -229,6 +269,28 @@ class AstBuilder extends DataTypeAstBuilder
     )
   }
 
+  override def visitDeclareCondition(ctx: DeclareConditionContext): ErrorCondition = {
+    val conditionName = ctx.multipartIdentifier().getText
+    val conditionValue = Option(ctx.stringLit()).map(_.getText.replace("'", "")).getOrElse("45000")
+
+    val sqlStateRegex = "^[A-Za-z0-9]{5}$".r
+    assert(sqlStateRegex.findFirstIn(conditionValue).isDefined)
+
+    ErrorCondition(conditionName, conditionValue)
+  }
+
+  override def visitDeclareHandler(ctx: DeclareHandlerContext): ErrorHandler = {
+    val conditions = visit(ctx.conditionValueList()).asInstanceOf[Seq[String]]
+    val handlerType = Option(ctx.EXIT()).map(_ => HandlerType.EXIT).getOrElse(HandlerType.CONTINUE)
+
+    val body = Option(ctx.compoundBody()).map(visit).getOrElse {
+      val logicalPlan = visit(ctx.statement()).asInstanceOf[LogicalPlan]
+      CompoundBody(Seq(SingleStatement(parsedPlan = logicalPlan)))
+    }.asInstanceOf[CompoundBody]
+
+    ErrorHandler(conditions, body, handlerType)
+  }
+
   override def visitSingleStatement(ctx: SingleStatementContext): LogicalPlan = withOrigin(ctx) {
     Option(ctx.statement().asInstanceOf[ParserRuleContext])
       .orElse(Option(ctx.setResetStatement().asInstanceOf[ParserRuleContext]))

sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/parser/SqlScriptingLogicalOperators.scala

@@ -17,6 +17,9 @@
 
 package org.apache.spark.sql.catalyst.parser
 
+import scala.collection.mutable
+
+import org.apache.spark.sql.catalyst.parser.HandlerType.HandlerType
 import org.apache.spark.sql.catalyst.plans.logical.LogicalPlan
 import org.apache.spark.sql.catalyst.trees.{CurrentOrigin, Origin, WithOrigin}
 
@@ -54,10 +57,39 @@ case class SingleStatement(parsedPlan: LogicalPlan)
  * @param label Label set to CompoundBody by user or UUID otherwise.
  *              It can be None in case when CompoundBody is not part of BeginEndCompoundBlock
  *              for example when CompoundBody is inside loop or conditional block.
+ * @param handlers Collection of handlers defined in the compound body.
+ * @param conditions Map of Condition Name - Sql State values declared in the compound body.
  */
 case class CompoundBody(
     collection: Seq[CompoundPlanStatement],
-    label: Option[String]) extends CompoundPlanStatement
+    label: Option[String] = None,
+    handlers: Seq[ErrorHandler] = Seq.empty,
+    conditions: mutable.HashMap[String, String] = mutable.HashMap()) extends CompoundPlanStatement
+
+/**
+ * Logical operator for an error condition.
+ * @param conditionName Name of the error condition.
+ * @param value SQLSTATE or Error Code.
+ */
+case class ErrorCondition(
+    conditionName: String,
+    value: String) extends CompoundPlanStatement
+
+object HandlerType extends Enumeration {
+    type HandlerType = Value
+    val EXIT, CONTINUE = Value
+}
+
+/**
+ * Logical operator for an error condition.
+ * @param conditions Name of the error condition variable for which the handler is built.
+ * @param body CompoundBody of the handler.
+ * @param handlerType Type of the handler (CONTINUE or EXIT).
+ */
+case class ErrorHandler(
+    conditions: Seq[String],
+    body: CompoundBody,
+    handlerType: HandlerType) extends CompoundPlanStatement
 
 /**
  * Logical operator for IF ELSE statement.