KYUUBI_SESSION_CONNECTION_URL_KEY

Author: turboFei

kyuubi-server/src/main/scala/org/apache/kyuubi/session/KyuubiBatchSession.scala

@@ -23,7 +23,7 @@ import scala.collection.JavaConverters._
 
 import org.apache.kyuubi.client.util.BatchUtils._
 import org.apache.kyuubi.config.{KyuubiConf, KyuubiReservedKeys}
-import org.apache.kyuubi.config.KyuubiReservedKeys.KYUUBI_BATCH_PRIORITY
+import org.apache.kyuubi.config.KyuubiReservedKeys.{KYUUBI_BATCH_PRIORITY, KYUUBI_SESSION_CONNECTION_URL_KEY}
 import org.apache.kyuubi.engine.KyuubiApplicationManager
 import org.apache.kyuubi.engine.spark.SparkProcessBuilder
 import org.apache.kyuubi.events.{EventBus, KyuubiSessionEvent}
@@ -54,6 +54,8 @@ class KyuubiBatchSession(
     conf,
     sessionManager) {
   override val sessionType: SessionType = SessionType.BATCH
+  override val connectionUrl: String =
+    metadata.map(_.kyuubiInstance).getOrElse(conf.getOrElse(KYUUBI_SESSION_CONNECTION_URL_KEY, ""))
 
   override val handle: SessionHandle = {
     val batchId = metadata.map(_.identifier).getOrElse(conf(KYUUBI_BATCH_ID_KEY))

kyuubi-server/src/main/scala/org/apache/kyuubi/session/KyuubiSessionManager.scala

@@ -312,18 +312,7 @@ class KyuubiSessionManager private (name: String) extends SessionManager(name) {
         kyuubiInstance,
         0,
         Int.MaxValue).map { metadata =>
-        createBatchSession(
-          metadata.username,
-          "anonymous",
-          metadata.ipAddress,
-          metadata.requestConf,
-          metadata.engineType,
-          Option(metadata.requestName),
-          metadata.resource,
-          metadata.className,
-          metadata.requestArgs,
-          Some(metadata),
-          fromRecovery = true)
+        createBatchSessionFromRecovery(metadata)
       }).getOrElse(Seq.empty)
     }
   }
@@ -341,25 +330,26 @@ class KyuubiSessionManager private (name: String) extends SessionManager(name) {
           getBatchMetadata(batchId)
             .filter(m =>
               m.kyuubiInstance == kyuubiInstance && batchStatesToRecovery.contains(m.opState))
-            .flatMap { metadata =>
-              Some(
-                createBatchSession(
-                  metadata.username,
-                  "anonymous",
-                  metadata.ipAddress,
-                  metadata.requestConf,
-                  metadata.engineType,
-                  Option(metadata.requestName),
-                  metadata.resource,
-                  metadata.className,
-                  metadata.requestArgs,
-                  Some(metadata),
-                  fromRecovery = true))
-            }
+            .flatMap { metadata => Some(createBatchSessionFromRecovery(metadata)) }
       }
     }
   }
 
+  private def createBatchSessionFromRecovery(metadata: Metadata): KyuubiBatchSession = {
+    createBatchSession(
+      metadata.username,
+      "anonymous",
+      metadata.ipAddress,
+      metadata.requestConf,
+      metadata.engineType,
+      Option(metadata.requestName),
+      metadata.resource,
+      metadata.className,
+      metadata.requestArgs,
+      Some(metadata),
+      fromRecovery = true)
+  }
+
   def reassignBatchSessions(
       kyuubiInstance: String,
       newKyuubiInstance: String): Seq[String] = {

kyuubi-server/src/test/scala/org/apache/kyuubi/server/api/v1/BatchesResourceSuite.scala

@@ -714,12 +714,12 @@ abstract class BatchesResourceSuiteBase extends KyuubiFunSuite
     assert(session2.createTime === batchMetadata2.createTime)
 
     eventually(timeout(10.seconds)) {
-      assert(session1.batchJobSubmissionOp.getStatus.state === OperationState.RUNNING ||
-        session1.batchJobSubmissionOp.getStatus.state === OperationState.FINISHED)
+      val batch1State = session1.batchJobSubmissionOp.getStatus.state
+      assert(batch1State === OperationState.RUNNING || OperationState.isTerminal(batch1State))
       assert(session1.batchJobSubmissionOp.builder.processLaunched)
 
-      assert(session2.batchJobSubmissionOp.getStatus.state === OperationState.RUNNING ||
-        session2.batchJobSubmissionOp.getStatus.state === OperationState.FINISHED)
+      val batch2State = session2.batchJobSubmissionOp.getStatus.state
+      assert(batch2State === OperationState.RUNNING || OperationState.isTerminal(batch2State))
       assert(!session2.batchJobSubmissionOp.builder.processLaunched)
     }