[KYUUBI #7248] Ensure jdbc engine statements are canceled when receive cancel operation

Author: ruanwenjun

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,26 @@ class ExecuteStatement(
     super.validateFetchOrientation(order)
   }
 
+  override def cancel(): Unit = withLockRequired {
+    if (!isTerminalState(state)) {
+      setState(OperationState.CANCELED)
+      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

@@ -75,4 +75,23 @@ class OperationWithEngineSuite extends MySQLOperationSuite with HiveJDBCTestHelp
       assert(tFetchResultsResp.getStatus.getStatusCode === TStatusCode.SUCCESS_STATUS)
     }
   }
+
+  test("Mysql - JDBC ExecuteStatement operation cancel should kill SQL statement") {
+    withSessionHandle { (client, handle) =>
+      val tExecuteStatementReq = new TExecuteStatementReq()
+      tExecuteStatementReq.setSessionHandle(handle)
+      tExecuteStatementReq.setStatement("SELECT sleep(1200)")
+      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 1200s
+    }
+  }
 }

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

@@ -75,4 +75,26 @@ class StarRocksOperationWithEngineSuite extends StarRocksOperationSuite with Hiv
       assert(tFetchResultsResp.getStatus.getStatusCode === TStatusCode.SUCCESS_STATUS)
     }
   }
+
+  test("starrocks - JDBC ExecuteStatement operation cancel should kill SQL statement") {
+    withSessionHandle { (client, handle) =>
+      val tExecuteStatementReq = new TExecuteStatementReq()
+      tExecuteStatementReq.setSessionHandle(handle)
+      tExecuteStatementReq.setStatement("SELECT sleep(1200)")
+      tExecuteStatementReq.setRunAsync(true)
+      val tExecuteStatementResp = client.ExecuteStatement(tExecuteStatementReq)
+
+      // todo:we need to improvement the cancel mechanism in starrocks server side,
+      //  once a query submitted, we should be able to cancel it immediately.
+      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 1200s
+    }
+
+  }
 }