[SPARK-23653][SQL] Show sql statement in spark SQL UI

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

@@ -635,6 +635,7 @@ class SparkSession private(
    * @since 2.0.0
    */
   def sql(sqlText: String): DataFrame = {
+    SQLExecution.setSqlText(sqlText)
     Dataset.ofRows(self, sessionState.sqlParser.parsePlan(sqlText))
   }
 

sql/core/src/main/scala/org/apache/spark/sql/execution/SQLExecution.scala

@@ -34,6 +34,16 @@ object SQLExecution {
 
   private val executionIdToQueryExecution = new ConcurrentHashMap[Long, QueryExecution]()
 
+  private val executionIdToSqlText = new ConcurrentHashMap[Long, String]()
+
+  def setSqlText(sqlText: String): Unit = {
+    executionIdToSqlText.putIfAbsent(_nextExecutionId.get(), sqlText)
+  }
+
+  def getSqlText(executionId: Long): String = {
+    executionIdToSqlText.get(executionId)
+  }
+
   def getQueryExecution(executionId: Long): QueryExecution = {
     executionIdToQueryExecution.get(executionId)
   }
@@ -63,16 +73,17 @@ object SQLExecution {
     val oldExecutionId = sc.getLocalProperty(EXECUTION_ID_KEY)
     val executionId = SQLExecution.nextExecutionId
     sc.setLocalProperty(EXECUTION_ID_KEY, executionId.toString)
+    val sqlText = getSqlText(executionId)
     executionIdToQueryExecution.put(executionId, queryExecution)
     try {
       // sparkContext.getCallSite() would first try to pick up any call site that was previously
       // set, then fall back to Utils.getCallSite(); call Utils.getCallSite() directly on
       // streaming queries would give us call site like "run at <unknown>:0"
       val callSite = sparkSession.sparkContext.getCallSite()
-
       sparkSession.sparkContext.listenerBus.post(SparkListenerSQLExecutionStart(
         executionId, callSite.shortForm, callSite.longForm, queryExecution.toString,
-        SparkPlanInfo.fromSparkPlan(queryExecution.executedPlan), System.currentTimeMillis()))
+        SparkPlanInfo.fromSparkPlan(queryExecution.executedPlan),
+        System.currentTimeMillis(), sqlText))
       try {
         body
       } finally {

sql/core/src/main/scala/org/apache/spark/sql/execution/ui/ExecutionPage.scala

@@ -78,6 +78,7 @@ class ExecutionPage(parent: SQLTab) extends WebUIPage("execution") with Logging
 
       summary ++
         planVisualization(metrics, graph) ++
+        showSQLText(executionUIData.sqlText) ++
         physicalPlanDescription(executionUIData.physicalPlanDescription)
     }.getOrElse {
       <div>No information to display for query {executionId}</div>
@@ -120,6 +121,25 @@ class ExecutionPage(parent: SQLTab) extends WebUIPage("execution") with Logging
   private def jobURL(jobId: Long): String =
     "%s/jobs/job?id=%s".format(UIUtils.prependBaseUri(parent.basePath), jobId)
 
+  private def showSQLText(sqlText: String): Seq[Node] = {
+    <div>
+      <span style="cursor: pointer;" onclick="clickShowSQLText();">
+        <span id="sql-text-arrow" class="arrow-closed"></span>
+        <a>SQL text</a>
+      </span>
+    </div>
+      <div id="sql-text-details" style="display: none;">
+        <pre>{sqlText}</pre>
+      </div>
+      <script>
+        function clickShowSQLText() {{
+        $('#sql-text-details').toggle();
+        $('#sql-text-arrow').toggleClass('arrow-open').toggleClass('arrow-closed');
+        }}
+      </script>
+        <br/>
+  }
+
   private def physicalPlanDescription(physicalPlanDescription: String): Seq[Node] = {
     <div>
       <span style="cursor: pointer;" onclick="clickPhysicalPlanDetails();">

sql/core/src/main/scala/org/apache/spark/sql/execution/ui/SQLAppStatusListener.scala

@@ -230,7 +230,7 @@ class SQLAppStatusListener(
 
   private def onExecutionStart(event: SparkListenerSQLExecutionStart): Unit = {
     val SparkListenerSQLExecutionStart(executionId, description, details,
-      physicalPlanDescription, sparkPlanInfo, time) = event
+      physicalPlanDescription, sparkPlanInfo, time, sqlText) = event
 
     def toStoredNodes(nodes: Seq[SparkPlanGraphNode]): Seq[SparkPlanGraphNodeWrapper] = {
       nodes.map {
@@ -265,6 +265,7 @@ class SQLAppStatusListener(
     exec.physicalPlanDescription = physicalPlanDescription
     exec.metrics = sqlPlanMetrics
     exec.submissionTime = time
+    exec.sqlText = sqlText
     update(exec)
   }
 
@@ -351,6 +352,7 @@ private class LiveExecutionData(val executionId: Long) extends LiveEntity {
   var jobs = Map[Int, JobExecutionStatus]()
   var stages = Set[Int]()
   var driverAccumUpdates = Map[Long, Long]()
+  var sqlText: String = null
 
   @volatile var metricsValues: Map[Long, String] = null
 
@@ -369,7 +371,8 @@ private class LiveExecutionData(val executionId: Long) extends LiveEntity {
       completionTime,
       jobs,
       stages,
-      metricsValues)
+      metricsValues,
+      sqlText)
   }
 
 }

sql/core/src/main/scala/org/apache/spark/sql/execution/ui/SQLAppStatusStore.scala

@@ -91,7 +91,8 @@ class SQLExecutionUIData(
      * from the SQL listener instance.
      */
     @JsonDeserialize(keyAs = classOf[JLong])
-    val metricValues: Map[Long, String]) {
+    val metricValues: Map[Long, String],
+    val sqlText: String) {
 
   @JsonIgnore @KVIndex("completionTime")
   private def completionTimeIndex: Long = completionTime.map(_.getTime).getOrElse(-1L)

sql/core/src/main/scala/org/apache/spark/sql/execution/ui/SQLListener.scala

@@ -34,7 +34,8 @@ case class SparkListenerSQLExecutionStart(
     details: String,
     physicalPlanDescription: String,
     sparkPlanInfo: SparkPlanInfo,
-    time: Long)
+    time: Long,
+    sqlText: String)
   extends SparkListenerEvent
 
 @DeveloperApi

sql/core/src/test/scala/org/apache/spark/sql/execution/SQLJsonProtocolSuite.scala

@@ -41,12 +41,13 @@ class SQLJsonProtocolSuite extends SparkFunSuite {
         |    "metadata":{},
         |    "metrics":[]
         |  },
-        |  "time":0
+        |  "time":0,
+        |  "sqlText":"select 1 as a"
         |}
       """.stripMargin
     val reconstructedEvent = JsonProtocol.sparkEventFromJson(parse(SQLExecutionStartJsonString))
     val expectedEvent = SparkListenerSQLExecutionStart(0, "test desc", "test detail", "test plan",
-      new SparkPlanInfo("TestNode", "test string", Nil, Nil), 0)
+      new SparkPlanInfo("TestNode", "test string", Nil, Nil), 0, "select 1 as a")
     assert(reconstructedEvent == expectedEvent)
   }
 }

sql/core/src/test/scala/org/apache/spark/sql/execution/ui/SQLAppStatusListenerSuite.scala

@@ -166,7 +166,8 @@ class SQLAppStatusListenerSuite extends SparkFunSuite with SharedSQLContext with
       "test",
       df.queryExecution.toString,
       SparkPlanInfo.fromSparkPlan(df.queryExecution.executedPlan),
-      System.currentTimeMillis()))
+      System.currentTimeMillis(),
+      "select 1 as a"))
 
     listener.onJobStart(SparkListenerJobStart(
       jobId = 0,
@@ -298,7 +299,8 @@ class SQLAppStatusListenerSuite extends SparkFunSuite with SharedSQLContext with
       "test",
       df.queryExecution.toString,
       SparkPlanInfo.fromSparkPlan(df.queryExecution.executedPlan),
-      System.currentTimeMillis()))
+      System.currentTimeMillis(),
+      "select 1 as a"))
     listener.onJobStart(SparkListenerJobStart(
       jobId = 0,
       time = System.currentTimeMillis(),
@@ -327,7 +329,8 @@ class SQLAppStatusListenerSuite extends SparkFunSuite with SharedSQLContext with
       "test",
       df.queryExecution.toString,
       SparkPlanInfo.fromSparkPlan(df.queryExecution.executedPlan),
-      System.currentTimeMillis()))
+      System.currentTimeMillis(),
+      "select 1 as a"))
     listener.onJobStart(SparkListenerJobStart(
       jobId = 0,
       time = System.currentTimeMillis(),
@@ -367,7 +370,8 @@ class SQLAppStatusListenerSuite extends SparkFunSuite with SharedSQLContext with
       "test",
       df.queryExecution.toString,
       SparkPlanInfo.fromSparkPlan(df.queryExecution.executedPlan),
-      System.currentTimeMillis()))
+      System.currentTimeMillis(),
+      "select 1 as a"))
     listener.onJobStart(SparkListenerJobStart(
       jobId = 0,
       time = System.currentTimeMillis(),
@@ -396,7 +400,8 @@ class SQLAppStatusListenerSuite extends SparkFunSuite with SharedSQLContext with
       "test",
       df.queryExecution.toString,
       SparkPlanInfo.fromSparkPlan(df.queryExecution.executedPlan),
-      System.currentTimeMillis()))
+      System.currentTimeMillis(),
+      "select 1 as a"))
 
     var stageId = 0
     def twoStageJob(jobId: Int): Unit = {
@@ -521,21 +526,24 @@ class SQLAppStatusListenerSuite extends SparkFunSuite with SharedSQLContext with
     val df = createTestDataFrame
     // Start execution 1 and execution 2
     time += 1
+    val sqlText = "select 1 as a"
     listener.onOtherEvent(SparkListenerSQLExecutionStart(
       1,
       "test",
       "test",
       df.queryExecution.toString,
       SparkPlanInfo.fromSparkPlan(df.queryExecution.executedPlan),
-      time))
+      time,
+      sqlText))
     time += 1
     listener.onOtherEvent(SparkListenerSQLExecutionStart(
       2,
       "test",
       "test",
       df.queryExecution.toString,
       SparkPlanInfo.fromSparkPlan(df.queryExecution.executedPlan),
-      time))
+      time,
+      sqlText))
 
     // Stop execution 2 before execution 1
     time += 1
@@ -551,7 +559,8 @@ class SQLAppStatusListenerSuite extends SparkFunSuite with SharedSQLContext with
       "test",
       df.queryExecution.toString,
       SparkPlanInfo.fromSparkPlan(df.queryExecution.executedPlan),
-      time))
+      time,
+      sqlText))
     assert(statusStore.executionsCount === 2)
     assert(statusStore.execution(2) === None)
   }