Capture exceptions in async ExecuteStatement

apache · May 26, 2021 · 8031528 · 8031528
1 parent c8b2ec4
commit 8031528
Show file tree

Hide file tree

Showing 2 changed files with 9 additions and 4 deletions.
diff --git a/kyuubi-main/src/main/scala/org/apache/kyuubi/operation/ExecuteStatement.scala b/kyuubi-main/src/main/scala/org/apache/kyuubi/operation/ExecuteStatement.scala
@@ -91,8 +91,10 @@ class ExecuteStatement(
       isComplete = true
       remoteState match {
         case INITIALIZED_STATE | PENDING_STATE | RUNNING_STATE =>
-          isComplete = false
-          statusResp = client.GetOperationStatus(statusReq)
+          try {
+            statusResp = client.GetOperationStatus(statusReq)
+            isComplete = false
+          } catch onError(rethrow = false)
 
         case FINISHED_STATE =>
           setState(OperationState.FINISHED)

diff --git a/kyuubi-main/src/main/scala/org/apache/kyuubi/operation/KyuubiOperation.scala b/kyuubi-main/src/main/scala/org/apache/kyuubi/operation/KyuubiOperation.scala
@@ -44,7 +44,8 @@ abstract class KyuubiOperation(
     ThriftUtils.verifyTStatus(tStatus)
   }
 
-  protected def onError(action: String = "operating"): PartialFunction[Throwable, Unit] = {
+  protected def onError(action: String = "operating",
+                        rethrow: Boolean = true): PartialFunction[Throwable, Unit] = {
     case e: Throwable =>
       state.synchronized {
         if (isTerminalState(state)) {
@@ -66,7 +67,9 @@ abstract class KyuubiOperation(
               KyuubiSQLException(s"Error $action $opType: ${e.getMessage}", e)
           }
           setOperationException(ke)
-          throw ke
+          if (rethrow) {
+            throw ke
+          }
         }
       }
   }