[incubator-kyuubi] branch master updated: [KYUUBI #1920] Skip the Plan parsing of UseStatement

[ya...@apache.org]

This is an automated email from the ASF dual-hosted git repository.

yao pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/incubator-kyuubi.git


The following commit(s) were added to refs/heads/master by this push:
     new cec25d9  [KYUUBI #1920] Skip the Plan parsing of UseStatement
cec25d9 is described below

commit cec25d985f8ac55254475c0eb97773deebd7f095
Author: odone <od...@gmail.com>
AuthorDate: Tue Feb 22 16:45:38 2022 +0800

    [KYUUBI #1920] Skip the Plan parsing of UseStatement
    
    Close #1920
    
    ### _Why are the changes needed?_
    
    ### _How was this patch tested?_
    - [x] Add some test cases that check the changes thoroughly including negative and positive cases if possible
    
    - [ ] Add screenshots for manual tests if appropriate
    
    - [ ] [Run test](https://kyuubi.apache.org/docs/latest/develop_tools/testing.html#running-tests) locally before make a pull request
    
    Closes #1921 from iodone/kyuubi-1902.
    
    Closes #1920
    
    79d7ba68 [odone] [KYUUBI #1920] Added CacheTableAsSelect to be skiped
    571003de [odone] [KYUUBI #1920] Added CacheTableStatement to be skiped
    889cb2b9 [odone] [KYUUBI #1920] Added CacheTableCommand and CreateView to be skiped
    456dbf5f [odone] [KYUUBI #1920] Update test
    90001690 [odone] [KYUUBI #1920] Update
    80e0ceff [odone] [KYUUBI #1920] Update
    
    Authored-by: odone <od...@gmail.com>
    Signed-off-by: Kent Yao <ya...@apache.org>
---
 .../engine/spark/operation/PlanOnlyStatement.scala | 23 +++++++++++------
 .../kyuubi/operation/PlanOnlyOperationSuite.scala  | 29 ++++++++++++++++++++++
 2 files changed, 45 insertions(+), 7 deletions(-)

diff --git a/externals/kyuubi-spark-sql-engine/src/main/scala/org/apache/kyuubi/engine/spark/operation/PlanOnlyStatement.scala b/externals/kyuubi-spark-sql-engine/src/main/scala/org/apache/kyuubi/engine/spark/operation/PlanOnlyStatement.scala
index 304b031..a4abf88 100644
--- a/externals/kyuubi-spark-sql-engine/src/main/scala/org/apache/kyuubi/engine/spark/operation/PlanOnlyStatement.scala
+++ b/externals/kyuubi-spark-sql-engine/src/main/scala/org/apache/kyuubi/engine/spark/operation/PlanOnlyStatement.scala
@@ -38,23 +38,32 @@ class PlanOnlyStatement(
   private val operationLog: OperationLog = OperationLog.createOperationLog(session, getHandle)
   override def getOperationLog: Option[OperationLog] = Option(operationLog)
 
-  override protected def resultSchema: StructType =
+  override protected def resultSchema: StructType = {
     if (result == null) {
       new StructType().add("plan", "string")
-    } else {
-      result.schema
-    }
+    } else if (result.isEmpty) {
+      new StructType().add("result", "string")
+    } else result.schema
+  }
 
-  private def isSetOrReset(plan: LogicalPlan): Boolean = {
+  private def shouldDirectRun(plan: LogicalPlan): Boolean = {
     val className = plan.getClass.getSimpleName
-    className == "SetCommand" || className == "ResetCommand"
+    className == "SetCommand" ||
+    className == "ResetCommand" ||
+    className == "UseStatement" ||
+    className == "SetNamespaceCommand" ||
+    className == "CacheTableStatement" ||
+    className == "CacheTableCommand" ||
+    className == "CacheTableAsSelect" ||
+    className == "CreateViewStatement" ||
+    className == "CreateViewCommand"
   }
 
   override protected def runInternal(): Unit = {
     try {
       val parsed = spark.sessionState.sqlParser.parsePlan(statement)
       parsed match {
-        case cmd if isSetOrReset(cmd) =>
+        case cmd if shouldDirectRun(cmd) =>
           result = spark.sql(statement)
           iter = new ArrayFetchIterator(result.collect())
         case plan => mode match {
diff --git a/kyuubi-server/src/test/scala/org/apache/kyuubi/operation/PlanOnlyOperationSuite.scala b/kyuubi-server/src/test/scala/org/apache/kyuubi/operation/PlanOnlyOperationSuite.scala
index 22d1f4e..2c526d1 100644
--- a/kyuubi-server/src/test/scala/org/apache/kyuubi/operation/PlanOnlyOperationSuite.scala
+++ b/kyuubi-server/src/test/scala/org/apache/kyuubi/operation/PlanOnlyOperationSuite.scala
@@ -79,6 +79,35 @@ class PlanOnlyOperationSuite extends WithKyuubiServer with HiveJDBCTestHelper {
     }
   }
 
+  test("KYUUBI #1920: Plan only operations with Usestatement or SetNamespaceCommand skiped") {
+    withSessionConf()(Map(KyuubiConf.OPERATION_PLAN_ONLY.key -> NONE.toString))(Map.empty) {
+      withDatabases("test_database") { statement =>
+        statement.execute("create database test_database")
+        statement.execute(s"set ${KyuubiConf.OPERATION_PLAN_ONLY.key}=optimize")
+        val result = statement.executeQuery("use test_database")
+        assert(!result.next(), "In contrast to PlanOnly mode, it will returns an empty result")
+      }
+    }
+  }
+
+  test("KYUUBI #1920: Plan only operations with CacheTable skiped") {
+    withSessionConf()(Map(KyuubiConf.OPERATION_PLAN_ONLY.key -> OPTIMIZE.toString))(Map.empty) {
+      withJdbcStatement() { statement =>
+        val result = statement.executeQuery("cache table cached_table as select 1")
+        assert(!result.next(), "In contrast to PlanOnly mode, it will returns an empty result")
+      }
+    }
+  }
+
+  test("KYUUBI #1920: Plan only operations with CreateViewStatement or CreateViewCommand skiped") {
+    withSessionConf()(Map(KyuubiConf.OPERATION_PLAN_ONLY.key -> OPTIMIZE.toString))(Map.empty) {
+      withJdbcStatement() { statement =>
+        val result = statement.executeQuery("create temp view temp_view as select 1")
+        assert(!result.next(), "In contrast to PlanOnly mode, it will returns an empty result")
+      }
+    }
+  }
+
   private def getOperationPlanWithStatement(statement: Statement): String = {
     val resultSet = statement.executeQuery("select 1 where true")
     assert(resultSet.next())