[KYUUBI #7323] Support timeout at PlanOnlyMode

Author: ruanwenjun

externals/kyuubi-spark-sql-engine/src/main/scala/org/apache/kyuubi/engine/spark/operation/PlanOnlyStatement.scala

@@ -18,6 +18,7 @@
 package org.apache.kyuubi.engine.spark.operation
 
 import java.lang.{Boolean => JBoolean}
+import java.util.concurrent.{TimeoutException, TimeUnit}
 
 import com.fasterxml.jackson.databind.ObjectMapper
 import com.fasterxml.jackson.module.scala.DefaultScalaModule
@@ -33,11 +34,12 @@ import org.apache.kyuubi.KyuubiSQLException
 import org.apache.kyuubi.config.KyuubiConf.{LINEAGE_PARSER_PLUGIN_PROVIDER, OPERATION_PLAN_ONLY_EXCLUDES, OPERATION_PLAN_ONLY_OUT_STYLE}
 import org.apache.kyuubi.engine.spark.KyuubiSparkUtil.getSessionConf
 import org.apache.kyuubi.engine.spark.operation.PlanOnlyStatement._
-import org.apache.kyuubi.operation.{AnalyzeMode, ArrayFetchIterator, ExecutionMode, IterableFetchIterator, JsonStyle, LineageMode, OperationHandle, OptimizeMode, OptimizeWithStatsMode, ParseMode, PhysicalMode, PlainStyle, PlanOnlyMode, PlanOnlyStyle, UnknownMode, UnknownStyle}
+import org.apache.kyuubi.operation.{AnalyzeMode, ArrayFetchIterator, ExecutionMode, IterableFetchIterator, JsonStyle, LineageMode, OperationHandle, OperationState, OptimizeMode, OptimizeWithStatsMode, ParseMode, PhysicalMode, PlainStyle, PlanOnlyMode, PlanOnlyStyle, UnknownMode, UnknownStyle}
 import org.apache.kyuubi.operation.PlanOnlyMode.{notSupportedModeError, unknownModeError}
 import org.apache.kyuubi.operation.PlanOnlyStyle.{notSupportedStyleError, unknownStyleError}
 import org.apache.kyuubi.operation.log.OperationLog
 import org.apache.kyuubi.session.Session
+import org.apache.kyuubi.util.ThreadUtils
 import org.apache.kyuubi.util.reflect.DynMethods
 
 /**
@@ -47,6 +49,7 @@ class PlanOnlyStatement(
     session: Session,
     override val statement: String,
     mode: PlanOnlyMode,
+    queryTimeout: Long,
     override protected val handle: OperationHandle)
   extends SparkOperation(session) {
 
@@ -76,27 +79,48 @@ class PlanOnlyStatement(
 
   override protected def runInternal(): Unit =
     try {
-      withLocalProperties {
-        SQLConf.withExistingConf(spark.sessionState.conf) {
-          val parsed = spark.sessionState.sqlParser.parsePlan(statement)
-          parsed match {
-            case cmd if planExcludes.contains(cmd.getClass.getSimpleName) =>
-              result = spark.sql(statement)
-              iter = new ArrayFetchIterator(result.collect())
-
-            case plan => style match {
-                case PlainStyle => explainWithPlainStyle(plan)
-                case JsonStyle => explainWithJsonStyle(plan)
-                case UnknownStyle => unknownStyleError(style)
-                case other => throw notSupportedStyleError(other, "Spark SQL")
-              }
-          }
+      if (queryTimeout > 0) {
+        val timeoutExecutor =
+          ThreadUtils.newDaemonSingleThreadScheduledExecutor("query-timeout-thread", false)
+        val future = timeoutExecutor.submit(new Runnable {
+          override def run(): Unit = doParsePlan()
+        })
+        try {
+          future.get(queryTimeout, TimeUnit.SECONDS)
+        } catch {
+          case _: TimeoutException =>
+            future.cancel(true)
+            cleanup(OperationState.TIMEOUT)
+        } finally {
+          timeoutExecutor.shutdownNow()
         }
+      } else {
+        doParsePlan()
       }
     } catch {
       onError()
     }
 
+  private def doParsePlan(): Unit = {
+    withLocalProperties {
+      SQLConf.withExistingConf(spark.sessionState.conf) {
+        val parsed = spark.sessionState.sqlParser.parsePlan(statement)
+        parsed match {
+          case cmd if planExcludes.contains(cmd.getClass.getSimpleName) =>
+            result = spark.sql(statement)
+            iter = new ArrayFetchIterator(result.collect())
+
+          case plan => style match {
+              case PlainStyle => explainWithPlainStyle(plan)
+              case JsonStyle => explainWithJsonStyle(plan)
+              case UnknownStyle => unknownStyleError(style)
+              case other => throw notSupportedStyleError(other, "Spark SQL")
+            }
+        }
+      }
+    }
+  }
+
   private def explainWithPlainStyle(plan: LogicalPlan): Unit = {
     mode match {
       case ParseMode =>

externals/kyuubi-spark-sql-engine/src/main/scala/org/apache/kyuubi/engine/spark/operation/SparkSQLOperationManager.scala

@@ -109,7 +109,7 @@ class SparkSQLOperationManager private (name: String) extends OperationManager(n
                     opHandle)
               }
             case mode =>
-              new PlanOnlyStatement(session, statement, mode, opHandle)
+              new PlanOnlyStatement(session, statement, mode, queryTimeout, opHandle)
           }
         case OperationLanguages.SCALA =>
           val repl = sessionToRepl.getOrElseUpdate(session.handle, KyuubiSparkILoop(spark))

kyuubi-server/src/test/scala/org/apache/kyuubi/engine/spark/SparkSqlEngineSuite.scala

@@ -17,6 +17,8 @@
 
 package org.apache.kyuubi.engine.spark
 
+import java.sql.SQLTimeoutException
+
 import org.apache.kyuubi.WithKyuubiServer
 import org.apache.kyuubi.config.KyuubiConf
 import org.apache.kyuubi.config.KyuubiConf._
@@ -149,5 +151,38 @@ class SparkSqlEngineSuite extends WithKyuubiServer with HiveJDBCTestHelper {
     }
   }
 
+  test("KYUUBI-7323: Support timeout at PlanOnlyMode") {
+    val tableName = "t_timeout"
+    val sessionConfMap = Map("kyuubi.operation.plan.only.mode=analyze" -> "analyze")
+    withSessionConf(sessionConfMap)(Map.empty)(Map.empty) {
+
+      withJdbcStatement(tableName) { statement =>
+        statement.setQueryTimeout(2)
+        val layers = 30
+
+        val cteBuilder = new StringBuilder()
+        cteBuilder.append("WITH p AS (SELECT dt, tid, count(1) as pv, DENSE_RANK() " +
+          "over(order by dt) as dt_rank FROM t GROUP BY 1, 2), ")
+        cteBuilder.append("test_d1 AS (SELECT dt, tid FROM (SELECT dt, tid, row_number() " +
+          "over(order by pv desc) as r FROM p WHERE dt_rank=1) WHERE r<=100)")
+
+        for (i <- 2 to layers) {
+          cteBuilder.append(s", test_d$i AS ( " +
+            s"SELECT dt, tid FROM (SELECT p.dt, p.tid, row_number() over(order by pv desc) as r " +
+            s"FROM p LEFT JOIN test_d${i - 1} prev ON p.tid = prev.tid " +
+            s"WHERE p.dt_rank=$i AND prev.tid IS NULL) WHERE r<=100 " +
+            s"UNION ALL SELECT dt, tid FROM test_d${i - 1})")
+        }
+
+        val complexSql = cteBuilder.toString() + s" SELECT * FROM test_d$layers"
+
+        val e = intercept[SQLTimeoutException] {
+          statement.executeQuery(complexSql)
+        }
+        assert(e.getMessage.equals("Query timed out after 2 seconds"))
+      }
+    }
+  }
+
   override protected def jdbcUrl: String = getJdbcUrl
 }