Merge pull request #221 from jumpserver/pr@new_terminal@feat_console_stability

fix(console): improve cancellation

Author: Zxc4wonder

backend/framework/src/main/java/org/jumpserver/chen/framework/console/QueryConsole.java

@@ -2,10 +2,12 @@
 
 import com.alibaba.druid.sql.parser.ParserException;
 import com.google.gson.Gson;
+import com.google.gson.JsonParseException;
 import lombok.extern.slf4j.Slf4j;
 import org.apache.commons.lang3.StringUtils;
 import org.jumpserver.chen.framework.console.action.DataViewAction;
 import org.jumpserver.chen.framework.console.action.QueryConsoleAction;
+import org.jumpserver.chen.framework.console.action.SQLChunkData;
 import org.jumpserver.chen.framework.console.dataview.DataView;
 import org.jumpserver.chen.framework.console.dataview.QueryDataViewTableEditContextFactory;
 import org.jumpserver.chen.framework.console.dataview.UpdateDataView;
@@ -359,16 +361,27 @@ private void onAction(QueryConsoleAction action) {
     private int expectedChunks = -1;
 
     private void handleSQLChunk(QueryConsoleAction action) {
-        var data = (Map<String, Object>) action.getData();
-        var chunk = (String) data.get("chunk");
-        var index = (Integer) data.get("index");
-        var total = (Integer) data.get("total");
+        SQLChunkData data;
+        try {
+            data = GSON.fromJson(GSON.toJson(action.getData()), SQLChunkData.class);
+        } catch (JsonParseException e) {
+            this.getConsoleLogger().error("invalid sql chunk");
+            this.resetSQLChunks();
+            return;
+        }
 
-        if (chunk == null || index == null || total == null || total <= 0) {
+        if (data == null
+                || data.getChunk() == null
+                || data.getIndex() == null
+                || data.getTotal() == null
+                || data.getTotal() <= 0) {
             this.getConsoleLogger().error("invalid sql chunk");
             this.resetSQLChunks();
             return;
         }
+        var chunk = data.getChunk();
+        var index = data.getIndex();
+        var total = data.getTotal();
         if (expectedChunks == -1) {
             expectedChunks = total;
         }
@@ -617,9 +630,9 @@ public void onCancel() {
         this.getState().setExecutionStatus(EXECUTION_STATUS_CANCELLED);
         try {
             var plan = this.currentPlan;
-            if (plan != null && plan.getStatement() != null) {
+            if (plan != null) {
                 plan.cancel();
-                this.getConsoleLogger().error("cancel query: %s", plan.getTargetSQL());
+                this.getConsoleLogger().warn("cancel query: %s", plan.getTargetSQL());
             }
         } catch (SQLException | RuntimeException e) {
             log.error("cancel failed ", e);
@@ -775,9 +788,9 @@ public void onSQL(String sql) {
         } catch (SQLException e) {
             if (!StringUtils.equals(this.getState().getExecutionStatus(), EXECUTION_STATUS_CANCELLED)) {
                 this.getState().setExecutionStatus(EXECUTION_STATUS_ERROR);
+                this.getConsoleLogger().error("%s: %s", MessageUtils.get("ExecuteError"), e.getMessage());
+                this.getPacketIO().sendPacket("message", Message.error(MessageUtils.get("ExecuteError"), e.getMessage()));
             }
-            this.getConsoleLogger().error("%s: %s", MessageUtils.get("ExecuteError"), e.getMessage());
-            this.getPacketIO().sendPacket("message", Message.error(MessageUtils.get("ExecuteError"), e.getMessage()));
         } finally {
             if (StringUtils.equals(this.getState().getExecutionStatus(), EXECUTION_STATUS_RUNNING)) {
                 this.getState().setExecutionStatus(EXECUTION_STATUS_SUCCESS);
@@ -893,15 +906,13 @@ private DataView runSingleSQL(String sql, ACLResult aclResult) throws SQLExcepti
                     .createPlan(SQL.of(sourceSQL));
             plan.setAclResult(aclResult);
             plan.setSqlQueryParams(sqlQueryParams);
-            plan.generateTargetSQL();
-            this.getConsoleLogger().info("execute sql: %s", plan.getTargetSQL());
-
             this.currentPlan = plan;
-
             this.getState().setCanCancel(true);
             this.stateManager.commit();
 
             try {
+                plan.generateTargetSQL();
+                this.getConsoleLogger().info("execute sql: %s", plan.getTargetSQL());
                 var result = plan.executeWithAudit();
                 this.getConsoleLogger().success(result);
                 return result;

backend/framework/src/main/java/org/jumpserver/chen/framework/console/action/SQLChunkData.java

@@ -0,0 +1,10 @@
+package org.jumpserver.chen.framework.console.action;
+
+import lombok.Data;
+
+@Data
+public class SQLChunkData {
+    private String chunk;
+    private Integer index;
+    private Integer total;
+}

backend/framework/src/main/java/org/jumpserver/chen/framework/datasource/sql/SQLExecutePlan.java

@@ -27,6 +27,7 @@ public class SQLExecutePlan {
     private String targetSQL;
     private final DbType druidDbType;
     private volatile Statement statement;
+    private volatile boolean cancelled;
     private Connection connection;
     private ACLResult aclResult;
 
@@ -85,7 +86,10 @@ public SQLQueryResult executeWithAudit() throws SQLException {
     }
 
 
-    public Statement createStatement() throws SQLException {
+    public synchronized Statement createStatement() throws SQLException {
+        if (this.cancelled) {
+            throw new SQLException(MessageUtils.get("ExecutionCanceled"));
+        }
         if (this.statement == null || this.statement.isClosed()) {
             this.statement = this.connection.createStatement();
         }
@@ -97,7 +101,14 @@ public SQLStatement getTargetSQLStatement() {
     }
 
     public void cancel() throws SQLException {
-        this.statement.cancel();
+        Statement currentStatement;
+        synchronized (this) {
+            this.cancelled = true;
+            currentStatement = this.statement;
+        }
+        if (currentStatement != null && !currentStatement.isClosed()) {
+            currentStatement.cancel();
+        }
     }
 
     public void close() {