fix code

apache · Jan 16, 2024 · cca773a · cca773a
1 parent 9fcc2c6
commit cca773a
Show file tree

Hide file tree

Showing 3 changed files with 17 additions and 13 deletions.
diff --git a/kyuubi-server/src/main/scala/org/apache/kyuubi/engine/EngineRef.scala b/kyuubi-server/src/main/scala/org/apache/kyuubi/engine/EngineRef.scala
@@ -339,10 +339,10 @@ private[kyuubi] class EngineRef(
 
   def getServiceNodes(
       discoveryClient: DiscoveryClient,
-      hostPort: (String, Int)): Seq[ServiceNodeInfo] = {
+      hostPort: (String, Int)): Option[ServiceNodeInfo] = {
     tryWithLock(discoveryClient) {
       val serviceNodes = discoveryClient.getServiceNodesInfo(engineSpace)
-      serviceNodes.filter { sn => (sn.host, sn.port) == hostPort }
+      serviceNodes.filter { sn => (sn.host, sn.port) == hostPort }.headOption
     }
   }
 

diff --git a/kyuubi-server/src/main/scala/org/apache/kyuubi/session/KyuubiSessionImpl.scala b/kyuubi-server/src/main/scala/org/apache/kyuubi/session/KyuubiSessionImpl.scala
@@ -120,7 +120,7 @@ class KyuubiSessionImpl(
     engineLastAlive = System.currentTimeMillis()
   }
 
-  def listZkEngineNodes: Seq[ServiceNodeInfo] = {
+  def getEngineNode: Option[ServiceNodeInfo] = {
     withDiscoveryClient(sessionConf) { discoveryClient =>
       engine.getServiceNodes(discoveryClient, _client.engineHostPort)
     }

diff --git a/kyuubi-server/src/main/scala/org/apache/kyuubi/sql/plan/command/DescribeEngine.scala b/kyuubi-server/src/main/scala/org/apache/kyuubi/sql/plan/command/DescribeEngine.scala
@@ -35,18 +35,22 @@ import org.apache.kyuubi.sql.schema.{Column, Row, Schema}
 case class DescribeEngine() extends RunnableCommand {
 
   override def run(kyuubiSession: KyuubiSession): Unit = {
-    val rows = Seq(kyuubiSession.asInstanceOf[KyuubiSessionImpl]).flatMap { session =>
+    val rows = Seq(kyuubiSession.asInstanceOf[KyuubiSessionImpl]).map { session =>
       lazy val client = session.client
-      session.listZkEngineNodes.map { nodeInfo =>
-        val values = new ListBuffer[String]()
-        values += client.engineId.getOrElse("")
-        values += client.engineName.getOrElse("")
-        values += client.engineUrl.getOrElse("")
-        values += s"${nodeInfo.host}:${nodeInfo.port}"
-        values += nodeInfo.version.getOrElse("")
-        values += nodeInfo.attributes.filter(_._1.contains("memory")).mkString(";")
-        Row(values.toList)
+      val values = new ListBuffer[String]()
+      values += client.engineId.getOrElse("")
+      values += client.engineName.getOrElse("")
+      values += client.engineUrl.getOrElse("")
+      session.getEngineNode match {
+        case Some(nodeInfo) =>
+          values += s"${nodeInfo.host}:${nodeInfo.port}"
+          values += nodeInfo.version.getOrElse("")
+          values += nodeInfo.attributes.mkString(",")
+        case None =>
+          values += ("", "", "")
       }
+
+      Row(values.toList)
     }
     iter = new IterableFetchIterator(rows)
   }