diff --git a/externals/kyuubi-flink-sql-engine/src/main/scala/org/apache/kyuubi/engine/flink/operation/PlanOnlyStatement.scala b/externals/kyuubi-flink-sql-engine/src/main/scala/org/apache/kyuubi/engine/flink/operation/PlanOnlyStatement.scala index 1284bfd73e6..b3b076f1288 100644 --- a/externals/kyuubi-flink-sql-engine/src/main/scala/org/apache/kyuubi/engine/flink/operation/PlanOnlyStatement.scala +++ b/externals/kyuubi-flink-sql-engine/src/main/scala/org/apache/kyuubi/engine/flink/operation/PlanOnlyStatement.scala @@ -67,7 +67,10 @@ class PlanOnlyStatement( statement) resultSet = ResultSetUtil.fromResultFetcher(resultFetcher, resultMaxRows, resultFetchTimeout); - case _ => explainOperation(statement) + shutdownTimeoutMonitor + case _ => + explainOperation(statement) + shutdownTimeoutMonitor } } catch { onError() diff --git a/externals/kyuubi-spark-sql-engine/src/main/scala/org/apache/kyuubi/engine/spark/operation/ExecutePython.scala b/externals/kyuubi-spark-sql-engine/src/main/scala/org/apache/kyuubi/engine/spark/operation/ExecutePython.scala index a3e090d2326..593a78e42b4 100644 --- a/externals/kyuubi-spark-sql-engine/src/main/scala/org/apache/kyuubi/engine/spark/operation/ExecutePython.scala +++ b/externals/kyuubi-spark-sql-engine/src/main/scala/org/apache/kyuubi/engine/spark/operation/ExecutePython.scala @@ -127,6 +127,7 @@ class ExecutePython( val ke = KyuubiSQLException("Error submitting python in background", rejected) setOperationException(ke) + shutdownTimeoutMonitor throw ke } } else { diff --git a/externals/kyuubi-spark-sql-engine/src/main/scala/org/apache/kyuubi/engine/spark/operation/ExecuteScala.scala b/externals/kyuubi-spark-sql-engine/src/main/scala/org/apache/kyuubi/engine/spark/operation/ExecuteScala.scala index e8335e549ea..a5fd68804f0 100644 --- a/externals/kyuubi-spark-sql-engine/src/main/scala/org/apache/kyuubi/engine/spark/operation/ExecuteScala.scala +++ b/externals/kyuubi-spark-sql-engine/src/main/scala/org/apache/kyuubi/engine/spark/operation/ExecuteScala.scala @@ -154,6 +154,7 @@ class ExecuteScala( val ke = KyuubiSQLException("Error submitting scala in background", rejected) setOperationException(ke) + shutdownTimeoutMonitor throw ke } } else { diff --git a/externals/kyuubi-spark-sql-engine/src/main/scala/org/apache/kyuubi/engine/spark/operation/ExecuteStatement.scala b/externals/kyuubi-spark-sql-engine/src/main/scala/org/apache/kyuubi/engine/spark/operation/ExecuteStatement.scala index 53f02b2fbf6..40657828c98 100644 --- a/externals/kyuubi-spark-sql-engine/src/main/scala/org/apache/kyuubi/engine/spark/operation/ExecuteStatement.scala +++ b/externals/kyuubi-spark-sql-engine/src/main/scala/org/apache/kyuubi/engine/spark/operation/ExecuteStatement.scala @@ -124,6 +124,7 @@ class ExecuteStatement( val ke = KyuubiSQLException("Error submitting query in background, query rejected", rejected) setOperationException(ke) + shutdownTimeoutMonitor throw ke } } else { diff --git a/externals/kyuubi-trino-engine/src/main/scala/org/apache/kyuubi/engine/trino/operation/ExecuteStatement.scala b/externals/kyuubi-trino-engine/src/main/scala/org/apache/kyuubi/engine/trino/operation/ExecuteStatement.scala index 4f1b42e1d1b..3ecc72b772a 100644 --- a/externals/kyuubi-trino-engine/src/main/scala/org/apache/kyuubi/engine/trino/operation/ExecuteStatement.scala +++ b/externals/kyuubi-trino-engine/src/main/scala/org/apache/kyuubi/engine/trino/operation/ExecuteStatement.scala @@ -76,6 +76,7 @@ class ExecuteStatement( val ke = KyuubiSQLException("Error submitting query in background, query rejected", rejected) setOperationException(ke) + shutdownTimeoutMonitor throw ke } } else { diff --git a/kyuubi-server/src/main/scala/org/apache/kyuubi/operation/ExecuteStatement.scala b/kyuubi-server/src/main/scala/org/apache/kyuubi/operation/ExecuteStatement.scala index 02c22e4f6a7..c9f4a74cff4 100644 --- a/kyuubi-server/src/main/scala/org/apache/kyuubi/operation/ExecuteStatement.scala +++ b/kyuubi-server/src/main/scala/org/apache/kyuubi/operation/ExecuteStatement.scala @@ -172,7 +172,10 @@ class ExecuteStatement( try { val opHandle = sessionManager.submitBackgroundOperation(asyncOperation) setBackgroundHandle(opHandle) - } catch onError("submitting query in background, query rejected") + } catch { + shutdownTimeoutMonitor + onError("submitting query in background, query rejected") + } if (!shouldRunAsync) getBackgroundHandle.get() }