From 586df4ee78faff5c8180abf4bde2c3f0bf4dc053 Mon Sep 17 00:00:00 2001
From: "xu.wenchun" <xu.wenchun@immomo.com>
Date: Mon, 13 Mar 2017 01:50:41 +0800
Subject: [PATCH] Support the --hivevar in the spark-sql command line and in
 the beeline

---
 .../hive/thriftserver/server/SparkSQLOperationManager.scala | 4 ++++
 sql/hive-thriftserver/src/test/resources/TestSQL.sql        | 1 +
 .../org/apache/spark/sql/hive/thriftserver/CliSuite.scala   | 6 ++++++
 3 files changed, 11 insertions(+)
 create mode 100644 sql/hive-thriftserver/src/test/resources/TestSQL.sql

diff --git a/sql/hive-thriftserver/src/main/scala/org/apache/spark/sql/hive/thriftserver/server/SparkSQLOperationManager.scala b/sql/hive-thriftserver/src/main/scala/org/apache/spark/sql/hive/thriftserver/server/SparkSQLOperationManager.scala
index 49ab66400934..c2439114408c 100644
--- a/sql/hive-thriftserver/src/main/scala/org/apache/spark/sql/hive/thriftserver/server/SparkSQLOperationManager.scala
+++ b/sql/hive-thriftserver/src/main/scala/org/apache/spark/sql/hive/thriftserver/server/SparkSQLOperationManager.scala
@@ -20,6 +20,8 @@ package org.apache.spark.sql.hive.thriftserver.server
 import java.util.{Map => JMap}
 import java.util.concurrent.ConcurrentHashMap
 
+import scala.collection.JavaConverters._
+
 import org.apache.hive.service.cli._
 import org.apache.hive.service.cli.operation.{ExecuteStatementOperation, Operation, OperationManager}
 import org.apache.hive.service.cli.session.HiveSession
@@ -50,6 +52,8 @@ private[thriftserver] class SparkSQLOperationManager()
     require(sqlContext != null, s"Session handle: ${parentSession.getSessionHandle} has not been" +
       s" initialized or had already closed.")
     val sessionState = sqlContext.sessionState.asInstanceOf[HiveSessionState]
+    parentSession.getSessionState.getHiveVariables.asScala.foreach(kv =>
+      sqlContext.conf.setConfString(kv._1, kv._2))
     val runInBackground = async && sessionState.hiveThriftServerAsync
     val operation = new SparkExecuteStatementOperation(parentSession, statement, confOverlay,
       runInBackground)(sqlContext, sessionToActivePool)
diff --git a/sql/hive-thriftserver/src/test/resources/TestSQL.sql b/sql/hive-thriftserver/src/test/resources/TestSQL.sql
new file mode 100644
index 000000000000..59c5c80afa8d
--- /dev/null
+++ b/sql/hive-thriftserver/src/test/resources/TestSQL.sql
@@ -0,0 +1 @@
+select date_add('${today}', -1);
\ No newline at end of file
diff --git a/sql/hive-thriftserver/src/test/scala/org/apache/spark/sql/hive/thriftserver/CliSuite.scala b/sql/hive-thriftserver/src/test/scala/org/apache/spark/sql/hive/thriftserver/CliSuite.scala
index d3cec11bd756..4bbaa4fbb140 100644
--- a/sql/hive-thriftserver/src/test/scala/org/apache/spark/sql/hive/thriftserver/CliSuite.scala
+++ b/sql/hive-thriftserver/src/test/scala/org/apache/spark/sql/hive/thriftserver/CliSuite.scala
@@ -283,4 +283,10 @@ class CliSuite extends SparkFunSuite with BeforeAndAfterAll with Logging {
       "SET conf3;" -> "conftest"
     )
   }
+
+  test("apply hivevar from cli command") {
+    val sqlFile = Thread.currentThread().getContextClassLoader.getResource("TestSQL.sql")
+    runCliWithin(2.minute, Seq("--hivevar", "today=2017-3-13", "-f",
+      sqlFile.getFile))(s"" -> "today=2017-3-13")
+  }
 }