Skip to content

Commit fb0fe89

Browse files
baijiangjiebaijiangjie
authored
Merge pull request #81 from jumpserver/dev
v4.6.0
2 parents 9eec12f + 6634e54 commit fb0fe89

8 files changed

Lines changed: 160 additions & 59 deletions

File tree

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

Lines changed: 63 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -34,6 +34,7 @@
3434
import java.util.ArrayList;
3535
import java.util.HashMap;
3636
import java.util.Map;
37+
import java.util.concurrent.ConcurrentHashMap;
3738
import java.util.concurrent.CountDownLatch;
3839
import java.util.concurrent.atomic.AtomicBoolean;
3940

@@ -141,6 +142,7 @@ public void handle(Packet packet) {
141142
}
142143
}
143144

145+
144146
private void onAction(QueryConsoleAction action) {
145147
switch (action.getAction()) {
146148
case QueryConsoleAction.ACTION_RUN_SQL -> {
@@ -153,6 +155,13 @@ private void onAction(QueryConsoleAction action) {
153155
this.getState().setInQuery(false);
154156
this.stateManager.commit();
155157
}
158+
case QueryConsoleAction.ACTION_RUN_SQL_CHUNK -> {
159+
this.handleSQLChunk(action);
160+
}
161+
case QueryConsoleAction.ACTION_RUN_SQL_COMPLETE -> {
162+
this.handleSQLComplete();
163+
}
164+
156165
case QueryConsoleAction.ACTION_RUN_SQL_FILE -> {
157166
this.getState().setInQuery(true);
158167
this.stateManager.commit();
@@ -177,6 +186,60 @@ private void onAction(QueryConsoleAction action) {
177186
}
178187
}
179188

189+
private final ConcurrentHashMap<Integer, String> sqlChunks = new ConcurrentHashMap<>();
190+
private CountDownLatch latch;
191+
private int expectedChunks = -1;
192+
193+
private void handleSQLChunk(QueryConsoleAction action) {
194+
var data = (Map<String, Object>) action.getData();
195+
var chunk = (String) data.get("chunk");
196+
var index = (Integer) data.get("index");
197+
var total = (Integer) data.get("total");
198+
199+
synchronized (this) {
200+
if (expectedChunks == -1) {
201+
expectedChunks = total;
202+
latch = new CountDownLatch(total);
203+
}
204+
}
205+
206+
if (sqlChunks.putIfAbsent(index, chunk) == null) {
207+
latch.countDown();
208+
}
209+
}
210+
211+
/**
212+
* 处理分段 SQL 接收完成
213+
*/
214+
private void handleSQLComplete() {
215+
try {
216+
// 等待所有分段接收完成
217+
latch.await();
218+
219+
// 按照索引顺序合并所有分段
220+
StringBuilder sqlBuilder = new StringBuilder();
221+
for (int i = 0; i < expectedChunks; i++) {
222+
sqlBuilder.append(sqlChunks.get(i));
223+
}
224+
225+
// 合并完成后清理缓存
226+
var sql = sqlBuilder.toString();
227+
sqlChunks.clear();
228+
expectedChunks = -1;
229+
230+
// 执行完整 SQL
231+
this.getState().setInQuery(true);
232+
this.stateManager.commit();
233+
234+
this.onSQL(sql);
235+
236+
} catch (InterruptedException e) {
237+
Thread.currentThread().interrupt();
238+
} finally {
239+
this.getState().setInQuery(false);
240+
this.stateManager.commit();
241+
}
242+
}
180243

181244
private void onDataViewAction(DataViewAction action) {
182245
var dataView = this.dataViews.get(action.getDataView());

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

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -7,6 +7,8 @@
77
@EqualsAndHashCode(callSuper = true)
88
public class QueryConsoleAction extends Action {
99
public static final String ACTION_RUN_SQL = "run_sql";
10+
public static final String ACTION_RUN_SQL_CHUNK = "run_sql_chunk";
11+
public static final String ACTION_RUN_SQL_COMPLETE = "run_sql_complete";
1012
public static final String ACTION_RUN_SQL_FILE = "run_sql_file";
1113
public static final String ACTION_CANCEL = "cancel";
1214
public static final String ACTION_CHANGE_CURRENT_CONTEXT = "change_current_context";

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

Lines changed: 43 additions & 48 deletions
Original file line numberDiff line numberDiff line change
@@ -20,8 +20,10 @@
2020
import java.nio.file.Files;
2121
import java.sql.Clob;
2222
import java.sql.SQLException;
23+
import java.text.SimpleDateFormat;
2324
import java.time.LocalDateTime;
2425
import java.time.format.DateTimeFormatter;
26+
import java.util.Date;
2527
import java.util.HashMap;
2628
import java.util.List;
2729
import java.util.Map;
@@ -101,13 +103,17 @@ private void fullData(SQLQueryResult result) {
101103
this.data.getData().clear();
102104

103105
this.getStateManager().getState().setTotal(result.getTotal());
106+
this.fullDataViewData(this.data, result);
107+
}
108+
104109

110+
private void fullDataViewData(DataViewData viewData, SQLQueryResult result) {
105111

106-
this.data.setFields(result.getFields());
112+
viewData.setFields(result.getFields());
107113

108114
Map<String, Integer> fieldNumMap = new HashMap<>();
109115

110-
this.data.getFields().forEach(field -> {
116+
viewData.getFields().forEach(field -> {
111117
if (fieldNumMap.containsKey(field.getName())) {
112118
var fieldName = field.getName();
113119
var num = fieldNumMap.get(field.getName());
@@ -118,16 +124,16 @@ private void fullData(SQLQueryResult result) {
118124
}
119125
});
120126

121-
122127
for (List<Object> row : result.getData()) {
123128
Map<String, Object> map = new HashMap<>();
124129
for (int i = 0; i < row.size(); i++) {
125-
map.put(this.data.getFields().get(i).getName(), row.get(i));
130+
map.put(viewData.getFields().get(i).getName(), row.get(i));
126131
}
127-
this.data.getData().add(map);
132+
viewData.getData().add(map);
128133
}
129134
}
130135

136+
131137
private static void writeString(BufferedWriter writer, Object object) throws IOException {
132138
var str = object.toString();
133139

@@ -137,6 +143,35 @@ private static void writeString(BufferedWriter writer, Object object) throws IOE
137143
writer.write(str);
138144
}
139145

146+
private void writeCSVData(BufferedWriter writer, DataViewData viewData) throws IOException, SQLException {
147+
148+
for (Field field : viewData.getFields()) {
149+
writeString(writer, field.getName());
150+
writer.write(",");
151+
}
152+
for (Map<String, Object> row : viewData.getData()) {
153+
for (Field field : viewData.getFields()) {
154+
var obj = row.get(field.getName());
155+
if (obj == null) {
156+
writer.write("NULL");
157+
writer.write(",");
158+
} else if (obj instanceof Clob clob) {
159+
writer.write(CodeUtils.escapeCsvValue(clob.getSubString(1, (int) clob.length())));
160+
writer.write(",");
161+
} else if (obj instanceof Date) {
162+
SimpleDateFormat fmt = new SimpleDateFormat("yyyy-MM-dd HH:mm:ss");
163+
writeString(writer, fmt.format(obj));
164+
} else {
165+
writeString(writer, row.get(field.getName()));
166+
writer.write(",");
167+
}
168+
}
169+
writer.newLine();
170+
}
171+
172+
writer.newLine();
173+
}
174+
140175
public void export(String scope) throws SQLException {
141176
var session = SessionManager.getCurrentSession();
142177

@@ -155,56 +190,16 @@ public void export(String scope) throws SQLException {
155190
var writer = Files.newBufferedWriter(f.toPath());
156191

157192
if (scope.equals("current")) {
158-
for (Field field : this.data.getFields()) {
159-
writeString(writer, field.getName());
160-
writer.write(",");
161-
}
162-
writer.newLine();
163-
164-
for (Map<String, Object> row : this.data.getData()) {
165-
for (Field field : this.data.getFields()) {
166-
if (row.get(field.getName()) == null) {
167-
writer.write("NULL");
168-
writer.write(",");
169-
} else if (row.get(field.getName()) instanceof Clob clob) {
170-
writer.write(CodeUtils.escapeCsvValue(clob.getSubString(1, (int) clob.length())));
171-
writer.write(",");
172-
} else {
173-
writeString(writer, row.get(field.getName()));
174-
writer.write(",");
175-
}
176-
}
177-
writer.newLine();
178-
}
193+
this.writeCSVData(writer, this.data);
179194
command.setOutput(String.format("%d rows exported", this.data.getData().size()));
180195
}
181196

182197
if (scope.equals("all")) {
183198
SQLQueryParams queryParams = new SQLQueryParams();
184199
queryParams.setLimit(-1);
185200
var result = this.loadDataInterface.loadData(queryParams);
186-
187-
for (Field field : result.getFields()) {
188-
writer.write(field.getName());
189-
writer.write(",");
190-
}
191-
writer.newLine();
192-
193-
for (List<Object> row : result.getData()) {
194-
for (Object o : row) {
195-
if (o == null) {
196-
writer.write("NULL");
197-
writer.write(",");
198-
} else if (o instanceof Clob clob) {
199-
writer.write(CodeUtils.escapeCsvValue(clob.getSubString(1, (int) clob.length())));
200-
writer.write(",");
201-
} else {
202-
writer.write(o.toString());
203-
writer.write(",");
204-
}
205-
}
206-
writer.newLine();
207-
}
201+
var viewData = new DataViewData();
202+
this.fullDataViewData(viewData, result);
208203
command.setOutput(String.format("%d rows exported", result.getData().size()));
209204
}
210205
writer.flush();

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

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -10,6 +10,7 @@
1010
import com.alibaba.druid.sql.ast.statement.SQLUpdateStatement;
1111
import lombok.Getter;
1212
import lombok.extern.slf4j.Slf4j;
13+
import org.apache.commons.lang3.StringUtils;
1314
import org.jumpserver.chen.framework.datasource.ConnectionManager;
1415
import org.jumpserver.chen.framework.datasource.entity.resource.Field;
1516
import org.jumpserver.chen.framework.datasource.sql.*;
@@ -150,7 +151,10 @@ private void executeStatement(SQLExecutePlan plan, Statement statement, SQLQuery
150151

151152
for (int i = 1; i <= resultSet.getMetaData().getColumnCount(); i++) {
152153
Field field = new Field();
153-
field.setName(resultSet.getMetaData().getColumnName(i));
154+
155+
var fieldName = StringUtils.isNotEmpty(resultSet.getMetaData().getColumnLabel(i)) ?
156+
resultSet.getMetaData().getColumnLabel(i) : resultSet.getMetaData().getColumnName(i);
157+
field.setName(fieldName);
154158
result.getFields().add(field);
155159
}
156160

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

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,7 @@
11
package org.jumpserver.chen.modules.oracle;
22

3+
import com.alibaba.druid.DbType;
4+
import com.alibaba.druid.sql.SQLUtils;
35
import org.jumpserver.chen.framework.datasource.ConnectionManager;
46
import org.jumpserver.chen.framework.datasource.base.BaseSQLActuator;
57
import org.jumpserver.chen.framework.datasource.sql.SQL;
@@ -59,4 +61,11 @@ public SQLExecutePlan createPlan(SQL sql) throws SQLException {
5961
this.beforeCreatePlan(sql);
6062
return super.createPlan(sql);
6163
}
64+
65+
@Override
66+
public List<String> parseSQL(SQL sql) {
67+
return SQLUtils.parseStatements(sql.getSql(), DbType.ali_oracle).stream()
68+
.map(stmt -> SQLUtils.toSQLString(stmt, DbType.ali_oracle))
69+
.toList();
70+
}
6271
}

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

Lines changed: 20 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,6 @@
11
package org.jumpserver.chen.modules.postgresql;
22

3+
import org.apache.commons.lang3.StringUtils;
34
import org.jumpserver.chen.framework.datasource.ConnectionManager;
45
import org.jumpserver.chen.framework.datasource.base.BaseSQLActuator;
56
import org.jumpserver.chen.framework.datasource.sql.SQL;
@@ -15,30 +16,45 @@ public PostgresqlActuator(ConnectionManager connectionManager) {
1516
super(connectionManager);
1617
}
1718

19+
private String dbName;
20+
1821
public PostgresqlActuator(PostgresqlActuator sqlActuator, Connection connection) {
1922
super(sqlActuator, connection);
2023
}
2124

2225
@Override
2326
public String getCurrentSchema() throws SQLException {
24-
var result = this.execute(SQL.of("SELECT CURRENT_SCHEMA()"));
25-
return (String) result.getData().get(0).get(0);
27+
var result = this.execute(SQL.of("SELECT current_schema()"));
28+
return this.formatSchemaName((String) result.getData().get(0).get(0));
2629
}
2730

2831
@Override
2932
public List<String> getSchemas() throws SQLException {
3033
var result = this.execute(SQL.of("SELECT SCHEMA_NAME FROM INFORMATION_SCHEMA.SCHEMATA"));
31-
return result.getData().stream().map(row -> (String) row.get(0)).toList();
34+
return result.getData().stream().map(row -> (String) row.get(0)).toList().stream().map(this::formatSchemaName).toList();
3235
}
3336

3437
@Override
3538
public void changeSchema(String schema) throws SQLException {
36-
this.execute(SQL.of("SET SEARCH_PATH TO '?';", schema));
39+
var ss = schema.split("\\.");
40+
this.execute(SQL.of("SET SEARCH_PATH TO '?';", ss[1]));
3741
}
3842

3943
@Override
4044
public SQLExecutePlan createPlan(String schema, String table, SQLQueryParams sqlQueryParams) throws SQLException {
4145
var sql = SQL.of("select * from \"?\".\"?\"", schema, table);
4246
return this.createPlan(sql, sqlQueryParams);
4347
}
48+
49+
private String formatSchemaName(String schema) {
50+
try {
51+
if (StringUtils.isEmpty(this.dbName)) {
52+
var result = this.execute(SQL.of("SELECT current_database();"));
53+
this.dbName = (String) result.getData().get(0).get(0);
54+
}
55+
} catch (SQLException e) {
56+
throw new RuntimeException(e);
57+
}
58+
return String.format("%s.%s", this.dbName, schema);
59+
}
4460
}

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

Lines changed: 18 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -253,7 +253,24 @@ export default {
253253
},
254254
onRun() {
255255
const sql = this.selectionValue || this.statement
256-
this.$emit('action', { action: 'run_sql', data: sql })
256+
const CHUNK_SIZE = 4096
257+
258+
if (sql.length <= CHUNK_SIZE) {
259+
this.$emit('action', { action: 'run_sql', data: sql })
260+
} else {
261+
const totalChunks = Math.ceil(sql.length / CHUNK_SIZE)
262+
for (let i = 0; i < totalChunks; i++) {
263+
const chunk = sql.slice(i * CHUNK_SIZE, (i + 1) * CHUNK_SIZE)
264+
this.$emit('action', {
265+
action: 'run_sql_chunk',
266+
data: { chunk, index: i, total: totalChunks }
267+
})
268+
}
269+
this.$emit('action', {
270+
action: 'run_sql_complete',
271+
data: { total: totalChunks }
272+
})
273+
}
257274
},
258275
onStop() {
259276
this.$emit('action', { action: 'cancel' })

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

Lines changed: 0 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -8,7 +8,6 @@
88
:state="state"
99
:subjects="subjects"
1010
@action="onEditorAction"
11-
@run="onRunSql"
1211
/>
1312
</template>
1413
<template slot="paneR">
@@ -153,11 +152,7 @@ export default {
153152
onDataViewAction(action) {
154153
this.ws.send(JSON.stringify({ type: 'data_view_action', data: action }))
155154
},
156-
onRunSql(sql) {
157-
this.ws.send(JSON.stringify({ type: 'sql', data: sql }))
158-
},
159155
onCloseDataView(name) {
160-
console.log(name)
161156
this.ws.send(JSON.stringify({ type: 'close_data_view', data: name }))
162157
},
163158
onLimitChange(limit) {

0 commit comments

Comments
 (0)