set SQLSession for async execution

sql/core/src/main/scala/org/apache/spark/sql/SQLContext.scala

@@ -916,6 +916,11 @@ class SQLContext(@transient val sparkContext: SparkContext)
     tlSession.remove()
   }
 
+  protected[sql] def setSession(session: SQLSession): Unit = {
+    detachSession()
+    tlSession.set(session)
+  }
+
   protected[sql] class SQLSession {
     // Note that this is a lazy val so we can override the default value in subclasses.
     protected[sql] lazy val conf: SQLConf = new SQLConf

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

@@ -154,6 +154,7 @@ private[hive] class SparkExecuteStatementOperation(
       val hiveConf = getConfigForOperation()
       val sparkServiceUGI = ShimLoader.getHadoopShims.getUGIForConf(hiveConf)
       val sessionHive = getCurrentHive()
+      val currentSqlSession = hiveContext.currentSession
 
       // Runnable impl to call runInternal asynchronously,
       // from a different thread
@@ -164,6 +165,7 @@ private[hive] class SparkExecuteStatementOperation(
             override def run(): Object = {
 
               // User information is part of the metastore client member in Hive
+              hiveContext.setSession(currentSqlSession)
               Hive.set(sessionHive)
               SessionState.setCurrentSessionState(parentSessionState)
               try {

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

@@ -19,11 +19,14 @@ package org.apache.spark.sql.hive.thriftserver
 
 import java.io.File
 import java.net.URL
-import java.sql.{Date, DriverManager, Statement}
+import java.nio.charset.StandardCharsets
+import java.nio.file.{Files, Paths}
+import java.sql.{Date, DriverManager, SQLException, Statement}
 
 import scala.collection.mutable.ArrayBuffer
 import scala.concurrent.duration._
-import scala.concurrent.{Await, Promise}
+import scala.concurrent.{Await, Promise, future}
+import scala.concurrent.ExecutionContext.Implicits.global
 import scala.sys.process.{Process, ProcessLogger}
 import scala.util.{Random, Try}
 
@@ -338,6 +341,30 @@ class HiveThriftBinaryServerSuite extends HiveThriftJdbcTest {
       }
     )
   }
+
+  test("test jdbc cancel") {
+    withJdbcStatement { statement =>
+      val queries = Seq(
+        "DROP TABLE IF EXISTS test_map",
+        "CREATE TABLE test_map(key INT, value STRING)",
+        s"LOAD DATA LOCAL INPATH '${TestData.smallKv}' OVERWRITE INTO TABLE test_map")
+
+      queries.foreach(statement.execute)
+
+      val f = future { Thread.sleep(3000); statement.cancel(); }
+
+      val join = "SELECT COUNT(*) FROM test_map " + List.fill(10)("join test_map").mkString(" ")
+      val e = intercept[SQLException] {
+        statement.executeQuery(join)
+      }
+      assert(e.getMessage contains "cancelled")
+      Await.result(f, Duration.Inf)
+
+      val rs1 = statement.executeQuery("SELECT COUNT(*) FROM test_map")
+      rs1.next()
+      assert(5 == rs1.getInt(1))
+    }
+  }
 }
 
 class HiveThriftHttpServerSuite extends HiveThriftJdbcTest {