commit code

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

@@ -138,7 +138,7 @@ object HiveThriftServer2 extends Logging {
   }
 
   private[thriftserver] object ExecutionState extends Enumeration {
-    val STARTED, COMPILED, FAILED, FINISHED, CLOSED = Value
+    val STARTED, COMPILED, CANCELED, FAILED, FINISHED, CLOSED = Value
     type ExecutionState = Value
   }
 
@@ -239,6 +239,14 @@ object HiveThriftServer2 extends Logging {
       executionList(id).state = ExecutionState.COMPILED
     }
 
+    def onStatementCanceled(id: String): Unit = synchronized {
+      executionList(id).finishTimestamp = System.currentTimeMillis
+      executionList(id).state = ExecutionState.CANCELED
+      totalRunning -= 1
+      // for progress bar
+      trimExecutionIfNecessary()
+    }
+
     def onStatementError(id: String, errorMessage: String, errorTrace: String): Unit = {
       synchronized {
         executionList(id).finishTimestamp = System.currentTimeMillis

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

@@ -70,6 +70,14 @@ private[hive] class SparkExecuteStatementOperation(
     }
   }
 
+  def cancelStatement(): Unit = {
+    // RDDs will be cleaned automatically upon garbage collection.
+    logDebug(s"CANCELING $statementId")
+    cleanup(OperationState.CANCELED)
+    sqlContext.sparkContext.clearJobGroup()
+    HiveThriftServer2.listener.onStatementCanceled(statementId)
+  }
+
   override def close(): Unit = {
     // RDDs will be cleaned automatically upon garbage collection.
     logDebug(s"CLOSING $statementId")