Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

v4.6.0 #81

Merged
merged 10 commits into from
Jan 15, 2025
Original file line number Diff line number Diff line change
Expand Up @@ -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;

Expand Down Expand Up @@ -141,6 +142,7 @@ public void handle(Packet packet) {
}
}


private void onAction(QueryConsoleAction action) {
switch (action.getAction()) {
case QueryConsoleAction.ACTION_RUN_SQL -> {
Expand All @@ -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();
Expand All @@ -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());
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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";
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -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());
Expand All @@ -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();

Expand All @@ -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();

Expand All @@ -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();
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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.*;
Expand Down Expand Up @@ -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);
}

Expand Down
Original file line number Diff line number Diff line change
@@ -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;
Expand Down Expand Up @@ -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();
}
}
Original file line number Diff line number Diff line change
@@ -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;
Expand All @@ -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);
}
}
19 changes: 18 additions & 1 deletion frontend/src/components/Main/Explore/QueryConsole/CodeEditor.vue
Original file line number Diff line number Diff line change
Expand Up @@ -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' })
Expand Down
5 changes: 0 additions & 5 deletions frontend/src/components/Main/Explore/QueryConsole/index.vue
Original file line number Diff line number Diff line change
Expand Up @@ -8,7 +8,6 @@
:state="state"
:subjects="subjects"
@action="onEditorAction"
@run="onRunSql"
/>
</template>
<template slot="paneR">
Expand Down Expand Up @@ -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) {
Expand Down
Loading