addition of connection close in HBaseRelation

yzhou2001 · yzhou2001 · commit 065d974912ec · 2016-02-02T17:18:10.000-08:00
diff --git a/src/main/scala/org/apache/spark/sql/hbase/HBaseCatalog.scala b/src/main/scala/org/apache/spark/sql/hbase/HBaseCatalog.scala
@@ -181,6 +181,7 @@ private[hbase] class HBaseCatalog(@transient hbaseContext: SQLContext,
             """The directory of a certain regionserver is not accessible,
             |please add 'cd ~' before 'start regionserver' in your regionserver start script.""")
         }
+        metadataTable.close()
       } else {
         deploySuccessfully_internal = Some(true)
       }
diff --git a/src/main/scala/org/apache/spark/sql/hbase/HBaseRelation.scala b/src/main/scala/org/apache/spark/sql/hbase/HBaseRelation.scala
@@ -160,12 +160,18 @@ private[hbase] case class HBaseRelation(
   logger.debug(s"HBaseRelation config has zkPort="
     + s"${getConf.get("hbase.zookeeper.property.clientPort")}")
 
-  @transient lazy val connection_ =
+  @transient var internal_connection : Connection = _
+  
+  def connection_ = {
     if (connection == null) {
-      ConnectionFactory.createConnection(getConf)
+      if (internal_connection == null) {
+        internal_connection = ConnectionFactory.createConnection(getConf)
+      }
+      internal_connection
     } else {
       connection
     }
+  }
 
   @transient private var htable_ : Table = _
 
@@ -188,11 +194,15 @@ private[hbase] case class HBaseRelation(
     refs.indexWhere(_.exprId == partitionKeys(keyIndex).exprId)
   }
 
-  def closeHTable() = {
+  def close() = {
     if (htable_ != null) {
       htable_.close()
       htable_ = null
     }
+    if (connection == null && internal_connection != null) {
+      internal_connection.close()
+      internal_connection = null
+    }
   }
 
   // corresponding logical relation
@@ -712,7 +722,7 @@ private[hbase] case class HBaseRelation(
     if (puts.nonEmpty) {
       htable.put(puts.toList)
     }
-    closeHTable()
+    close()
   }
 
 
diff --git a/src/main/scala/org/apache/spark/sql/hbase/HBaseSQLReaderRDD.scala b/src/main/scala/org/apache/spark/sql/hbase/HBaseSQLReaderRDD.scala
@@ -194,7 +194,7 @@ class HBaseSQLReaderRDD(val relation: HBaseRelation,
       def close() = {
         try {
           scanner.close()
-          relation.closeHTable()
+          relation.close()
         } catch {
           case e: Exception => logWarning("Exception in scanner.close", e)
         }
diff --git a/src/main/scala/org/apache/spark/sql/hbase/SparkSqlRegionObserver.scala b/src/main/scala/org/apache/spark/sql/hbase/SparkSqlRegionObserver.scala
@@ -89,7 +89,7 @@ class HBaseCoprocessorSQLReaderRDD(var relation: HBaseRelation,
       def close() = {
         try {
           scanner.close()
-          relation.closeHTable()
+          relation.close()
         } catch {
           case e: Exception => logWarning("Exception in scanner.close", e)
         }
diff --git a/src/main/scala/org/apache/spark/sql/hbase/execution/hbaseCommands.scala b/src/main/scala/org/apache/spark/sql/hbase/execution/hbaseCommands.scala
@@ -284,7 +284,7 @@ case class BulkLoadIntoTableCommand(
           if (recordsWritten > 0) {
             load.doBulkLoad(targetPath, relation.connection_.getAdmin, relation.htable,
               relation.connection_.getRegionLocator(relation.hTableName))
-            relation.closeHTable()
+            relation.close()
           }
         }
         1
@@ -304,7 +304,7 @@ case class BulkLoadIntoTableCommand(
       load.doBulkLoad(tablePath, relation.connection_.getAdmin, relation.htable,
       relation.connection_.getRegionLocator(relation.hTableName))
     }
-    relation.closeHTable()
+    relation.close()
     logDebug(s"finish BulkLoad on table ${relation.htable.getName}:" +
       s" ${System.currentTimeMillis()}")
     Seq.empty[Row]
diff --git a/src/test/scala/org/apache/spark/sql/hbase/TestHbase.scala b/src/test/scala/org/apache/spark/sql/hbase/TestHbase.scala
@@ -29,8 +29,7 @@ object TestHbase {
   def hsc: HBaseSQLContext = {
     if (hsc_ == null) {
       hsc_ = new HBaseSQLContext(new SparkContext("local", "TestSQLContext", new SparkConf(true)
-        .set("spark.hadoop.hbase.zookeeper.quorum", "localhost")
-        .set("spark.hadoop.dfs.replication", "1")))
+        .set("spark.hadoop.hbase.zookeeper.quorum", "localhost")))
     }
     hsc_
   }
@@ -50,6 +49,7 @@ object TestHbase {
 
   def stop: Unit = {
     hsc_.catalog.stopAdmin()
+    hsc_.catalog.connection.close()
     hsc.sparkContext.stop()
     hsc_ = null
     testUtil.cleanupDataTestDirOnTestFS()

Original file line number	Diff line number	Diff line change
`@@ -181,6 +181,7 @@ private[hbase] class HBaseCatalog(@transient hbaseContext: SQLContext,`
`181`	`181`	`"""The directory of a certain regionserver is not accessible,`
`182`	`182`	`\|please add 'cd ~' before 'start regionserver' in your regionserver start script.""")`
`183`	`183`	`}`
	`184`	`+ metadataTable.close()`
`184`	`185`	`} else {`
`185`	`186`	`deploySuccessfully_internal = Some(true)`
`186`	`187`	`}`
Original file line number	Diff line number	Diff line change
`@@ -194,7 +194,7 @@ class HBaseSQLReaderRDD(val relation: HBaseRelation,`
`194`	`194`	`def close() = {`
`195`	`195`	`try {`
`196`	`196`	`scanner.close()`
`197`		`- relation.closeHTable()`
	`197`	`+ relation.close()`
`198`	`198`	`} catch {`
`199`	`199`	`case e: Exception => logWarning("Exception in scanner.close", e)`
`200`	`200`	`}`
Original file line number	Diff line number	Diff line change
`@@ -89,7 +89,7 @@ class HBaseCoprocessorSQLReaderRDD(var relation: HBaseRelation,`
`89`	`89`	`def close() = {`
`90`	`90`	`try {`
`91`	`91`	`scanner.close()`
`92`		`- relation.closeHTable()`
	`92`	`+ relation.close()`
`93`	`93`	`} catch {`
`94`	`94`	`case e: Exception => logWarning("Exception in scanner.close", e)`
`95`	`95`	`}`
Original file line number	Diff line number	Diff line change
`@@ -284,7 +284,7 @@ case class BulkLoadIntoTableCommand(`
`284`	`284`	`if (recordsWritten > 0) {`
`285`	`285`	`load.doBulkLoad(targetPath, relation.connection_.getAdmin, relation.htable,`
`286`	`286`	`relation.connection_.getRegionLocator(relation.hTableName))`
`287`		`- relation.closeHTable()`
	`287`	`+ relation.close()`
`288`	`288`	`}`
`289`	`289`	`}`
`290`	`290`	`1`
`@@ -304,7 +304,7 @@ case class BulkLoadIntoTableCommand(`
`304`	`304`	`load.doBulkLoad(tablePath, relation.connection_.getAdmin, relation.htable,`
`305`	`305`	`relation.connection_.getRegionLocator(relation.hTableName))`
`306`	`306`	`}`
`307`		`- relation.closeHTable()`
	`307`	`+ relation.close()`
`308`	`308`	`logDebug(s"finish BulkLoad on table ${relation.htable.getName}:" +`
`309`	`309`	`s" ${System.currentTimeMillis()}")`
`310`	`310`	`Seq.empty[Row]`