Merge pull request #197 from jumpserver/pr@dev@fix_gson_issues

fix: make large SQL chunk transfer reliable

Author: Zxc4wonder

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

@@ -37,9 +37,8 @@
 import java.util.LinkedHashMap;
 import java.util.List;
 import java.util.Map;
-import java.util.concurrent.ConcurrentHashMap;
+import java.util.Optional;
 import java.util.concurrent.CountDownLatch;
-import java.util.concurrent.TimeUnit;
 import java.util.concurrent.atomic.AtomicBoolean;
 
 @Slf4j
@@ -51,6 +50,7 @@ public class QueryConsole extends AbstractConsole {
     private StateManager<QueryConsoleState> stateManager;
     private final Map<String, DataView> dataViews = new HashMap<>();
     private volatile Map<String, String> allowedContexts = Map.of();
+    private final SQLChunkTransferManager sqlChunkTransfers = new SQLChunkTransferManager();
 
     private static final Gson GSON = new Gson();
 
@@ -168,7 +168,7 @@ private void onAction(QueryConsoleAction action) {
                 this.handleSQLChunk(action);
             }
             case QueryConsoleAction.ACTION_RUN_SQL_COMPLETE -> {
-                this.handleSQLComplete();
+                this.handleSQLComplete(action);
             }
 
             case QueryConsoleAction.ACTION_RUN_SQL_FILE -> {
@@ -184,6 +184,7 @@ private void onAction(QueryConsoleAction action) {
 
 
             case QueryConsoleAction.ACTION_CANCEL -> {
+                this.sqlChunkTransfers.cancelAll();
                 this.onCancel();
                 this.getState().setInQuery(false);
                 this.stateManager.commit();
@@ -195,65 +196,26 @@ 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();
+        Optional<String> sql;
+        try {
+            sql = this.sqlChunkTransfers.receiveChunk(action.getData());
+        } catch (IllegalArgumentException | IllegalStateException e) {
+            this.getConsoleLogger().error("Invalid SQL chunk transfer: %s", e.getMessage());
+            return;
         }
+        sql.ifPresent(this::onSQL);
     }
 
-    /**
-     * 处理分段 SQL 接收完成
-     */
-    private void handleSQLComplete() {
+    private void handleSQLComplete(QueryConsoleAction action) {
+        Optional<String> sql;
         try {
-
-            // 等待所有分段接收完成
-            boolean completed = latch.await(10, TimeUnit.SECONDS); // 超时10秒
-
-            if (!completed) {
-                this.getConsoleLogger().error("read sql message timeout!！");
-                return;
-            }
-
-            // 按照索引顺序合并所有分段
-            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();
+            sql = this.sqlChunkTransfers.receiveComplete(action.getData());
+        } catch (IllegalArgumentException | IllegalStateException e) {
+            this.getConsoleLogger().error("Invalid SQL chunk transfer: %s", e.getMessage());
+            return;
         }
+        sql.ifPresent(this::onSQL);
     }
 
     private void onDataViewAction(DataViewAction action) {
@@ -538,6 +500,7 @@ private void sendDataView(DataView dataView, boolean clearOthers) {
 
     @Override
     public void close() {
+        this.sqlChunkTransfers.close();
         if (this.currentPlan != null) {
             // flush
             var session = SessionManager.getCurrentSession();