[SPARK-33526][SQL] Add config to control if cancel invoke interrupt task on thriftserver

sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala

@@ -937,6 +937,16 @@ object SQLConf {
       .timeConf(TimeUnit.SECONDS)
       .createWithDefault(0L)
 
+  val THRIFTSERVER_FORCE_CANCEL =
+    buildConf("spark.sql.thriftServer.forceCancel")
+      .doc("When true, all the job of query will be cancelled and running tasks will be" +
+        "interrupted. When false, all the job of query will be cancelled but running task" +
+        "will be remained until finished. Note that, this config must be set before query" +
+        "otherwise it doesn't help.")
+      .version("3.2.0")
+      .booleanConf
+      .createWithDefault(false)
+
   val THRIFTSERVER_UI_STATEMENT_LIMIT =
     buildConf("spark.sql.thriftserver.ui.retainedStatements")
       .doc("The number of SQL statements kept in the JDBC/ODBC web UI history.")

sql/hive-thriftserver/src/main/scala/org/apache/spark/sql/hive/thriftserver/SparkExecuteStatementOperation.scala

@@ -63,6 +63,8 @@ private[hive] class SparkExecuteStatementOperation(
     }
   }
 
+  private val forceCancel = sqlContext.conf.getConf(SQLConf.THRIFTSERVER_FORCE_CANCEL)
+
   private val substitutorStatement = SQLConf.withExistingConf(sqlContext.conf) {
     new VariableSubstitution().substitute(statement)
   }
@@ -131,7 +133,7 @@ private[hive] class SparkExecuteStatementOperation(
 
   def getNextRowSet(order: FetchOrientation, maxRowsL: Long): RowSet = withLocalProperties {
     try {
-      sqlContext.sparkContext.setJobGroup(statementId, substitutorStatement)
+      sqlContext.sparkContext.setJobGroup(statementId, substitutorStatement, forceCancel)
       getNextRowSetInternal(order, maxRowsL)
     } finally {
       sqlContext.sparkContext.clearJobGroup()
@@ -321,7 +323,7 @@ private[hive] class SparkExecuteStatementOperation(
         parentSession.getSessionState.getConf.setClassLoader(executionHiveClassLoader)
       }
 
-      sqlContext.sparkContext.setJobGroup(statementId, substitutorStatement)
+      sqlContext.sparkContext.setJobGroup(statementId, substitutorStatement, forceCancel)
       result = sqlContext.sql(statement)
       logDebug(result.queryExecution.toString())
       HiveThriftServer2.eventManager.onStatementParsed(statementId,

sql/hive-thriftserver/src/test/scala/org/apache/spark/sql/hive/thriftserver/ThriftServerWithSparkContextSuite.scala

@@ -18,9 +18,14 @@
 package org.apache.spark.sql.hive.thriftserver
 
 import java.sql.SQLException
+import java.util.concurrent.atomic.AtomicBoolean
 
 import org.apache.hive.service.cli.HiveSQLException
 
+import org.apache.spark.TaskKilled
+import org.apache.spark.scheduler.{SparkListener, SparkListenerTaskEnd}
+import org.apache.spark.sql.internal.SQLConf
+
 trait ThriftServerWithSparkContextSuite extends SharedThriftServer {
 
   test("the scratch dir will be deleted during server start but recreated with new operation") {
@@ -79,6 +84,45 @@ trait ThriftServerWithSparkContextSuite extends SharedThriftServer {
         "java.lang.NumberFormatException: invalid input syntax for type numeric: 1.2"))
     }
   }
+
+  test("SPARK-33526: Add config to control if cancel invoke interrupt task on thriftserver") {
+    withJdbcStatement { statement =>
+      val forceCancel = new AtomicBoolean(false)
+      val listener = new SparkListener {
+        override def onTaskEnd(taskEnd: SparkListenerTaskEnd): Unit = {
+          taskEnd.reason match {
+            case _: TaskKilled =>
+              if (forceCancel.get()) {
+                assert(System.currentTimeMillis() - taskEnd.taskInfo.launchTime < 1000)
+              } else {
+                assert(System.currentTimeMillis() - taskEnd.taskInfo.launchTime >= 2000)
+              }
+            case _ =>
+          }
+        }
+      }
+
+      spark.sparkContext.addSparkListener(listener)
+      try {
+        statement.execute(s"SET ${SQLConf.THRIFTSERVER_QUERY_TIMEOUT.key}=1")
+        statement.execute(s"SET ${SQLConf.THRIFTSERVER_FORCE_CANCEL.key}=false")
+        forceCancel.set(false)
+        val e1 = intercept[SQLException] {
+          statement.execute("select java_method('java.lang.Thread', 'sleep', 3000L)")
+        }.getMessage
+        assert(e1.contains("Query timed out"))
+
+        statement.execute(s"SET ${SQLConf.THRIFTSERVER_FORCE_CANCEL.key}=true")
+        forceCancel.set(true)
+        val e2 = intercept[SQLException] {
+          statement.execute("select java_method('java.lang.Thread', 'sleep', 3000L)")
+        }.getMessage
+        assert(e2.contains("Query timed out"))
+      } finally {
+        spark.sparkContext.removeSparkListener(listener)
+      }
+    }
+  }
 }