[JDBC] Adaptive call fetchLaunchEngineResult

jiaoqingbo · Sep 26, 2023 · ca9bd53 · ca9bd53
1 parent 8ef6ca3
commit ca9bd53
Show file tree

Hide file tree

Showing 3 changed files with 10 additions and 1 deletion.
diff --git a/kyuubi-common/src/main/scala/org/apache/kyuubi/config/KyuubiReservedKeys.scala b/kyuubi-common/src/main/scala/org/apache/kyuubi/config/KyuubiReservedKeys.scala
@@ -38,6 +38,8 @@ object KyuubiReservedKeys {
     "kyuubi.session.engine.launch.handle.guid"
   final val KYUUBI_SESSION_ENGINE_LAUNCH_HANDLE_SECRET =
     "kyuubi.session.engine.launch.handle.secret"
+  final val KYUUBI_SESSION_ENGINE_LAUNCH_SUPPORT_RESULT =
+    "kyuubi.session.engine.launch.support.result"
   final val KYUUBI_OPERATION_SET_CURRENT_CATALOG = "kyuubi.operation.set.current.catalog"
   final val KYUUBI_OPERATION_GET_CURRENT_CATALOG = "kyuubi.operation.get.current.catalog"
   final val KYUUBI_OPERATION_SET_CURRENT_DATABASE = "kyuubi.operation.set.current.database"

diff --git a/kyuubi-hive-jdbc/src/main/java/org/apache/kyuubi/jdbc/hive/KyuubiConnection.java b/kyuubi-hive-jdbc/src/main/java/org/apache/kyuubi/jdbc/hive/KyuubiConnection.java
@@ -104,6 +104,7 @@ public class KyuubiConnection implements SQLConnection, KyuubiLoggable {
   private Thread engineLogThread;
   private boolean engineLogInflight = true;
   private volatile boolean launchEngineOpCompleted = false;
+  private boolean launchEngineOpSupportResult = false;
   private String engineId = "";
   private String engineName = "";
   private String engineUrl = "";
@@ -770,6 +771,10 @@ private void openSession() throws SQLException {
       String launchEngineOpHandleSecret =
           openRespConf.get("kyuubi.session.engine.launch.handle.secret");
 
+      launchEngineOpSupportResult =
+          Boolean.parseBoolean(
+              openRespConf.getOrDefault("kyuubi.session.engine.launch.support.result", "false"));
+
       if (launchEngineOpHandleGuid != null && launchEngineOpHandleSecret != null) {
         try {
           byte[] guidBytes = Base64.getMimeDecoder().decode(launchEngineOpHandleGuid);
@@ -1353,7 +1358,7 @@ public void waitLaunchEngineToComplete() throws SQLException {
   }
 
   private void fetchLaunchEngineResult() {
-    if (launchEngineOpHandle == null) return;
+    if (launchEngineOpHandle == null || !launchEngineOpSupportResult) return;
 
     TFetchResultsReq tFetchResultsReq =
         new TFetchResultsReq(

diff --git a/kyuubi-server/src/main/scala/org/apache/kyuubi/server/KyuubiTBinaryFrontendService.scala b/kyuubi-server/src/main/scala/org/apache/kyuubi/server/KyuubiTBinaryFrontendService.scala
@@ -92,6 +92,8 @@ final class KyuubiTBinaryFrontendService(
         KYUUBI_SESSION_ENGINE_LAUNCH_HANDLE_SECRET,
         Base64.getMimeEncoder.encodeToString(opHandleIdentifier.getSecret))
 
+      respConfiguration.put(KYUUBI_SESSION_ENGINE_LAUNCH_SUPPORT_RESULT, true.toString)
+
       resp.setSessionHandle(sessionHandle.toTSessionHandle)
       resp.setConfiguration(respConfiguration)
       resp.setStatus(OK_STATUS)