jdbc ExecuteStatement should set setCurrentOperationLog

lsm1 · Feb 20, 2025 · 4924266 · 4924266
1 parent 0f4906d
commit 4924266
Showing 1 changed file with 10 additions and 0 deletions.
diff --git a/...jdbc-engine/src/main/scala/org/apache/kyuubi/engine/jdbc/operation/ExecuteStatement.scala b/...jdbc-engine/src/main/scala/org/apache/kyuubi/engine/jdbc/operation/ExecuteStatement.scala
@@ -43,6 +43,16 @@ class ExecuteStatement(
 
   @volatile private var jdbcStatement: Statement = _
 
+  override protected def beforeRun(): Unit = {
+    OperationLog.setCurrentOperationLog(operationLog)
+    setState(OperationState.PENDING)
+    setHasResultSet(true)
+  }
+
+  override protected def afterRun(): Unit = {
+    OperationLog.removeCurrentOperationLog()
+  }
+
   override protected def runInternal(): Unit = {
     addTimeoutMonitor(queryTimeout)
     if (shouldRunAsync) {