Fix AppState when Engine connection is terminated

kyuubi-server/src/main/scala/org/apache/kyuubi/operation/BatchJobSubmission.scala

@@ -154,13 +154,20 @@ class BatchJobSubmission(
  engineId = appInfo.id,
  engineName = appInfo.name,
  engineUrl = appInfo.url.orNull,
- engineState = appInfo.state.toString,
+ engineState = getAppState(state, appInfo.state).toString,
  engineError = appInfo.error,
  endTime = endTime)
  session.sessionManager.updateMetadata(metadataToUpdate)
  }
  }
 
+ private def getAppState(operState: OperationState, appState: ApplicationState): Unit = {
+ if (state == ERROR && appState != ApplicationState.FAILED) {
+ ApplicationState.UNKNOWN
+ }
+ appState
+ }
+
  override def getOperationLog: Option[OperationLog] = Option(_operationLog)
 
  // we can not set to other state if it is canceled

kyuubi-server/src/test/scala/org/apache/kyuubi/WithKyuubiServerOnYarn.scala

@@ -193,20 +193,23 @@ class KyuubiOperationYarnClusterSuite extends WithKyuubiServerOnYarn with HiveJD
  }
 
  test("fast fail the kyuubi connection on engine terminated") {
+ val batchId = UUID.randomUUID().toString
  withSessionConf(Map.empty)(Map(
  "spark.master" -> "yarn",
  "spark.submit.deployMode" -> "cluster",
  "spark.sql.defaultCatalog=spark_catalog" -> "spark_catalog",
  "spark.sql.catalog.spark_catalog.type" -> "invalid_type",
  ENGINE_INIT_TIMEOUT.key -> "PT10M",
- KYUUBI_BATCH_ID_KEY -> UUID.randomUUID().toString))(Map.empty) {
+ KYUUBI_BATCH_ID_KEY -> batchId))(Map.empty) {
  val startTime = System.currentTimeMillis()
  val exception = intercept[Exception] {
  withJdbcStatement() { _ => }
  }
  val elapsedTime = System.currentTimeMillis() - startTime
  assert(elapsedTime < 60 * 1000)
  assert(exception.getMessage contains "Could not open client transport with JDBC Uri")
+ assert(sessionManager.getBatchMetadata(batchId).map(_.state).contains("ERROR"))
+ assert(sessionManager.getBatchMetadata(batchId).map(_.appState).contains("UNKNOWN"))
  }
  }
 }