Merge pull request #81 from jumpserver/dev

v4.6.0

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";

backend/framework/src/main/java/org/jumpserver/chen/framework/console/dataview/DataView.java

@@ -20,8 +20,10 @@
 import java.nio.file.Files;
 import java.sql.Clob;
 import java.sql.SQLException;
+import java.text.SimpleDateFormat;
 import java.time.LocalDateTime;
 import java.time.format.DateTimeFormatter;
+import java.util.Date;
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
@@ -101,13 +103,17 @@ private void fullData(SQLQueryResult result) {
         this.data.getData().clear();
 
         this.getStateManager().getState().setTotal(result.getTotal());
+        this.fullDataViewData(this.data, result);
+    }
+
 
+    private void fullDataViewData(DataViewData viewData, SQLQueryResult result) {
 
-        this.data.setFields(result.getFields());
+        viewData.setFields(result.getFields());
 
         Map<String, Integer> fieldNumMap = new HashMap<>();
 
-        this.data.getFields().forEach(field -> {
+        viewData.getFields().forEach(field -> {
             if (fieldNumMap.containsKey(field.getName())) {
                 var fieldName = field.getName();
                 var num = fieldNumMap.get(field.getName());
@@ -118,16 +124,16 @@ private void fullData(SQLQueryResult result) {
             }
         });
 
-
         for (List<Object> row : result.getData()) {
             Map<String, Object> map = new HashMap<>();
             for (int i = 0; i < row.size(); i++) {
-                map.put(this.data.getFields().get(i).getName(), row.get(i));
+                map.put(viewData.getFields().get(i).getName(), row.get(i));
             }
-            this.data.getData().add(map);
+            viewData.getData().add(map);
         }
     }
 
+
     private static void writeString(BufferedWriter writer, Object object) throws IOException {
         var str = object.toString();
 
@@ -137,6 +143,35 @@ private static void writeString(BufferedWriter writer, Object object) throws IOE
         writer.write(str);
     }
 
+    private void writeCSVData(BufferedWriter writer, DataViewData viewData) throws IOException, SQLException {
+
+        for (Field field : viewData.getFields()) {
+            writeString(writer, field.getName());
+            writer.write(",");
+        }
+        for (Map<String, Object> row : viewData.getData()) {
+            for (Field field : viewData.getFields()) {
+                var obj = row.get(field.getName());
+                if (obj == null) {
+                    writer.write("NULL");
+                    writer.write(",");
+                } else if (obj instanceof Clob clob) {
+                    writer.write(CodeUtils.escapeCsvValue(clob.getSubString(1, (int) clob.length())));
+                    writer.write(",");
+                } else if (obj instanceof Date) {
+                    SimpleDateFormat fmt = new SimpleDateFormat("yyyy-MM-dd HH:mm:ss");
+                    writeString(writer, fmt.format(obj));
+                } else {
+                    writeString(writer, row.get(field.getName()));
+                    writer.write(",");
+                }
+            }
+            writer.newLine();
+        }
+
+        writer.newLine();
+    }
+
     public void export(String scope) throws SQLException {
         var session = SessionManager.getCurrentSession();
 
@@ -155,56 +190,16 @@ public void export(String scope) throws SQLException {
             var writer = Files.newBufferedWriter(f.toPath());
 
             if (scope.equals("current")) {
-                for (Field field : this.data.getFields()) {
-                    writeString(writer, field.getName());
-                    writer.write(",");
-                }
-                writer.newLine();
-
-                for (Map<String, Object> row : this.data.getData()) {
-                    for (Field field : this.data.getFields()) {
-                        if (row.get(field.getName()) == null) {
-                            writer.write("NULL");
-                            writer.write(",");
-                        } else if (row.get(field.getName()) instanceof Clob clob) {
-                            writer.write(CodeUtils.escapeCsvValue(clob.getSubString(1, (int) clob.length())));
-                            writer.write(",");
-                        } else {
-                            writeString(writer, row.get(field.getName()));
-                            writer.write(",");
-                        }
-                    }
-                    writer.newLine();
-                }
+                this.writeCSVData(writer, this.data);
                 command.setOutput(String.format("%d rows exported", this.data.getData().size()));
             }
 
             if (scope.equals("all")) {
                 SQLQueryParams queryParams = new SQLQueryParams();
                 queryParams.setLimit(-1);
                 var result = this.loadDataInterface.loadData(queryParams);
-
-                for (Field field : result.getFields()) {
-                    writer.write(field.getName());
-                    writer.write(",");
-                }
-                writer.newLine();
-
-                for (List<Object> row : result.getData()) {
-                    for (Object o : row) {
-                        if (o == null) {
-                            writer.write("NULL");
-                            writer.write(",");
-                        } else if (o instanceof Clob clob) {
-                            writer.write(CodeUtils.escapeCsvValue(clob.getSubString(1, (int) clob.length())));
-                            writer.write(",");
-                        } else {
-                            writer.write(o.toString());
-                            writer.write(",");
-                        }
-                    }
-                    writer.newLine();
-                }
+                var viewData = new DataViewData();
+                this.fullDataViewData(viewData, result);
                 command.setOutput(String.format("%d rows exported", result.getData().size()));
             }
             writer.flush();

backend/framework/src/main/java/org/jumpserver/chen/framework/datasource/base/BaseSQLActuator.java

@@ -10,6 +10,7 @@
 import com.alibaba.druid.sql.ast.statement.SQLUpdateStatement;
 import lombok.Getter;
 import lombok.extern.slf4j.Slf4j;
+import org.apache.commons.lang3.StringUtils;
 import org.jumpserver.chen.framework.datasource.ConnectionManager;
 import org.jumpserver.chen.framework.datasource.entity.resource.Field;
 import org.jumpserver.chen.framework.datasource.sql.*;
@@ -150,7 +151,10 @@ private void executeStatement(SQLExecutePlan plan, Statement statement, SQLQuery
 
                 for (int i = 1; i <= resultSet.getMetaData().getColumnCount(); i++) {
                     Field field = new Field();
-                    field.setName(resultSet.getMetaData().getColumnName(i));
+
+                    var fieldName = StringUtils.isNotEmpty(resultSet.getMetaData().getColumnLabel(i)) ?
+                            resultSet.getMetaData().getColumnLabel(i) : resultSet.getMetaData().getColumnName(i);
+                    field.setName(fieldName);
                     result.getFields().add(field);
                 }
 

backend/modules/src/main/java/org.jumpserver.chen.modules/oracle/OracleActuator.java

@@ -1,5 +1,7 @@
 package org.jumpserver.chen.modules.oracle;
 
+import com.alibaba.druid.DbType;
+import com.alibaba.druid.sql.SQLUtils;
 import org.jumpserver.chen.framework.datasource.ConnectionManager;
 import org.jumpserver.chen.framework.datasource.base.BaseSQLActuator;
 import org.jumpserver.chen.framework.datasource.sql.SQL;
@@ -59,4 +61,11 @@ public SQLExecutePlan createPlan(SQL sql) throws SQLException {
         this.beforeCreatePlan(sql);
         return super.createPlan(sql);
     }
+
+    @Override
+    public List<String> parseSQL(SQL sql) {
+        return SQLUtils.parseStatements(sql.getSql(), DbType.ali_oracle).stream()
+                .map(stmt -> SQLUtils.toSQLString(stmt, DbType.ali_oracle))
+                .toList();
+    }
 }

backend/modules/src/main/java/org.jumpserver.chen.modules/postgresql/PostgresqlActuator.java

@@ -1,5 +1,6 @@
 package org.jumpserver.chen.modules.postgresql;
 
+import org.apache.commons.lang3.StringUtils;
 import org.jumpserver.chen.framework.datasource.ConnectionManager;
 import org.jumpserver.chen.framework.datasource.base.BaseSQLActuator;
 import org.jumpserver.chen.framework.datasource.sql.SQL;
@@ -15,30 +16,45 @@ public PostgresqlActuator(ConnectionManager connectionManager) {
         super(connectionManager);
     }
 
+    private String dbName;
+
     public PostgresqlActuator(PostgresqlActuator sqlActuator, Connection connection) {
         super(sqlActuator, connection);
     }
 
     @Override
     public String getCurrentSchema() throws SQLException {
-        var result = this.execute(SQL.of("SELECT CURRENT_SCHEMA()"));
-        return (String) result.getData().get(0).get(0);
+        var result = this.execute(SQL.of("SELECT current_schema()"));
+        return this.formatSchemaName((String) result.getData().get(0).get(0));
     }
 
     @Override
     public List<String> getSchemas() throws SQLException {
         var result = this.execute(SQL.of("SELECT SCHEMA_NAME FROM INFORMATION_SCHEMA.SCHEMATA"));
-        return result.getData().stream().map(row -> (String) row.get(0)).toList();
+        return result.getData().stream().map(row -> (String) row.get(0)).toList().stream().map(this::formatSchemaName).toList();
     }
 
     @Override
     public void changeSchema(String schema) throws SQLException {
-        this.execute(SQL.of("SET SEARCH_PATH TO '?';", schema));
+        var ss = schema.split("\\.");
+        this.execute(SQL.of("SET SEARCH_PATH TO '?';", ss[1]));
     }
 
     @Override
     public SQLExecutePlan createPlan(String schema, String table, SQLQueryParams sqlQueryParams) throws SQLException {
         var sql = SQL.of("select * from \"?\".\"?\"", schema, table);
         return this.createPlan(sql, sqlQueryParams);
     }
+
+    private String formatSchemaName(String schema) {
+        try {
+            if (StringUtils.isEmpty(this.dbName)) {
+                var result = this.execute(SQL.of("SELECT current_database();"));
+                this.dbName = (String) result.getData().get(0).get(0);
+            }
+        } catch (SQLException e) {
+            throw new RuntimeException(e);
+        }
+        return String.format("%s.%s", this.dbName, schema);
+    }
 }

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) {