From cca773a8056c4434d42afccdcb3343d27a2d8522 Mon Sep 17 00:00:00 2001
From: wangjunbo <wangjunbo@qiyi.com>
Date: Tue, 16 Jan 2024 17:08:03 +0800
Subject: [PATCH] fix code

---
 .../org/apache/kyuubi/engine/EngineRef.scala  |  4 ++--
 .../kyuubi/session/KyuubiSessionImpl.scala    |  2 +-
 .../sql/plan/command/DescribeEngine.scala     | 24 +++++++++++--------
 3 files changed, 17 insertions(+), 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
index f03546388d0..0c9d5ba6744 100644
--- 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
index 2aadab01fcc..4e4d261141f 100644
--- 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
index c91557fd550..5942164b446 100644
--- 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)
   }