[CARMEL-6490] Lots of unfinished state sql in the query log (#1199)

* [CARMEL-6490] Lots of unfinished state sql in the query log

* Fix ut

Author: allenma

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

@@ -47,6 +47,7 @@ import org.apache.spark.sql.execution.adaptive.AdaptiveSparkPlanExec
 import org.apache.spark.sql.execution.datasources.LogicalRelation
 import org.apache.spark.sql.execution.exchange.ShuffleExchangeExec
 import org.apache.spark.sql.execution.ui.StatementStart
+import org.apache.spark.sql.hive.thriftserver.HiveThriftServer2.ExecutionState
 import org.apache.spark.sql.hive.thriftserver.errors.QueryLevelRestrictionErrors
 import org.apache.spark.sql.hive.thriftserver.ql.{QueryLogExtInfo, QueryLogObjectList}
 import org.apache.spark.sql.hive.thriftserver.ui.SparkListenerThriftServerQueryExit
@@ -135,6 +136,8 @@ private[hive] class SparkDownloadDataOperation(
   private var schemaStr: String = _
   private var totalDataSize: Long = 0
 
+  @volatile private var failDetails: String = ""
+
   override def close(): Unit = {
     // CARMEL-4662 Fix Download query state is incorrect.
     if (getStatus.getState eq OperationState.FINISHED) {
@@ -144,12 +147,25 @@ private[hive] class SparkDownloadDataOperation(
       SparkListenerThriftServerQueryExit(
         statementId,
         QueryLogObjectList(Option(result).map(_.queryExecution)),
-        QueryLogExtInfo(false, totalDataSize)))
+        QueryLogExtInfo(
+          finalQueryStateId,
+          failDetails,
+          System.currentTimeMillis(),
+          false, totalDataSize)))
     logInfo(s"CLOSING $statementId")
     cleanup(OperationState.CLOSED)
     sqlContext.sparkContext.clearJobGroup()
   }
 
+  private def finalQueryStateId: Int = {
+    getStatus.getState match {
+      case OperationState.ERROR => ExecutionState.FAILED.id
+      case OperationState.FINISHED => ExecutionState.FINISHED.id
+      case OperationState.CANCELED => ExecutionState.CANCELED.id
+      case _ => ExecutionState.FINISHED.id
+    }
+  }
+
   override def runInternal(): Unit = {
     setState(OperationState.PENDING)
     setHasResultSet(true)
@@ -301,8 +317,9 @@ private[hive] class SparkDownloadDataOperation(
       case NonFatal(e) =>
         logQueryError(s"Error executing query [$statementId]", e)
         setState(OperationState.ERROR)
+        failDetails = Utils.findFirstCause(e).toString
         HiveThriftServer2.eventManager.onStatementError(
-          statementId, Utils.findFirstCause(e).toString, Utils.exceptionString(e))
+          statementId, failDetails, Utils.exceptionString(e))
         val exception = new HiveSQLException(e)
         setOperationException(exception)
     } finally {
@@ -592,16 +609,20 @@ private[hive] class SparkDownloadDataOperation(
   }
 
   override def cancel(): Unit = {
+    failDetails = s"Canceling operation, stack trace:\n" +
+      Thread.currentThread().getStackTrace.mkString("\n")
     if (statementId != null) {
-      HiveThriftServer2.eventManager.onStatementCanceled(statementId,
-        s"Canceling operation, stack trace:\n"
-          + Thread.currentThread().getStackTrace.mkString("\n"))
+      HiveThriftServer2.eventManager.onStatementCanceled(statementId, failDetails)
     }
     HiveThriftServer2.listener.onQueryExit(
       SparkListenerThriftServerQueryExit(
         statementId,
         QueryLogObjectList(Option(result).map(_.queryExecution)),
-        QueryLogExtInfo(false, totalDataSize)))
+        QueryLogExtInfo(
+          ExecutionState.CANCELED.id,
+          failDetails,
+          System.currentTimeMillis(),
+          false, totalDataSize)))
     cleanup(OperationState.CANCELED)
   }
 

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

@@ -48,6 +48,7 @@ import org.apache.spark.sql.execution.datasources.{InsertIntoDataSourceCommand,
 import org.apache.spark.sql.execution.datasources.v2.{AppendDataExecV1, OverwriteByExpressionExecV1, V2TableWriteExec}
 import org.apache.spark.sql.execution.metric.SQLMetric
 import org.apache.spark.sql.execution.ui.StatementStart
+import org.apache.spark.sql.hive.thriftserver.HiveThriftServer2.ExecutionState
 import org.apache.spark.sql.hive.thriftserver.ql.{QueryLogExtInfo, QueryLogObjectList}
 import org.apache.spark.sql.hive.thriftserver.server.ResultShareGroup
 import org.apache.spark.sql.hive.thriftserver.ui.SparkListenerThriftServerQueryExit
@@ -83,6 +84,8 @@ private[hive] class SparkExecuteStatementOperation(
   private var iter: Iterator[SparkRow] = _
   private var dataTypes: Array[DataType] = _
 
+  @volatile private var failDetails: String = "" // store the query failed reason
+
   private lazy val resultSchema: TableSchema = {
     if (result == null || result.schema.isEmpty) {
       new TableSchema(Arrays.asList(new FieldSchema("Result", "string", "")))
@@ -371,17 +374,19 @@ private[hive] class SparkExecuteStatementOperation(
           e match {
             case hiveException: HiveSQLException =>
               setErrorStatus(OperationState.ERROR, hiveException)
+              failDetails = SparkUtils.findFirstCause(e).toString
               HiveThriftServer2.eventManager.onStatementError(
-                statementId, SparkUtils.findFirstCause(e).toString,
+                statementId, failDetails,
                 SparkUtils.exceptionString(hiveException))
               throw hiveException
             case _ =>
               val root = ExceptionUtils.getRootCause(e)
               val hiveException =
                 new HiveSQLException("Error running query: " + root.toString, root)
+              failDetails = SparkUtils.findFirstCause(e).toString
               setErrorStatus(OperationState.ERROR, hiveException)
               HiveThriftServer2.eventManager.onStatementError(
-                statementId, SparkUtils.findFirstCause(e).toString,
+                statementId, failDetails,
                 SparkUtils.exceptionString(root))
               throw hiveException
           }
@@ -398,12 +403,25 @@ private[hive] class SparkExecuteStatementOperation(
           SparkListenerThriftServerQueryExit(
             statementId,
             QueryLogObjectList(Option(result).map(_.queryExecution)),
-            QueryLogExtInfo(isSpill, finalRowsLength)))
+            QueryLogExtInfo(
+              finalQueryStateId,
+              failDetails,
+              System.currentTimeMillis(),
+              isSpill, finalRowsLength)))
       }
       sqlContext.sparkContext.clearJobGroup()
     }
   }
 
+  private def finalQueryStateId: Int = {
+    getStatus.getState match {
+      case OperationState.ERROR => ExecutionState.FAILED.id
+      case OperationState.FINISHED => ExecutionState.FINISHED.id
+      case OperationState.CANCELED => ExecutionState.CANCELED.id
+      case _ => ExecutionState.FINISHED.id
+    }
+  }
+
   private def withRetry[T](f: => T): T = {
     val maxRetry = 1
     var retryNum = 0
@@ -594,9 +612,9 @@ private[hive] class SparkExecuteStatementOperation(
         // 'InterruptedException', which could be misleading, the query final state
         // should be canceled instead of failed. Hence we need guarantee that the
         // cancellation event is posted before the query execution is actually canceled.
-        HiveThriftServer2.eventManager.onStatementCanceled(statementId,
-          s"Canceling operation, stack trace:\n"
-            + Thread.currentThread().getStackTrace.mkString("\n"))
+        failDetails = "Canceling operation, stack trace:\n" +
+          Thread.currentThread().getStackTrace.mkString("\n")
+        HiveThriftServer2.eventManager.onStatementCanceled(statementId, failDetails)
         cleanup()
       }
     }

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

@@ -129,9 +129,9 @@ private[thriftserver] class KafkaSink(
         execInfo.sessionId,
         execInfo.userName,
         sdf.format(execInfo.startTimestamp),
-        String.valueOf(execInfo.finishTimestamp - execInfo.startTimestamp),
-        String.valueOf(execInfo.state.id),
-        execInfo.detail,
+        String.valueOf(extInfo.finishedTime - execInfo.startTimestamp),
+        String.valueOf(extInfo.stateId),
+        extInfo.detail,
         execInfo.statement,
         queue,
         String.valueOf(execInfo.currentTasksTime),

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

@@ -164,4 +164,10 @@ case class QueryLogObject(
     other.db.equals(db) && other.table.equals(table)
 }
 
-case class QueryLogExtInfo(isSpill: Boolean, rowCount: Long)
+case class QueryLogExtInfo(
+  stateId: Int,
+  detail: String,
+  finishedTime: Long,
+  isSpill: Boolean,
+  rowCount: Long)
+

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

@@ -74,11 +74,15 @@ class QueryLogSuite extends SparkFunSuite with SharedSparkSession {
     val sink = new MockQueryLogSink
     val ql = new QueryLogger(spark.sqlContext.conf, sink)
     ql.logQuery("query1", mockExecutionInfo, null,
-      QueryLogObjectList(Some(qe)), QueryLogExtInfo(true, -1))
+      QueryLogObjectList(Some(qe)), QueryLogExtInfo(
+        3, "", System.currentTimeMillis(),
+        true, -1))
     assert(sink.queryObjList.isEmpty)
     val qe2 = sql("explain select col1 from test_table").queryExecution
     ql.logQuery("query2", mockExecutionInfo, null,
-      QueryLogObjectList(Some(qe2)), QueryLogExtInfo(true, -1))
+      QueryLogObjectList(Some(qe2)), QueryLogExtInfo(
+        3, "", System.currentTimeMillis(),
+        true, -1))
     assert(sink.queryObjList.isEmpty)
     ql.logExceptionQuery("exception", "1", "hdmi-default", "1", mockExecutionInfo)
     ql.logRunningQuery(
@@ -99,12 +103,16 @@ class QueryLogSuite extends SparkFunSuite with SharedSparkSession {
           val sink = new MockQueryLogSink
           val ql = new QueryLogger(spark.sqlContext.conf, sink)
           ql.logQuery("query1", mockExecutionInfo, null,
-            QueryLogObjectList(Some(qe)), QueryLogExtInfo(true, -1))
+            QueryLogObjectList(Some(qe)), QueryLogExtInfo(
+              3, "", System.currentTimeMillis(),
+              true, -1))
           assert(sink.queryObjList.size == 1)
           assert(sink.queryObjList.map(_.table).toSet.equals(Set("test_table")))
           val qe2 = sql("insert into temp select * from test_table").queryExecution
           ql.logQuery("query2", mockExecutionInfo, null,
-            QueryLogObjectList(Some(qe2)), QueryLogExtInfo(true, -1))
+            QueryLogObjectList(Some(qe2)), QueryLogExtInfo(
+              3, "", System.currentTimeMillis(),
+              true, -1))
           assert(sink.queryObjList.map(_.table).toSet.equals(Set("test_table", "test_table")))
         }
       }
@@ -136,7 +144,9 @@ class QueryLogSuite extends SparkFunSuite with SharedSparkSession {
         val sink = new MockQueryLogSink
         val ql = new QueryLogger(spark.sqlContext.conf, sink)
         ql.logQuery("query1", mockExecutionInfo, null,
-          QueryLogObjectList(Some(qe)), QueryLogExtInfo(true, -1))
+          QueryLogObjectList(Some(qe)), QueryLogExtInfo(
+            3, "", System.currentTimeMillis(),
+            true, -1))
         assert(sink.queryObjList.size == 1)
         val q1 = QueryLogObject("default", "test_table", true)
         assert(sink.queryObjList.contains(q1))
@@ -151,7 +161,9 @@ class QueryLogSuite extends SparkFunSuite with SharedSparkSession {
         val sink = new MockQueryLogSink
         val ql = new QueryLogger(spark.sqlContext.conf, sink)
         ql.logQuery("query1", mockExecutionInfo, null,
-          QueryLogObjectList(Some(qe)), QueryLogExtInfo(true, -1))
+          QueryLogObjectList(Some(qe)), QueryLogExtInfo(
+            3, "", System.currentTimeMillis(),
+            true, -1))
         assert(sink.queryObjList.size == 2)
         val q1 = QueryLogObject("default", "test_table", false)
         val q2 = QueryLogObject("default", "view1", true)
@@ -168,7 +180,9 @@ class QueryLogSuite extends SparkFunSuite with SharedSparkSession {
         val sink = new MockQueryLogSink
         val ql = new QueryLogger(spark.sqlContext.conf, sink)
         ql.logQuery("query1", mockExecutionInfo, null,
-          QueryLogObjectList(Some(qe)), QueryLogExtInfo(true, -1))
+          QueryLogObjectList(Some(qe)), QueryLogExtInfo(
+            3, "", System.currentTimeMillis(),
+            true, -1))
         assert(sink.queryObjList.size == 2)
         val q1 = QueryLogObject("default", "test_table", false)
         val q2 = QueryLogObject("default", "view2", true)