From f2ac40e390d7992ae2736f13f76e695f0006e9d1 Mon Sep 17 00:00:00 2001
From: zhang-arvin <arvin.zhang@htx-inc.com>
Date: Sun, 30 Aug 2026 23:37:53 +0800
Subject: [PATCH] [BugFix] Fix JDBC kill task not cleaning up DDL state (#5320)

---
 .../executor/JDBCEngineConnExecutor.scala     | 78 +++++++++++++++++++
 1 file changed, 78 insertions(+)

diff --git a/linkis-engineconn-plugins/jdbc/src/main/scala/org/apache/linkis/manager/engineplugin/jdbc/executor/JDBCEngineConnExecutor.scala b/linkis-engineconn-plugins/jdbc/src/main/scala/org/apache/linkis/manager/engineplugin/jdbc/executor/JDBCEngineConnExecutor.scala
index 9d05100f3dd..565341c07a4 100644
--- a/linkis-engineconn-plugins/jdbc/src/main/scala/org/apache/linkis/manager/engineplugin/jdbc/executor/JDBCEngineConnExecutor.scala
+++ b/linkis-engineconn-plugins/jdbc/src/main/scala/org/apache/linkis/manager/engineplugin/jdbc/executor/JDBCEngineConnExecutor.scala
@@ -76,6 +76,7 @@ import org.springframework.util.CollectionUtils
 import java.sql.{Connection, ResultSet, Statement}
 import java.util
 import java.util.concurrent.ConcurrentHashMap
+import java.util.regex.{Matcher, Pattern}
 
 import scala.collection.mutable.ArrayBuffer
 
@@ -90,6 +91,8 @@ class JDBCEngineConnExecutor(override val outputPrintLimit: Int, val id: Int)
 
   private val connectionCache: util.Map[String, Connection] = new util.HashMap[String, Connection]()
 
+  private val taskCodeMap: util.Map[String, String] = new ConcurrentHashMap[String, String]()
+
   override def init(): Unit = {
     logger.info("jdbc executor start init.")
     setCodeParser(new SQLCodeParser)
@@ -176,6 +179,8 @@ class JDBCEngineConnExecutor(override val outputPrintLimit: Int, val id: Int)
   ): ExecuteResponse = {
 
     val taskId = engineExecutorContext.getJobId.get
+    // Store the code for potential DDL cleanup in killTask
+    taskCodeMap.put(taskId, code)
     val connection: Connection = getConnection(engineExecutorContext)
     var statement: Statement = null
     var resultSet: ResultSet = null
@@ -486,9 +491,82 @@ class JDBCEngineConnExecutor(override val outputPrintLimit: Int, val id: Int)
     logger.info("All query task has killed successfully.")
   }
 
+  /**
+   * Pattern to extract CREATE TABLE/VIEW statements and their object names.
+   * Captures: CREATE [TEMPORARY] TABLE [IF NOT EXISTS] <objectName>
+   *           CREATE [OR REPLACE] VIEW [IF NOT EXISTS] <objectName>
+   */
+  private val DDL_PATTERN: Pattern = Pattern.compile(
+    "(?i)CREATE\\s+(TEMPORARY\\s+)?(TABLE|VIEW)\\s+(IF\\s+NOT\\s+EXISTS\\s+|OR\\s+REPLACE\\s+)?([`\"\\[]?[\\w.-]+[`\"\\]]?)",
+    Pattern.CASE_INSENSITIVE
+  )
+
+  /**
+   * Extract cleanup DDL statements from the original SQL code.
+   * For each CREATE TABLE/VIEW found, generates a corresponding DROP IF EXISTS statement.
+   */
+  private def extractDDLStatements(code: String): List[String] = {
+    val matcher: Matcher = DDL_PATTERN.matcher(code)
+    val statements = scala.collection.mutable.ListBuffer[String]()
+    while (matcher.find()) {
+      val objectType = matcher.group(2).toUpperCase
+      val objectName = matcher.group(4)
+      if (StringUtils.isNotBlank(objectName)) {
+        objectType match {
+          case "TABLE" =>
+            statements += s"DROP TABLE IF EXISTS $objectName"
+          case "VIEW" =>
+            statements += s"DROP VIEW IF EXISTS $objectName"
+          case _ =>
+        }
+      }
+    }
+    statements.toList
+  }
+
+  /**
+   * Execute cleanup DDL statements to roll back partially created objects
+   * when a task is killed mid-execution.
+   */
+  private def cleanUpDDLState(taskId: String, code: String): Unit = {
+    val ddls = extractDDLStatements(code)
+    if (ddls.nonEmpty && connectionCache.containsKey(taskId)) {
+      val connection = connectionCache.get(taskId)
+      var statement: Statement = null
+      Utils.tryCatch {
+        statement = connection.createStatement()
+        ddls.foreach { ddl =>
+          logger.info(s"Executing cleanup DDL for task $taskId: $ddl")
+          Utils.tryCatch {
+            statement.execute(ddl)
+          } { case e: Throwable =>
+            logger.warn(s"Cleanup DDL failed for task $taskId: $ddl", e)
+          }
+        }
+      } { case e: Throwable =>
+        logger.warn(s"Failed to execute cleanup DDL for task $taskId", e)
+      }
+      Utils.tryFinally(null)(_ => {
+        if (statement != null) {
+          Utils.tryAndWarn(statement.close())
+        }
+      })
+    }
+  }
+
   override def killTask(taskId: String): Unit = {
     logger.info(s"Killing jdbc query task $taskId")
     connectionManager.cancelStatement(taskId)
+    // Clean up DDL state: if the task was executing DDL (CREATE TABLE/VIEW),
+    // attempt to drop the partially created objects
+    val code = taskCodeMap.remove(taskId)
+    if (StringUtils.isNotBlank(code)) {
+      Utils.tryCatch {
+        cleanUpDDLState(taskId, code)
+      } { case e: Throwable =>
+        logger.warn(s"Failed to clean up DDL state for task $taskId", e)
+      }
+    }
     super.killTask(taskId)
     logger.info(s"The query task $taskId has killed successfully.")
   }