add status assertion in new ut case

Author: ruanwenjun

externals/kyuubi-jdbc-engine/src/main/scala/org/apache/kyuubi/engine/jdbc/operation/ExecuteStatement.scala

@@ -129,6 +129,8 @@ class ExecuteStatement(
       if (jdbcStatement != null) {
         dialect.cancelStatement(jdbcStatement)
         jdbcStatement = null
+      } else {
+        warn(s"The cancel operation $statementId might be ignore due to jdbcStatement is null.")
       }
     }
   }

externals/kyuubi-jdbc-engine/src/test/scala/org/apache/kyuubi/engine/jdbc/mysql/OperationWithEngineSuite.scala

@@ -23,6 +23,7 @@ import org.apache.kyuubi.config.KyuubiConf
 import org.apache.kyuubi.engine.jdbc.connection.ConnectionProvider
 import org.apache.kyuubi.operation.HiveJDBCTestHelper
 import org.apache.kyuubi.shaded.hive.service.rpc.thrift._
+import org.apache.kyuubi.shaded.hive.service.rpc.thrift.TOperationState.{CANCELED_STATE, RUNNING_STATE}
 
 class OperationWithEngineSuite extends MySQLOperationSuite with HiveJDBCTestHelper {
 
@@ -87,16 +88,24 @@ class OperationWithEngineSuite extends MySQLOperationSuite with HiveJDBCTestHelp
         // The SQL will sleep 120s
         tExecuteStatementReq.setStatement("SELECT sleep(120)")
         tExecuteStatementReq.setRunAsync(true)
-        val tExecuteStatementResp = client.ExecuteStatement(tExecuteStatementReq)
-        assert(tExecuteStatementResp.getStatus.getStatusCode === TStatusCode.SUCCESS_STATUS)
-
-        Thread.sleep(1000) // wait for statement to start executing
-
-        val tCancelOperationReq = new TCancelOperationReq()
-        tCancelOperationReq.setOperationHandle(tExecuteStatementResp.getOperationHandle)
-        val TCancelOperationReq = client.CancelOperation(tCancelOperationReq)
-        assert(TCancelOperationReq.getStatus.getStatusCode === TStatusCode.SUCCESS_STATUS)
-        // If the statement is not cancelled successfully, will block here until 120s
+        val executeResp = client.ExecuteStatement(tExecuteStatementReq)
+        assert(executeResp.getStatus.getStatusCode === TStatusCode.SUCCESS_STATUS)
+
+        assertOperationStatusIn(
+          client,
+          executeResp.getOperationHandle,
+          Set(RUNNING_STATE),
+          5)
+
+        val cancelResp =
+          client.CancelOperation(new TCancelOperationReq(executeResp.getOperationHandle))
+        assert(cancelResp.getStatus.getStatusCode === TStatusCode.SUCCESS_STATUS)
+
+        assertOperationStatusIn(
+          client,
+          executeResp.getOperationHandle,
+          Set(CANCELED_STATE),
+          5)
       }
     }
   }

externals/kyuubi-jdbc-engine/src/test/scala/org/apache/kyuubi/engine/jdbc/starrocks/StarRocksOperationWithEngineSuite.scala

@@ -24,6 +24,7 @@ import org.apache.kyuubi.config.KyuubiConf
 import org.apache.kyuubi.engine.jdbc.connection.ConnectionProvider
 import org.apache.kyuubi.operation.HiveJDBCTestHelper
 import org.apache.kyuubi.shaded.hive.service.rpc.thrift._
+import org.apache.kyuubi.shaded.hive.service.rpc.thrift.TOperationState.{CANCELED_STATE, RUNNING_STATE}
 
 class StarRocksOperationWithEngineSuite extends StarRocksOperationSuite with HiveJDBCTestHelper {
 
@@ -83,21 +84,28 @@ class StarRocksOperationWithEngineSuite extends StarRocksOperationSuite with Hiv
   test("starrocks - JDBC ExecuteStatement cancel operation should kill SQL statement") {
     failAfter(20.seconds) {
       withSessionHandle { (client, handle) =>
-        val tExecuteStatementReq = new TExecuteStatementReq()
-        tExecuteStatementReq.setSessionHandle(handle)
+        val executeReq = new TExecuteStatementReq()
+        executeReq.setSessionHandle(handle)
         // The SQL will sleep 120s
-        tExecuteStatementReq.setStatement("SELECT sleep(120)")
-        tExecuteStatementReq.setRunAsync(true)
-        val tExecuteStatementResp = client.ExecuteStatement(tExecuteStatementReq)
-
-        Thread.sleep(1000) // wait for statement to start executing
-
-        val tCancelOperationReq = new TCancelOperationReq()
-        tCancelOperationReq.setOperationHandle(tExecuteStatementResp.getOperationHandle)
-
-        val tFetchResultsResp = client.CancelOperation(tCancelOperationReq)
-        assert(tFetchResultsResp.getStatus.getStatusCode === TStatusCode.SUCCESS_STATUS)
-        // If the statement is not cancelled successfully, will block here until 120s
+        executeReq.setStatement("SELECT sleep(120)")
+        executeReq.setRunAsync(true)
+        val executeResp = client.ExecuteStatement(executeReq)
+
+        assertOperationStatusIn(
+          client,
+          executeResp.getOperationHandle,
+          Set(RUNNING_STATE),
+          5)
+
+        val cancelResp =
+          client.CancelOperation(new TCancelOperationReq(executeResp.getOperationHandle))
+        assert(cancelResp.getStatus.getStatusCode === TStatusCode.SUCCESS_STATUS)
+
+        assertOperationStatusIn(
+          client,
+          executeResp.getOperationHandle,
+          Set(CANCELED_STATE),
+          5)
       }
     }
   }

kyuubi-common/src/test/scala/org/apache/kyuubi/operation/HiveJDBCTestHelper.scala

@@ -124,11 +124,27 @@ trait HiveJDBCTestHelper extends JDBCTestHelper {
   }
 
   def waitForOperationToComplete(client: Iface, op: TOperationHandle): Unit = {
-    val req = new TGetOperationStatusReq(op)
-    var state = client.GetOperationStatus(req).getOperationState
-    eventually(timeout(90.seconds), interval(100.milliseconds)) {
-      state = client.GetOperationStatus(req).getOperationState
-      assert(!Set(INITIALIZED_STATE, PENDING_STATE, RUNNING_STATE).contains(state))
+    assertOperationStatusIn(
+      client,
+      op,
+      Set(
+        FINISHED_STATE,
+        CANCELED_STATE,
+        CLOSED_STATE,
+        ERROR_STATE,
+        UKNOWN_STATE,
+        TIMEDOUT_STATE),
+      90)
+  }
+
+  def assertOperationStatusIn(
+      client: Iface,
+      op: TOperationHandle,
+      status: Set[TOperationState],
+      timeoutInSeconds: Int): Unit = {
+    eventually(timeout(timeoutInSeconds.seconds), interval(100.milliseconds)) {
+      val state = client.GetOperationStatus(new TGetOperationStatusReq(op)).getOperationState
+      assert(status.contains(state))
     }
   }