[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,29 @@ 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
+      }
+    }
+  }
+
   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,6 +16,9 @@
  */
 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
@@ -75,4 +78,26 @@ 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 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
+      }
+    }
+  }
 }

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

@@ -16,6 +16,10 @@
  */
 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
@@ -75,4 +79,26 @@ 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 tExecuteStatementReq = new TExecuteStatementReq()
+        tExecuteStatementReq.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
+      }
+    }
+  }
 }