Merge pull request #77 from jumpserver/pr@dev@feat_limit_sql_buffer

feat:  优化 sql 执行大小

Author: Aaron3S

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

@@ -34,6 +34,7 @@
 import java.util.ArrayList;
 import java.util.HashMap;
 import java.util.Map;
+import java.util.concurrent.ConcurrentHashMap;
 import java.util.concurrent.CountDownLatch;
 import java.util.concurrent.atomic.AtomicBoolean;
 
@@ -141,6 +142,7 @@ public void handle(Packet packet) {
         }
     }
 
+
     private void onAction(QueryConsoleAction action) {
         switch (action.getAction()) {
             case QueryConsoleAction.ACTION_RUN_SQL -> {
@@ -153,6 +155,13 @@ private void onAction(QueryConsoleAction action) {
                 this.getState().setInQuery(false);
                 this.stateManager.commit();
             }
+            case QueryConsoleAction.ACTION_RUN_SQL_CHUNK -> {
+                this.handleSQLChunk(action);
+            }
+            case QueryConsoleAction.ACTION_RUN_SQL_COMPLETE -> {
+                this.handleSQLComplete();
+            }
+
             case QueryConsoleAction.ACTION_RUN_SQL_FILE -> {
                 this.getState().setInQuery(true);
                 this.stateManager.commit();
@@ -177,6 +186,60 @@ private void onAction(QueryConsoleAction action) {
         }
     }
 
+    private final ConcurrentHashMap<Integer, String> sqlChunks = new ConcurrentHashMap<>();
+    private CountDownLatch latch;
+    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");
+
+        synchronized (this) {
+            if (expectedChunks == -1) {
+                expectedChunks = total;
+                latch = new CountDownLatch(total);
+            }
+        }
+
+        if (sqlChunks.putIfAbsent(index, chunk) == null) {
+            latch.countDown();
+        }
+    }
+
+    /**
+     * 处理分段 SQL 接收完成
+     */
+    private void handleSQLComplete() {
+        try {
+            // 等待所有分段接收完成
+            latch.await();
+
+            // 按照索引顺序合并所有分段
+            StringBuilder sqlBuilder = new StringBuilder();
+            for (int i = 0; i < expectedChunks; i++) {
+                sqlBuilder.append(sqlChunks.get(i));
+            }
+
+            // 合并完成后清理缓存
+            var sql = sqlBuilder.toString();
+            sqlChunks.clear();
+            expectedChunks = -1;
+
+            // 执行完整 SQL
+            this.getState().setInQuery(true);
+            this.stateManager.commit();
+
+            this.onSQL(sql);
+
+        } catch (InterruptedException e) {
+            Thread.currentThread().interrupt();
+        } finally {
+            this.getState().setInQuery(false);
+            this.stateManager.commit();
+        }
+    }
 
     private void onDataViewAction(DataViewAction action) {
         var dataView = this.dataViews.get(action.getDataView());

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

@@ -7,6 +7,8 @@
 @EqualsAndHashCode(callSuper = true)
 public class QueryConsoleAction extends Action {
     public static final String ACTION_RUN_SQL = "run_sql";
+    public static final String ACTION_RUN_SQL_CHUNK = "run_sql_chunk";
+    public static final String ACTION_RUN_SQL_COMPLETE = "run_sql_complete";
     public static final String ACTION_RUN_SQL_FILE = "run_sql_file";
     public static final String ACTION_CANCEL = "cancel";
     public static final String ACTION_CHANGE_CURRENT_CONTEXT = "change_current_context";

frontend/src/components/Main/Explore/QueryConsole/CodeEditor.vue

@@ -253,7 +253,24 @@ export default {
     },
     onRun() {
       const sql = this.selectionValue || this.statement
-      this.$emit('action', { action: 'run_sql', data: sql })
+      const CHUNK_SIZE = 4096
+
+      if (sql.length <= CHUNK_SIZE) {
+        this.$emit('action', { action: 'run_sql', data: sql })
+      } else {
+        const totalChunks = Math.ceil(sql.length / CHUNK_SIZE)
+        for (let i = 0; i < totalChunks; i++) {
+          const chunk = sql.slice(i * CHUNK_SIZE, (i + 1) * CHUNK_SIZE)
+          this.$emit('action', {
+            action: 'run_sql_chunk',
+            data: { chunk, index: i, total: totalChunks }
+          })
+        }
+        this.$emit('action', {
+          action: 'run_sql_complete',
+          data: { total: totalChunks }
+        })
+      }
     },
     onStop() {
       this.$emit('action', { action: 'cancel' })

frontend/src/components/Main/Explore/QueryConsole/index.vue

@@ -8,7 +8,6 @@
             :state="state"
             :subjects="subjects"
             @action="onEditorAction"
-            @run="onRunSql"
           />
         </template>
         <template slot="paneR">
@@ -153,11 +152,7 @@ export default {
     onDataViewAction(action) {
       this.ws.send(JSON.stringify({ type: 'data_view_action', data: action }))
     },
-    onRunSql(sql) {
-      this.ws.send(JSON.stringify({ type: 'sql', data: sql }))
-    },
     onCloseDataView(name) {
-      console.log(name)
       this.ws.send(JSON.stringify({ type: 'close_data_view', data: name }))
     },
     onLimitChange(limit) {