[SPARK-21993][SQL] Close sessionState when finish

sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/catalog/ExternalCatalog.scala

@@ -360,4 +360,6 @@ abstract class ExternalCatalog
       event: ExternalCatalogEvent): Unit = {
     listener.onEvent(event)
   }
+
+  def close(): Unit = { }
 }

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

@@ -114,9 +114,12 @@ class SparkSession private(
   @InterfaceStability.Unstable
   @transient
   lazy val sharedState: SharedState = {
+    sharedStateInitialized = true
     existingSharedState.getOrElse(new SharedState(sparkContext))
   }
 
+  var sharedStateInitialized: Boolean = false
+
   /**
    * Initial options for session. This options are applied once when sessionState is created.
    */
@@ -706,6 +709,9 @@ class SparkSession private(
    */
   def stop(): Unit = {
     sparkContext.stop()
+    if (sharedStateInitialized) {
+      sharedState.close()
+    }
   }
 
   /**

sql/core/src/main/scala/org/apache/spark/sql/internal/SharedState.scala

@@ -87,6 +87,7 @@ private[sql] class SharedState(val sparkContext: SparkContext) extends Logging {
    */
   val listener: SQLListener = createListenerAndUI(sparkContext)
 
+  var externalCatalogInitialized: Boolean = false
   /**
    * A catalog that interacts with external systems.
    */
@@ -115,6 +116,7 @@ private[sql] class SharedState(val sparkContext: SparkContext) extends Logging {
       }
     })
 
+    externalCatalogInitialized = true
     externalCatalog
   }
 
@@ -154,6 +156,12 @@ private[sql] class SharedState(val sparkContext: SparkContext) extends Logging {
     }
     SparkSession.sqlListener.get()
   }
+
+  def close(): Unit = {
+    if (externalCatalogInitialized) {
+      externalCatalog.close()
+    }
+  }
 }
 
 object SharedState extends Logging {

sql/hive-thriftserver/src/main/java/org/apache/hive/service/cli/CLIService.java

@@ -126,6 +126,11 @@ private void applyAuthorizationConfigPolicy(HiveConf newHiveConf) throws HiveExc
     ss.setIsHiveServerQuery(true);
     SessionState.start(ss);
     ss.applyAuthorizationPolicy();
+    try {
+      ss.close();
+    } catch (IOException e) {
+      LOG.error("Failed closing Hive session state.", e);
+    }
   }
 
   private void setupBlockedUdfs() {

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

@@ -109,6 +109,8 @@ object HiveThriftServer2 extends Logging {
       case e: Exception =>
         logError("Error starting HiveThriftServer2", e)
         System.exit(-1)
+    } finally {
+      executionHive.close()
     }
   }
 

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

@@ -124,7 +124,10 @@ private[hive] object SparkSQLCLIDriver extends Logging {
     SessionState.start(sessionState)
 
     // Clean up after we exit
-    ShutdownHookManager.addShutdownHook { () => SparkSQLEnv.stop() }
+    ShutdownHookManager.addShutdownHook { () =>
+      sessionState.close()
+      SparkSQLEnv.stop()
+    }
 
     val remoteMode = isRemoteMode(sessionState)
     // "-h" option has been passed, so connect to Hive thrift server.

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

@@ -22,7 +22,7 @@ import java.io.PrintStream
 import org.apache.spark.{SparkConf, SparkContext}
 import org.apache.spark.internal.Logging
 import org.apache.spark.sql.{SparkSession, SQLContext}
-import org.apache.spark.sql.hive.{HiveExternalCatalog, HiveUtils}
+import org.apache.spark.sql.hive.HiveUtils
 import org.apache.spark.util.Utils
 
 /** A singleton object for the master program. The slaves should not access this. */
@@ -31,6 +31,7 @@ private[hive] object SparkSQLEnv extends Logging {
 
   var sqlContext: SQLContext = _
   var sparkContext: SparkContext = _
+  var sparkSession: SparkSession = _
 
   def init() {
     if (sqlContext == null) {
@@ -45,16 +46,10 @@ private[hive] object SparkSQLEnv extends Logging {
       sparkConf
         .setAppName(maybeAppName.getOrElse(s"SparkSQL::${Utils.localHostName()}"))
 
-      val sparkSession = SparkSession.builder.config(sparkConf).enableHiveSupport().getOrCreate()
+      sparkSession = SparkSession.builder.config(sparkConf).enableHiveSupport().getOrCreate()
       sparkContext = sparkSession.sparkContext
       sqlContext = sparkSession.sqlContext
 
-      val metadataHive = sparkSession
-        .sharedState.externalCatalog.asInstanceOf[HiveExternalCatalog]
-        .client.newSession()
-      metadataHive.setOut(new PrintStream(System.out, true, "UTF-8"))
-      metadataHive.setInfo(new PrintStream(System.err, true, "UTF-8"))
-      metadataHive.setError(new PrintStream(System.err, true, "UTF-8"))
       sparkSession.conf.set("spark.sql.hive.version", HiveUtils.hiveExecutionVersion)
     }
   }
@@ -65,6 +60,7 @@ private[hive] object SparkSQLEnv extends Logging {
     // Stop the SparkContext
     if (SparkSQLEnv.sparkContext != null) {
       sparkContext.stop()
+      sparkSession.close()
       sparkContext = null
       sqlContext = null
     }

sql/hive/src/main/scala/org/apache/spark/sql/hive/HiveExternalCatalog.scala

@@ -1267,6 +1267,11 @@ private[spark] class HiveExternalCatalog(conf: SparkConf, hadoopConf: Configurat
     client.listFunctions(db, pattern)
   }
 
+  override def close(): Unit = {
+    super.close()
+    client.close()
+  }
+
 }
 
 object HiveExternalCatalog {

sql/hive/src/main/scala/org/apache/spark/sql/hive/HiveSessionStateBuilder.scala

@@ -42,7 +42,7 @@ class HiveSessionStateBuilder(session: SparkSession, parentState: Option[Session
    * Create a Hive aware resource loader.
    */
   override protected lazy val resourceLoader: HiveSessionResourceLoader = {
-    val client: HiveClient = externalCatalog.client.newSession()
+    val client: HiveClient = externalCatalog.client
     new HiveSessionResourceLoader(session, client)
   }
 

sql/hive/src/main/scala/org/apache/spark/sql/hive/client/HiveClient.scala

@@ -278,4 +278,6 @@ private[hive] trait HiveClient {
   /** Used for testing only.  Removes all metadata from this instance of Hive. */
   def reset(): Unit
 
+  /** Close this client. */
+  def close(): Unit
 }

sql/hive/src/main/scala/org/apache/spark/sql/hive/client/HiveClientImpl.scala

@@ -838,6 +838,10 @@ private[hive] class HiveClientImpl(
         client.dropDatabase(db, true, false, true)
       }
   }
+
+  def close(): Unit = {
+    state.close()
+  }
 }
 
 private[hive] object HiveClientImpl {

sql/hive/src/main/scala/org/apache/spark/sql/hive/test/TestHive.scala

@@ -198,7 +198,7 @@ private[hive] class TestHiveSparkSession(
     new TestHiveSessionStateBuilder(this, parentSessionState).build()
   }
 
-  lazy val metadataHive: HiveClient = sharedState.externalCatalog.client.newSession()
+  lazy val metadataHive: HiveClient = sharedState.externalCatalog.client
 
   override def newSession(): TestHiveSparkSession = {
     new TestHiveSparkSession(sc, Some(sharedState), None, loadTestTables)

sql/hive/src/test/scala/org/apache/spark/sql/hive/client/VersionsSuite.scala

@@ -578,11 +578,6 @@ class VersionsSuite extends SparkFunSuite with Logging {
       client.setError(new PrintStream(new ByteArrayOutputStream()))
     }
 
-    test(s"$version: newSession") {
-      val newClient = client.newSession()
-      assert(newClient != null)
-    }
-
     test(s"$version: withHiveState and addJar") {
       val newClassPath = "."
       client.addJar(newClassPath)