[KYUUBI #5797] Support to describe engine and restart engine with command in current session

kyuubi-server/src/main/antlr4/org/apache/kyuubi/sql/KyuubiSqlBaseLexer.g4

@@ -51,6 +51,10 @@ KYUUBIADMIN: 'KYUUBIADMIN';
 
 SESSION: 'SESSION';
 
+ENGINE: 'ENGINE';
+
+RESTART: 'RESTART';
+
 BACKQUOTED_IDENTIFIER
     : '`' ( ~'`' | '``' )* '`'
     ;

kyuubi-server/src/main/antlr4/org/apache/kyuubi/sql/KyuubiSqlBaseParser.g4

@@ -30,4 +30,6 @@ statement
 
 runnableCommand
     : (DESC | DESCRIBE) SESSION                                     #describeSession
+    | (DESC | DESCRIBE) ENGINE                                      #describeEngine
+    | RESTART ENGINE                                                #restartEngine
     ;

kyuubi-server/src/main/scala/org/apache/kyuubi/client/KyuubiSyncThriftClient.scala

@@ -42,6 +42,8 @@ import org.apache.kyuubi.util.{ThreadUtils, ThriftUtils}
 import org.apache.kyuubi.util.ThreadUtils.scheduleTolerableRunnableWithFixedDelay
 
 class KyuubiSyncThriftClient private (
+    host: String,
+    port: Int,
     protocol: TProtocol,
     engineAliveProbeProtocol: Option[TProtocol],
     engineAliveProbeInterval: Long,
@@ -77,6 +79,9 @@ class KyuubiSyncThriftClient private (
   @VisibleForTesting
   private[kyuubi] def getEngineAliveProbeProtocol: Option[TProtocol] = engineAliveProbeProtocol
 
+  private[kyuubi] def getHost: String = this.host
+  private[kyuubi] def getPort: Int = this.port
+
   private def shutdownAsyncRequestExecutor(): Unit = {
     Option(asyncRequestExecutor).filterNot(_.isShutdown).foreach(ThreadUtils.shutdown(_))
     asyncRequestInterrupted = true
@@ -483,6 +488,8 @@ private[kyuubi] object KyuubiSyncThriftClient extends Logging {
         None
       }
     new KyuubiSyncThriftClient(
+      host,
+      port,
       tProtocol,
       aliveProbeProtocol,
       aliveProbeInterval,

kyuubi-server/src/main/scala/org/apache/kyuubi/engine/EngineRef.scala

@@ -86,7 +86,7 @@ private[kyuubi] class EngineRef(
   // a small amount of time for timeout
   private val LOCK_TIMEOUT_SPAN_FACTOR = if (Utils.isTesting) 0.5 else 0.1
 
-  private var builder: ProcBuilder = _
+  @volatile private var builder: ProcBuilder = _
 
   private[kyuubi] def getEngineRefId(): String = engineRefId
 

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

@@ -18,10 +18,11 @@
 package org.apache.kyuubi.session
 
 import java.util.Base64
+import java.util.concurrent.locks.ReentrantReadWriteLock
 
 import scala.collection.JavaConverters._
 
-import org.apache.kyuubi.KyuubiSQLException
+import org.apache.kyuubi.{KyuubiSQLException, Utils}
 import org.apache.kyuubi.client.KyuubiSyncThriftClient
 import org.apache.kyuubi.config.KyuubiConf
 import org.apache.kyuubi.config.KyuubiConf._
@@ -80,7 +81,7 @@ class KyuubiSessionImpl(
     handle.identifier.toString,
     sessionManager.applicationManager,
     sessionManager.engineStartupProcessSemaphore)
-  private[kyuubi] val launchEngineOp = sessionManager.operationManager
+  @volatile private[kyuubi] var launchEngineOp = sessionManager.operationManager
     .newLaunchEngineOperation(this, sessionConf.get(SESSION_ENGINE_LAUNCH_ASYNC))
 
   private lazy val sessionUserSignBase64: String =
@@ -101,7 +102,18 @@ class KyuubiSessionImpl(
   }
 
   @volatile private var _client: KyuubiSyncThriftClient = _
-  def client: KyuubiSyncThriftClient = _client
+
+  private val lock = new ReentrantReadWriteLock()
+
+  private def withReadLockAcquired[T](block: => T): T = Utils.withLockRequired(lock.readLock()) {
+    block
+  }
+
+  private def withWriteLockAcquired[T](block: => T): T = Utils.withLockRequired(lock.writeLock()) {
+    block
+  }
+
+  def client: KyuubiSyncThriftClient = withReadLockAcquired { _client }
 
   @volatile private var _engineSessionHandle: SessionHandle = _
 
@@ -119,6 +131,17 @@ class KyuubiSessionImpl(
     engineLastAlive = System.currentTimeMillis()
   }
 
+  def reLaunchEngine(): Unit = handleSessionException {
+    engineLaunched = false
+    withDiscoveryClient(sessionConf) { discoveryClient =>
+      engine.deregister(discoveryClient, (_client.getHost, _client.getPort))
+    }
+    launchEngineOp = sessionManager.operationManager
+      .newLaunchEngineOperation(this, sessionConf.get(SESSION_ENGINE_LAUNCH_ASYNC))
+    runOperation(launchEngineOp)
+    engineLastAlive = System.currentTimeMillis()
+  }
+
   private[kyuubi] def openEngineSession(extraEngineLog: Option[OperationLog] = None): Unit =
     handleSessionException {
       withDiscoveryClient(sessionConf) { discoveryClient =>
@@ -158,17 +181,22 @@ class KyuubiSessionImpl(
             }
 
           try {
-            val passwd =
+            val passwd = {
               if (sessionManager.getConf.get(ENGINE_SECURITY_ENABLED)) {
                 InternalSecurityAccessor.get().issueToken()
               } else {
                 Option(password).filter(_.nonEmpty).getOrElse("anonymous")
               }
-            _client = KyuubiSyncThriftClient.createClient(user, passwd, host, port, sessionConf)
-            _engineSessionHandle =
-              _client.openSession(protocol, user, passwd, openEngineSessionConf)
-            logSessionInfo(s"Connected to engine [$host:$port]/[${client.engineId.getOrElse("")}]" +
-              s" with ${_engineSessionHandle}]")
+            }
+            withWriteLockAcquired{
+              if (_client != null) _client.closeSession()
+              _client = KyuubiSyncThriftClient.createClient(user, passwd, host, port, sessionConf)
+              _engineSessionHandle =
+                _client.openSession(protocol, user, passwd, openEngineSessionConf)
+              logSessionInfo(
+                s"Connected to engine [$host:$port]/[${client.engineId.getOrElse("")}]" +
+                  s" with ${_engineSessionHandle}]")
+            }
             shouldRetry = false
           } catch {
             case e: TTransportException
@@ -199,7 +227,7 @@ class KyuubiSessionImpl(
             attempt += 1
             if (shouldRetry && _client != null) {
               try {
-                _client.closeSession()
+                withWriteLockAcquired{ _client.closeSession() }
               } catch {
                 case e: Throwable =>
                   warn(
@@ -269,7 +297,9 @@ class KyuubiSessionImpl(
     super.close()
     sessionManager.credentialsManager.removeSessionCredentialsEpoch(handle.identifier.toString)
     try {
-      if (_client != null) _client.closeSession()
+      if (_client != null) {
+        withWriteLockAcquired { _client.closeSession() }
+      }
     } finally {
       openSessionError.foreach { _ => if (engine != null) engine.close() }
       sessionEvent.endTime = System.currentTimeMillis()
@@ -314,6 +344,7 @@ class KyuubiSessionImpl(
 
   def checkEngineConnectionAlive(): Boolean = {
     try {
+      if (!engineLaunched) return true
       if (Option(client).exists(_.engineConnectionClosed)) return false
       if (!aliveProbeEnabled) return true
       getInfo(TGetInfoType.CLI_DBMS_VER)

kyuubi-server/src/main/scala/org/apache/kyuubi/sql/parser/server/KyuubiAstBuilder.scala

@@ -20,7 +20,7 @@ package org.apache.kyuubi.sql.parser.server
 import org.apache.kyuubi.sql.{KyuubiSqlBaseParser, KyuubiSqlBaseParserBaseVisitor}
 import org.apache.kyuubi.sql.KyuubiSqlBaseParser.SingleStatementContext
 import org.apache.kyuubi.sql.plan.{KyuubiTreeNode, PassThroughNode}
-import org.apache.kyuubi.sql.plan.command.{DescribeSession, RunnableCommand}
+import org.apache.kyuubi.sql.plan.command.{DescribeEngine, DescribeSession, RestartEngine, RunnableCommand}
 
 class KyuubiAstBuilder extends KyuubiSqlBaseParserBaseVisitor[AnyRef] {
 
@@ -44,4 +44,14 @@ class KyuubiAstBuilder extends KyuubiSqlBaseParserBaseVisitor[AnyRef] {
       : RunnableCommand = {
     DescribeSession()
   }
+
+  override def visitDescribeEngine(ctx: KyuubiSqlBaseParser.DescribeEngineContext)
+      : RunnableCommand = {
+    DescribeEngine()
+  }
+
+  override def visitRestartEngine(ctx: KyuubiSqlBaseParser.RestartEngineContext)
+      : RunnableCommand = {
+    RestartEngine()
+  }
 }

kyuubi-server/src/main/scala/org/apache/kyuubi/sql/plan/command/DescribeEngine.scala

@@ -0,0 +1,64 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.kyuubi.sql.plan.command
+
+import scala.collection.mutable.ListBuffer
+
+import org.apache.kyuubi.operation.IterableFetchIterator
+import org.apache.kyuubi.session.{KyuubiSession, KyuubiSessionImpl}
+import org.apache.kyuubi.shaded.hive.service.rpc.thrift.TTypeId
+import org.apache.kyuubi.sql.schema.{Column, Row, Schema}
+
+/**
+ * A runnable node for description the current session.
+ *
+ * The syntax of using this command in SQL is:
+ * {{{
+ *   [DESC|DESCRIBE] ENGINE;
+ * }}}
+ */
+case class DescribeEngine() extends RunnableCommand {
+
+  override def run(kyuubiSession: KyuubiSession): Unit = {
+    val rows = Seq(kyuubiSession).map { session =>
+      val client = session.asInstanceOf[KyuubiSessionImpl].client
+      val values = new ListBuffer[String]()
+      values += client.engineId.getOrElse("")
+      values += client.engineName.getOrElse("")
+      values += client.engineUrl.getOrElse("")
+      Row(values.toList)
+    }
+    iter = new IterableFetchIterator(rows)
+  }
+
+  override def resultSchema: Schema = {
+    Schema(DescribeEngine.outputCols().toList)
+  }
+
+  override def name(): String = "Describe Engine Node"
+}
+
+object DescribeEngine {
+
+  def outputCols(): Seq[Column] = {
+    Seq(
+      Column("ENGINE_ID", TTypeId.STRING_TYPE, Some("Kyuubi engine identify")),
+      Column("ENGINE_NAME", TTypeId.STRING_TYPE, Some("Kyuubi engine name")),
+      Column("ENGINE_URL", TTypeId.STRING_TYPE, Some("Kyuubi engine url")))
+  }
+}

kyuubi-server/src/main/scala/org/apache/kyuubi/sql/plan/command/DescribeSession.scala

@@ -56,8 +56,8 @@ object DescribeSession {
 
   def outputCols(): Seq[Column] = {
     Seq(
-      Column("id", TTypeId.STRING_TYPE, Some("Kyuubi session identify")),
-      Column("user", TTypeId.STRING_TYPE, Some("Kyuubi session user")),
-      Column("type", TTypeId.STRING_TYPE, Some("Kyuubi session type")))
+      Column("SESSION_ID", TTypeId.STRING_TYPE, Some("Kyuubi session identify")),
+      Column("SESSION_USER", TTypeId.STRING_TYPE, Some("Kyuubi session user")),
+      Column("SESSION_TYPE", TTypeId.STRING_TYPE, Some("Kyuubi session type")))
   }
 }

kyuubi-server/src/main/scala/org/apache/kyuubi/sql/plan/command/RestartEngine.scala

@@ -0,0 +1,64 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.kyuubi.sql.plan.command
+
+import scala.collection.mutable.ListBuffer
+
+import org.apache.kyuubi.operation.IterableFetchIterator
+import org.apache.kyuubi.session.{KyuubiSession, KyuubiSessionImpl}
+import org.apache.kyuubi.shaded.hive.service.rpc.thrift.TTypeId
+import org.apache.kyuubi.sql.schema.{Column, Row, Schema}
+
+/**
+ * A runnable node for description the current session.
+ *
+ * The syntax of using this command in SQL is:
+ * {{{
+ *   RESTART ENGINE;
+ * }}}
+ */
+case class RestartEngine() extends RunnableCommand {
+
+  override def run(kyuubiSession: KyuubiSession): Unit = {
+    val rows = Seq(kyuubiSession.asInstanceOf[KyuubiSessionImpl]).map { session =>
+      session.reLaunchEngine()
+      val client = session.client
+      val values = new ListBuffer[String]()
+      values += client.engineId.getOrElse("")
+      values += client.engineName.getOrElse("")
+      values += client.engineUrl.getOrElse("")
+      Row(values.toList)
+    }
+    iter = new IterableFetchIterator(rows)
+  }
+
+  override def resultSchema: Schema = {
+    Schema(RestartEngine.outputCols().toList)
+  }
+
+  override def name(): String = "Restart Session Engine Node"
+}
+
+object RestartEngine {
+  def outputCols(): Seq[Column] = {
+    Seq(
+      Column("ENGINE_ID", TTypeId.STRING_TYPE, Some("Kyuubi engine identify")),
+      Column("ENGINE_NAME", TTypeId.STRING_TYPE, Some("Kyuubi engine name")),
+      Column("ENGINE_URL", TTypeId.STRING_TYPE, Some("Kyuubi engine url")))
+  }
+}

kyuubi-server/src/test/scala/org/apache/kyuubi/operation/parser/DescribeEngineSuite.scala

@@ -0,0 +1,35 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.kyuubi.operation.parser
+
+class DescribeEngineSuite extends ExecutedCommandExecSuite {
+
+  test("desc/describe kyuubi engine") {
+    Seq("DESC", "DESCRIBE").foreach { desc =>
+      withJdbcStatement() { statement =>
+        val resultSet = statement.executeQuery(s"KYUUBI $desc ENGINE")
+        assert(resultSet.next())
+
+        assert(resultSet.getMetaData.getColumnCount == 3)
+        assert(resultSet.getMetaData.getColumnName(1) == "ENGINE_ID")
+        assert(resultSet.getMetaData.getColumnName(2) == "ENGINE_NAME")
+        assert(resultSet.getMetaData.getColumnName(3) == "ENGINE_URL")
+      }
+    }
+  }
+}

kyuubi-server/src/test/scala/org/apache/kyuubi/operation/parser/DescribeSessionSuite.scala

@@ -19,15 +19,17 @@ package org.apache.kyuubi.operation.parser
 
 class DescribeSessionSuite extends ExecutedCommandExecSuite {
 
-  test("desc kyuubi session") {
-    withJdbcStatement() { statement =>
-      val resultSet = statement.executeQuery("KYUUBI DESC SESSION")
-      assert(resultSet.next())
+  test("desc/describe kyuubi session") {
+    Seq("DESC", "DESCRIBE").foreach { desc =>
+      withJdbcStatement() { statement =>
+        val resultSet = statement.executeQuery(s"KYUUBI $desc SESSION")
+        assert(resultSet.next())
 
-      assert(resultSet.getMetaData.getColumnCount == 3)
-      assert(resultSet.getMetaData.getColumnName(1) == "id")
-      assert(resultSet.getMetaData.getColumnName(2) == "user")
-      assert(resultSet.getMetaData.getColumnName(3) == "type")
+        assert(resultSet.getMetaData.getColumnCount == 3)
+        assert(resultSet.getMetaData.getColumnName(1) == "SESSION_ID")
+        assert(resultSet.getMetaData.getColumnName(2) == "SESSION_USER")
+        assert(resultSet.getMetaData.getColumnName(3) == "SESSION_TYPE")
+      }
     }
   }
 }