[KYUUBI #7248] JDBC engine should cancel the statement when receive cancel operation

externals/kyuubi-jdbc-engine/src/main/scala/org/apache/kyuubi/engine/jdbc/dialect/JdbcDialect.scala

@@ -81,6 +81,18 @@ abstract class JdbcDialect extends SupportServiceLoader with Logging {
   def getTRowSetGenerator(): JdbcTRowSetGenerator
 
   def getSchemaHelper(): SchemaHelper
+
+  def cancelStatement(jdbcStatement: Statement): Unit = {
+    if (jdbcStatement != null) {
+      jdbcStatement.cancel()
+    }
+  }
+
+  def closeStatement(jdbcStatement: Statement): Unit = {
+    if (jdbcStatement != null && !jdbcStatement.isClosed) {
+      jdbcStatement.close()
+    }
+  }
 }
 
 object JdbcDialects extends Logging {

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

@@ -120,12 +120,31 @@ class ExecuteStatement(
     super.validateFetchOrientation(order)
   }
 
+  override def cancel(): Unit = withLockRequired {
+    if (!isTerminalState(state)) {
+      setState(OperationState.CANCELED)
+      // TODO: If `shouldRunAsync` is true, the statement is initialized lazily.
+      // When a SQL is submitted and immediately canceled, `jdbcStatement` may still be null,
+      // which can lead to the cancellation not taking effect.
+      if (jdbcStatement != null) {
+        dialect.cancelStatement(jdbcStatement)
+        jdbcStatement = null
+      } else {
+        warn(s"The cancel operation $statementId might be ignore due to jdbcStatement is null.")
+      }
+    }
+  }
+
   override def cleanup(targetState: OperationState): Unit = withLockRequired {
     try {
       super.cleanup(targetState)
     } finally {
-      if (jdbcStatement != null && !jdbcStatement.isClosed) {
-        jdbcStatement.close()
+      if (jdbcStatement != null) {
+        if (targetState == OperationState.CANCELED) {
+          dialect.cancelStatement(jdbcStatement)
+        } else {
+          dialect.closeStatement(jdbcStatement)
+        }
         jdbcStatement = null
       }
     }

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

@@ -62,10 +62,6 @@ abstract class JdbcOperation(session: Session) extends AbstractOperation(session
     resp
   }
 
-  override def cancel(): Unit = {
-    cleanup(OperationState.CANCELED)
-  }
-
   override def close(): Unit = {
     cleanup(OperationState.CLOSED)
   }

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

@@ -16,10 +16,14 @@
  */
 package org.apache.kyuubi.engine.jdbc.mysql
 
+import org.scalatest.concurrent.TimeLimits.failAfter
+import org.scalatest.time.SpanSugar.convertIntToGrainOfTime
+
 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 {
 
@@ -75,4 +79,34 @@ class OperationWithEngineSuite extends MySQLOperationSuite with HiveJDBCTestHelp
       assert(tFetchResultsResp.getStatus.getStatusCode === TStatusCode.SUCCESS_STATUS)
     }
   }
+
+  test("MySQL - JDBC ExecuteStatement cancel operation should kill SQL statement") {
+    failAfter(20.seconds) {
+      withSessionHandle { (client, handle) =>
+        val tExecuteStatementReq = new TExecuteStatementReq()
+        tExecuteStatementReq.setSessionHandle(handle)
+        // The SQL will sleep 120s
+        tExecuteStatementReq.setStatement("SELECT sleep(120)")
+        tExecuteStatementReq.setRunAsync(true)
+        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

@@ -16,10 +16,15 @@
  */
 package org.apache.kyuubi.engine.jdbc.starrocks
 
+import scala.concurrent.duration.DurationInt
+
+import org.scalatest.concurrent.TimeLimits.failAfter
+
 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 {
 
@@ -75,4 +80,33 @@ class StarRocksOperationWithEngineSuite extends StarRocksOperationSuite with Hiv
       assert(tFetchResultsResp.getStatus.getStatusCode === TStatusCode.SUCCESS_STATUS)
     }
   }
+
+  test("starrocks - JDBC ExecuteStatement cancel operation should kill SQL statement") {
+    failAfter(20.seconds) {
+      withSessionHandle { (client, handle) =>
+        val executeReq = new TExecuteStatementReq()
+        executeReq.setSessionHandle(handle)
+        // The SQL will sleep 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))
     }
   }