From ec947b2ec50d135561622237ce72bbd657ab42ec Mon Sep 17 00:00:00 2001 From: "Crane.z" <1481445951@qq.com> Date: Mon, 17 Aug 2026 11:05:12 +0800 Subject: [PATCH 1/2] feat(console): support native SQL execution --- .../ConsoleStatementBoundaryScanner.java | 161 ++++++++++++++++++ .../chen/framework/console/QueryConsole.java | 122 ++++++++++--- .../datasource/base/BaseSQLActuator.java | 92 ++++++---- .../framework/datasource/sql/SQLActuator.java | 4 + 4 files changed, 321 insertions(+), 58 deletions(-) create mode 100644 backend/framework/src/main/java/org/jumpserver/chen/framework/console/ConsoleStatementBoundaryScanner.java diff --git a/backend/framework/src/main/java/org/jumpserver/chen/framework/console/ConsoleStatementBoundaryScanner.java b/backend/framework/src/main/java/org/jumpserver/chen/framework/console/ConsoleStatementBoundaryScanner.java new file mode 100644 index 0000000..51d72dc --- /dev/null +++ b/backend/framework/src/main/java/org/jumpserver/chen/framework/console/ConsoleStatementBoundaryScanner.java @@ -0,0 +1,161 @@ +package org.jumpserver.chen.framework.console; + +import java.sql.SQLException; + +/** + * Finds top-level SQL statement boundaries without parsing database syntax. + */ +final class ConsoleStatementBoundaryScanner { + private static final String MULTIPLE_STATEMENTS_ERROR = + "Console raw execution accepts only one top-level database command"; + + private ConsoleStatementBoundaryScanner() { + } + + static void requireSingleStatement(String sql) throws SQLException { + if (sql == null || sql.isBlank()) { + throw new SQLException("Console raw execution requires a database command"); + } + + int completedStatements = 0; + boolean hasStatementContent = false; + int blockCommentDepth = 0; + char quote = 0; + String dollarQuote = null; + + for (int index = 0; index < sql.length(); index++) { + char current = sql.charAt(index); + char next = index + 1 < sql.length() ? sql.charAt(index + 1) : 0; + + if (dollarQuote != null) { + if (sql.startsWith(dollarQuote, index)) { + index += dollarQuote.length() - 1; + dollarQuote = null; + } + continue; + } + + if (quote != 0) { + if (current == '\\' && index + 1 < sql.length()) { + index++; + continue; + } + if (current == quote) { + if (next == quote) { + index++; + } else { + quote = 0; + } + } + continue; + } + + if (blockCommentDepth > 0) { + if (current == '/' && next == '*') { + blockCommentDepth++; + index++; + } else if (current == '*' && next == '/') { + blockCommentDepth--; + index++; + } + continue; + } + + if (current == '-' && next == '-') { + index = skipLineComment(sql, index + 2); + continue; + } + if (current == '/' && next == '*') { + blockCommentDepth = 1; + index++; + continue; + } + if (current == '\'' || current == '"' || current == '`') { + quote = current; + hasStatementContent = true; + continue; + } + if (current == '$') { + String delimiter = dollarQuoteDelimiter(sql, index); + if (delimiter != null) { + dollarQuote = delimiter; + hasStatementContent = true; + index += delimiter.length() - 1; + continue; + } + } + if (current == ';') { + if (hasStatementContent) { + completedStatements++; + hasStatementContent = false; + if (completedStatements > 1) { + throw new SQLException(MULTIPLE_STATEMENTS_ERROR); + } + } + continue; + } + if (!Character.isWhitespace(current)) { + if (completedStatements > 0) { + throw new SQLException(MULTIPLE_STATEMENTS_ERROR); + } + hasStatementContent = true; + } + } + + if (hasStatementContent) { + completedStatements++; + } + if (completedStatements == 0) { + throw new SQLException("Console raw execution requires a database command"); + } + if (completedStatements > 1) { + throw new SQLException(MULTIPLE_STATEMENTS_ERROR); + } + } + + private static int skipLineComment(String sql, int index) { + while (index < sql.length()) { + char current = sql.charAt(index); + if (current == '\n' || current == '\r') { + return index; + } + index++; + } + return sql.length() - 1; + } + + private static String dollarQuoteDelimiter(String sql, int start) { + if (start > 0 && isIdentifierPart(sql.charAt(start - 1))) { + return null; + } + int end = sql.indexOf('$', start + 1); + if (end < 0) { + return null; + } + String tag = sql.substring(start + 1, end); + if (!tag.isEmpty()) { + if (!isTagStart(tag.charAt(0))) { + return null; + } + for (int index = 1; index < tag.length(); index++) { + if (!isTagPart(tag.charAt(index))) { + return null; + } + } + } + String delimiter = sql.substring(start, end + 1); + return sql.indexOf(delimiter, end + 1) >= 0 ? delimiter : null; + } + + private static boolean isTagStart(char value) { + return value == '_' || Character.isLetter(value); + } + + private static boolean isTagPart(char value) { + return isTagStart(value) || Character.isDigit(value); + } + + private static boolean isIdentifierPart(char value) { + return isTagPart(value) || value == '$'; + } +} diff --git a/backend/framework/src/main/java/org/jumpserver/chen/framework/console/QueryConsole.java b/backend/framework/src/main/java/org/jumpserver/chen/framework/console/QueryConsole.java index 1eeaac2..03249b3 100644 --- a/backend/framework/src/main/java/org/jumpserver/chen/framework/console/QueryConsole.java +++ b/backend/framework/src/main/java/org/jumpserver/chen/framework/console/QueryConsole.java @@ -31,6 +31,7 @@ import org.jumpserver.chen.framework.datasource.sql.SQL; import org.jumpserver.chen.framework.datasource.sql.SQLActuator; import org.jumpserver.chen.framework.datasource.sql.SQLExecutePlan; +import org.jumpserver.chen.framework.datasource.sql.SQLQueryResult; import org.jumpserver.chen.framework.i18n.MessageUtils; import org.jumpserver.chen.framework.jms.acl.ACLResult; import org.jumpserver.chen.framework.jms.entity.CommandRecord; @@ -52,6 +53,7 @@ import java.nio.file.StandardOpenOption; import java.sql.Connection; import java.sql.SQLException; +import java.sql.Statement; import java.util.ArrayList; import java.util.Collections; import java.util.HashMap; @@ -99,7 +101,7 @@ public class QueryConsole extends AbstractConsole { // SQL Server, Dameng or DB2, so the server-side probe is the real source of truth for whether // a user transaction is open before a DataView save commits or stages behind a savepoint. private volatile QueryTransactionStateInspector transactionStateInspector; - private volatile SQLExecutePlan currentPlan; + private volatile ActiveExecution currentExecution; private StateManager stateManager; private final Map dataViews = new LinkedHashMap<>(); // Manual context changes remain restricted to values returned by the current server-side actuator. @@ -633,10 +635,10 @@ private SaveChangesResult rejectedSave(DataView dataView, String reason) { public void onCancel() { this.getState().setExecutionStatus(EXECUTION_STATUS_CANCELLED); try { - var plan = this.currentPlan; - if (plan != null) { - plan.cancel(); - this.getConsoleLogger().warn("cancel query: %s", plan.getTargetSQL()); + var execution = this.currentExecution; + if (execution != null) { + execution.cancelAction().cancel(); + this.getConsoleLogger().warn("cancel query: %s", execution.sql()); } } catch (SQLException | RuntimeException e) { log.error("cancel failed ", e); @@ -767,22 +769,10 @@ public void onSQL(String sql) { var session = SessionManager.getCurrentSession(); try { - var stmts = this.getSqlActuator().parseSQL(SQL.of(sql)); - var clearOthers = true; - for (String stmt : stmts) { - var aclResult = session.checkACL(stmt, this.getConnection()); - if (!this.canExecuteStatement(session, stmt, aclResult)) { - break; - } - var dataView = this.runSingleSQL(stmt, aclResult); - if (!dataView.isHasTable()) { - this.getConsoleLogger().success("%s , %s: %d", - MessageUtils.get("ExecuteSuccess"), - MessageUtils.get("AffectedRows"), dataView.getUpdateCount()); - } else { - this.sendDataView(dataView, clearOthers); - clearOthers = false; - } + if (this.consoleMode) { + this.runRawConsoleSQL(sql, session); + } else { + this.runQuerySQL(sql, session); } this.ensureCurrentSchema(); } catch (ParserException e) { @@ -807,6 +797,75 @@ public void onSQL(String sql) { } } + private void runQuerySQL(String sql, Session session) throws SQLException { + var statements = this.getSqlActuator().parseSQL(SQL.of(sql)); + var clearOthers = true; + for (String statement : statements) { + var aclResult = session.checkACL(statement, this.getConnection()); + if (!this.canExecuteStatement(session, statement, aclResult)) { + break; + } + var dataView = this.runSingleSQL(statement, aclResult); + if (!dataView.isHasTable()) { + this.logAffectedRows(dataView); + } else { + this.sendDataView(dataView, clearOthers); + clearOthers = false; + } + } + } + + private void runRawConsoleSQL(String sql, Session session) throws SQLException { + ConsoleStatementBoundaryScanner.requireSingleStatement(sql); + Connection connection = this.getConnection(); + ACLResult aclResult = session.checkACL(sql, connection); + if (!this.canExecuteStatement(session, sql, aclResult)) { + return; + } + + DataView dataView = new DataView( + UUID.randomUUID().toString(), sql, this.getPacketIO(), this.getConsoleLogger() + ); + dataView.setSql(sql); + dataView.setLoadDataInterface((ignored) -> this.executeRawConsoleSQL(sql, aclResult, connection)); + dataView.loadData(); + if (!dataView.isHasTable()) { + this.logAffectedRows(dataView); + } else { + this.sendDataView(dataView, true); + } + } + + private SQLQueryResult executeRawConsoleSQL(String sql, ACLResult aclResult, Connection connection) + throws SQLException { + SQLActuator actuator = this.datasource.getConnectionManager().getSqlActuator().withConnection(connection); + SQLExecutePlan plan = actuator.createPlan(SQL.of(sql)); + plan.setAclResult(aclResult); + Statement statement = plan.createStatement(); + ActiveExecution execution = new ActiveExecution(sql, statement::cancel); + this.currentExecution = execution; + this.getState().setCanCancel(true); + this.stateManager.commit(); + try { + this.getConsoleLogger().info("execute raw sql: %s", sql); + SQLQueryResult result = actuator.executeRawWithAudit(plan); + this.getConsoleLogger().success(result); + return result; + } finally { + if (this.currentExecution == execution) { + this.currentExecution = null; + } + this.getState().setCanCancel(false); + this.stateManager.commit(); + } + } + + private void logAffectedRows(DataView dataView) { + this.getConsoleLogger().success("%s , %s: %d", + MessageUtils.get("ExecuteSuccess"), + MessageUtils.get("AffectedRows"), dataView.getUpdateCount()); + } + private void sendSQLError(String kind, String title, String message, String sql, SQLException exception) { var error = new LinkedHashMap(); error.put("kind", kind); @@ -928,7 +987,8 @@ private DataView runSingleSQL(String sql, ACLResult aclResult) throws SQLExcepti .createPlan(SQL.of(sourceSQL)); plan.setAclResult(aclResult); plan.setSqlQueryParams(sqlQueryParams); - this.currentPlan = plan; + ActiveExecution execution = new ActiveExecution(sourceSQL, plan::cancel); + this.currentExecution = execution; this.getState().setCanCancel(true); this.stateManager.commit(); @@ -939,7 +999,9 @@ private DataView runSingleSQL(String sql, ACLResult aclResult) throws SQLExcepti this.getConsoleLogger().success(result); return result; } finally { - this.currentPlan = null; + if (this.currentExecution == execution) { + this.currentExecution = null; + } this.getState().setCanCancel(false); this.stateManager.commit(); } @@ -1019,12 +1081,12 @@ public void close() { currentSession.getController().cancelDialogs(this.getPacketIO().getWsSession().getId()); } - var plan = this.currentPlan; - if (plan != null) { + var execution = this.currentExecution; + if (execution != null) { try { // flush var session = SessionManager.getCurrentSession(); - var lastCmd = plan.getTargetSQL(); + var lastCmd = execution.sql(); var cmdRecord = new CommandRecord(lastCmd); cmdRecord.setError("Abnormal exit"); if (session != null) { @@ -1044,4 +1106,12 @@ public void close() { } log.info("console closed"); } + + @FunctionalInterface + private interface ExecutionCancel { + void cancel() throws SQLException; + } + + private record ActiveExecution(String sql, ExecutionCancel cancelAction) { + } } diff --git a/backend/framework/src/main/java/org/jumpserver/chen/framework/datasource/base/BaseSQLActuator.java b/backend/framework/src/main/java/org/jumpserver/chen/framework/datasource/base/BaseSQLActuator.java index dbe23db..dedc8aa 100644 --- a/backend/framework/src/main/java/org/jumpserver/chen/framework/datasource/base/BaseSQLActuator.java +++ b/backend/framework/src/main/java/org/jumpserver/chen/framework/datasource/base/BaseSQLActuator.java @@ -133,12 +133,21 @@ public SQLQueryResult execute(SQL sql) throws SQLException { @Override public SQLQueryResult execute(SQLExecutePlan plan) throws SQLException { + return this.execute(plan, true); + } + + @Override + public SQLQueryResult executeRaw(SQLExecutePlan plan) throws SQLException { + return this.execute(plan, false); + } + + private SQLQueryResult execute(SQLExecutePlan plan, boolean enrichResult) throws SQLException { String sql = plan.getTargetSQL(); SQLQueryResult result = new SQLQueryResult(sql); result.setAclResult(plan.getAclResult()); try { Statement statement = plan.createStatement(); - this.executeStatement(plan, statement, result); + this.executeStatement(plan, statement, result, enrichResult); } finally { if (plan.getConnection() instanceof DruidPooledConnection) { plan.getConnection().close(); @@ -147,7 +156,12 @@ public SQLQueryResult execute(SQLExecutePlan plan) throws SQLException { return result; } - private void executeStatement(SQLExecutePlan plan, Statement statement, SQLQueryResult result) throws SQLException { + private void executeStatement( + SQLExecutePlan plan, + Statement statement, + SQLQueryResult result, + boolean enrichResult + ) throws SQLException { try (statement) { result.setStartTime(new Time(System.currentTimeMillis())); @@ -157,41 +171,26 @@ private void executeStatement(SQLExecutePlan plan, Statement statement, SQLQuery result.setQueryFinishedTime(new Time(System.currentTimeMillis())); if (hasResult) { - var resultSet = statement.getResultSet(); - var metaData = resultSet.getMetaData(); - var columnCount = metaData.getColumnCount(); - - for (int i = 1; i <= columnCount; i++) { - result.getFields().add(buildField(metaData, i)); - } - - while (resultSet.next()) { - List fs = new ArrayList<>(); - for (int i = 1; i <= columnCount; i++) { - try { - fs.add(this.normalizeJdbcValue(resultSet.getObject(i))); - } catch (NoClassDefFoundError e) { - log.error(e.getMessage()); - } - } - result.getData().add(fs); + try (ResultSet resultSet = statement.getResultSet()) { + this.readResultSet(resultSet, result); } - resultSet.close(); - markGeneratedColumns(plan.getConnection(), this.getDruidDbType(), result.getFields()); result.setFetchFinishedTime(new Time(System.currentTimeMillis())); - this.analyzeResultEditability(plan, result); - - // 数据脱敏 + if (enrichResult) { + markGeneratedColumns(plan.getConnection(), this.getDruidDbType(), result.getFields()); + this.analyzeResultEditability(plan, result); + } this.handleDataMasking(result); - - var total = this.count(plan); - if (total < 0) { - result.setTotal(result.getData().size()); + if (enrichResult) { + var total = this.count(plan); + if (total < 0) { + result.setTotal(result.getData().size()); + } else { + result.setPaged(true); + result.setTotal(total); + } } else { - result.setPaged(true); - result.setTotal(total); + result.setTotal(result.getData().size()); } - } else { result.setUpdateCount(statement.getUpdateCount()); } @@ -201,6 +200,25 @@ private void executeStatement(SQLExecutePlan plan, Statement statement, SQLQuery } } + private void readResultSet(ResultSet resultSet, SQLQueryResult result) throws SQLException { + var metaData = resultSet.getMetaData(); + var columnCount = metaData.getColumnCount(); + for (int index = 1; index <= columnCount; index++) { + result.getFields().add(buildField(metaData, index)); + } + while (resultSet.next()) { + List row = new ArrayList<>(); + for (int index = 1; index <= columnCount; index++) { + try { + row.add(this.normalizeJdbcValue(resultSet.getObject(index))); + } catch (NoClassDefFoundError e) { + log.error(e.getMessage()); + } + } + result.getData().add(row); + } + } + private void analyzeResultEditability(SQLExecutePlan plan, SQLQueryResult result) { var analyzer = new QueryResultEditabilityAnalyzer( this.getDruidDbType(), @@ -627,6 +645,16 @@ public SQLQueryResult executeWithAudit(SQLExecutePlan plan) throws SQLException } } + @Override + public SQLQueryResult executeRawWithAudit(SQLExecutePlan plan) throws SQLException { + var sess = SessionManager.getCurrentSession(); + try { + return sess.withAudit(plan.getTargetSQL(), () -> this.executeRaw(plan)); + } catch (CommandRejectException e) { + throw new SQLException(e.getMessage()); + } + } + public int count(SQL sql) throws SQLException { return this.count(this.createPlan(sql)); } diff --git a/backend/framework/src/main/java/org/jumpserver/chen/framework/datasource/sql/SQLActuator.java b/backend/framework/src/main/java/org/jumpserver/chen/framework/datasource/sql/SQLActuator.java index d019005..d2490a7 100644 --- a/backend/framework/src/main/java/org/jumpserver/chen/framework/datasource/sql/SQLActuator.java +++ b/backend/framework/src/main/java/org/jumpserver/chen/framework/datasource/sql/SQLActuator.java @@ -23,10 +23,14 @@ public interface SQLActuator { SQLQueryResult execute(SQLExecutePlan plan) throws SQLException; + SQLQueryResult executeRaw(SQLExecutePlan plan) throws SQLException; + SQLQueryResult execute(SQL sql) throws SQLException; SQLQueryResult executeWithAudit(SQLExecutePlan plan) throws SQLException; + SQLQueryResult executeRawWithAudit(SQLExecutePlan plan) throws SQLException; + SQLQueryResult executeWithAudit(SQL sql) throws SQLException; SQLActuator withConnection(Connection connection); From 045d15a5027d9eeca49f1baea50b87655069dbb1 Mon Sep 17 00:00:00 2001 From: "Crane.z" <1481445951@qq.com> Date: Mon, 17 Aug 2026 14:43:09 +0800 Subject: [PATCH 2/2] fix(console): harden raw SQL execution --- .../ConsoleStatementBoundaryScanner.java | 25 ++++- .../chen/framework/console/QueryConsole.java | 35 +++--- .../framework/console/dataview/DataView.java | 2 + .../console/state/DataViewState.java | 2 + .../datasource/base/BaseSQLActuator.java | 102 +++++++++++++++++- .../datasource/sql/SQLQueryResult.java | 12 +++ 6 files changed, 156 insertions(+), 22 deletions(-) diff --git a/backend/framework/src/main/java/org/jumpserver/chen/framework/console/ConsoleStatementBoundaryScanner.java b/backend/framework/src/main/java/org/jumpserver/chen/framework/console/ConsoleStatementBoundaryScanner.java index 51d72dc..53ef5a9 100644 --- a/backend/framework/src/main/java/org/jumpserver/chen/framework/console/ConsoleStatementBoundaryScanner.java +++ b/backend/framework/src/main/java/org/jumpserver/chen/framework/console/ConsoleStatementBoundaryScanner.java @@ -1,9 +1,12 @@ package org.jumpserver.chen.framework.console; +import com.alibaba.druid.DbType; +import com.alibaba.druid.sql.SQLUtils; + import java.sql.SQLException; /** - * Finds top-level SQL statement boundaries without parsing database syntax. + * Validates that raw console input is one database command without rewriting it. */ final class ConsoleStatementBoundaryScanner { private static final String MULTIPLE_STATEMENTS_ERROR = @@ -12,11 +15,29 @@ final class ConsoleStatementBoundaryScanner { private ConsoleStatementBoundaryScanner() { } - static void requireSingleStatement(String sql) throws SQLException { + static void requireSingleStatement(String sql, DbType dbType) throws SQLException { if (sql == null || sql.isBlank()) { throw new SQLException("Console raw execution requires a database command"); } + try { + var statements = SQLUtils.parseStatements(sql, dbType); + if (statements.size() == 1) { + return; + } + if (statements.size() > 1) { + throw new SQLException(MULTIPLE_STATEMENTS_ERROR); + } + } catch (RuntimeException ignored) { + // Raw mode exists for commands the dialect parser does not understand. The conservative + // fallback below recognizes only top-level delimiters and keeps quoted/dollar-quoted + // procedure bodies opaque; it never changes the command sent to the JDBC driver. + } + + requireSingleOpaqueStatement(sql); + } + + private static void requireSingleOpaqueStatement(String sql) throws SQLException { int completedStatements = 0; boolean hasStatementContent = false; int blockCommentDepth = 0; diff --git a/backend/framework/src/main/java/org/jumpserver/chen/framework/console/QueryConsole.java b/backend/framework/src/main/java/org/jumpserver/chen/framework/console/QueryConsole.java index 03249b3..830e34f 100644 --- a/backend/framework/src/main/java/org/jumpserver/chen/framework/console/QueryConsole.java +++ b/backend/framework/src/main/java/org/jumpserver/chen/framework/console/QueryConsole.java @@ -816,23 +816,30 @@ private void runQuerySQL(String sql, Session session) throws SQLException { } private void runRawConsoleSQL(String sql, Session session) throws SQLException { - ConsoleStatementBoundaryScanner.requireSingleStatement(sql); + ConsoleStatementBoundaryScanner.requireSingleStatement(sql, this.datasource.getDruidDbType()); Connection connection = this.getConnection(); ACLResult aclResult = session.checkACL(sql, connection); if (!this.canExecuteStatement(session, sql, aclResult)) { return; } - DataView dataView = new DataView( - UUID.randomUUID().toString(), sql, this.getPacketIO(), this.getConsoleLogger() - ); - dataView.setSql(sql); - dataView.setLoadDataInterface((ignored) -> this.executeRawConsoleSQL(sql, aclResult, connection)); - dataView.loadData(); - if (!dataView.isHasTable()) { - this.logAffectedRows(dataView); - } else { - this.sendDataView(dataView, true); + SQLQueryResult executionResult = this.executeRawConsoleSQL(sql, aclResult, connection); + if (executionResult.getResults().isEmpty()) { + this.getConsoleLogger().success("%s", MessageUtils.get("ExecuteSuccess")); + } + for (SQLQueryResult result : executionResult.getResults()) { + this.getConsoleLogger().success(result); + if (!result.isHasResultSet()) { + continue; + } + + DataView dataView = new DataView( + UUID.randomUUID().toString(), sql, this.getPacketIO(), this.getConsoleLogger() + ); + dataView.setSql(sql); + dataView.setLoadDataInterface((ignored) -> result); + dataView.loadData(); + this.sendDataView(dataView, false); } } @@ -847,10 +854,8 @@ private SQLQueryResult executeRawConsoleSQL(String sql, ACLResult aclResult, Con this.getState().setCanCancel(true); this.stateManager.commit(); try { - this.getConsoleLogger().info("execute raw sql: %s", sql); - SQLQueryResult result = actuator.executeRawWithAudit(plan); - this.getConsoleLogger().success(result); - return result; + this.getConsoleLogger().info("execute sql: %s", sql); + return actuator.executeRawWithAudit(plan); } finally { if (this.currentExecution == execution) { this.currentExecution = null; diff --git a/backend/framework/src/main/java/org/jumpserver/chen/framework/console/dataview/DataView.java b/backend/framework/src/main/java/org/jumpserver/chen/framework/console/dataview/DataView.java index d91e43a..11c1d19 100644 --- a/backend/framework/src/main/java/org/jumpserver/chen/framework/console/dataview/DataView.java +++ b/backend/framework/src/main/java/org/jumpserver/chen/framework/console/dataview/DataView.java @@ -131,6 +131,8 @@ private void fullData(SQLQueryResult result) { } this.state.setPaged(result.isPaged()); + this.state.setTruncated(result.isTruncated()); + this.state.setRowLimit(result.getRowLimit()); this.data.getFields().clear(); this.data.getData().clear(); diff --git a/backend/framework/src/main/java/org/jumpserver/chen/framework/console/state/DataViewState.java b/backend/framework/src/main/java/org/jumpserver/chen/framework/console/state/DataViewState.java index edb3187..59f30b4 100644 --- a/backend/framework/src/main/java/org/jumpserver/chen/framework/console/state/DataViewState.java +++ b/backend/framework/src/main/java/org/jumpserver/chen/framework/console/state/DataViewState.java @@ -12,6 +12,8 @@ public class DataViewState extends State { private int page; private int limit; private int total; + private boolean truncated; + private int rowLimit; private boolean pinned; private boolean paged; private String filter; diff --git a/backend/framework/src/main/java/org/jumpserver/chen/framework/datasource/base/BaseSQLActuator.java b/backend/framework/src/main/java/org/jumpserver/chen/framework/datasource/base/BaseSQLActuator.java index dedc8aa..ac141fa 100644 --- a/backend/framework/src/main/java/org/jumpserver/chen/framework/datasource/base/BaseSQLActuator.java +++ b/backend/framework/src/main/java/org/jumpserver/chen/framework/datasource/base/BaseSQLActuator.java @@ -44,6 +44,7 @@ public abstract class BaseSQLActuator implements SQLActuator { private final DbType druidDbType; // Keep large JDBC text values bounded so one cell cannot fail or stall the whole result view. private static final int MAX_TEXT_DISPLAY_LENGTH = 1024 * 1024; + private static final int RAW_RESULT_ROW_LIMIT = 1000; private static final String TRUNCATED_SUFFIX = "...[truncated]"; private ConnectionManager connectionManager; private Connection connection; @@ -138,7 +139,18 @@ public SQLQueryResult execute(SQLExecutePlan plan) throws SQLException { @Override public SQLQueryResult executeRaw(SQLExecutePlan plan) throws SQLException { - return this.execute(plan, false); + SQLQueryResult executionResult = new SQLQueryResult(plan.getTargetSQL()); + executionResult.setAclResult(plan.getAclResult()); + executionResult.setHasResultSet(false); + try { + Statement statement = plan.createStatement(); + this.executeRawStatement(plan, statement, executionResult); + } finally { + if (plan.getConnection() instanceof DruidPooledConnection) { + plan.getConnection().close(); + } + } + return executionResult; } private SQLQueryResult execute(SQLExecutePlan plan, boolean enrichResult) throws SQLException { @@ -172,7 +184,7 @@ private void executeStatement( if (hasResult) { try (ResultSet resultSet = statement.getResultSet()) { - this.readResultSet(resultSet, result); + this.readResultSet(resultSet, result, Integer.MAX_VALUE); } result.setFetchFinishedTime(new Time(System.currentTimeMillis())); if (enrichResult) { @@ -195,18 +207,94 @@ private void executeStatement( result.setUpdateCount(statement.getUpdateCount()); } result.setEndTime(new Time(System.currentTimeMillis())); + } catch (SQLException e) { + throw e; } catch (Exception e) { - throw new SQLException(e.getMessage()); + throw new SQLException(e.getMessage(), e); + } + } + + private void executeRawStatement( + SQLExecutePlan plan, + Statement statement, + SQLQueryResult executionResult + ) throws SQLException { + long executionStartedAt = System.currentTimeMillis(); + executionResult.setStartTime(new Time(executionStartedAt)); + + try (statement) { + try { + // Ask the driver for one sentinel row so truncation can be detected without rewriting SQL. + statement.setMaxRows(RAW_RESULT_ROW_LIMIT + 1); + } catch (SQLException e) { + log.debug("JDBC driver does not support Statement.setMaxRows", e); + } + + long resultStartedAt = executionStartedAt; + boolean hasResultSet = statement.execute(plan.getTargetSQL()); + while (true) { + long queryFinishedAt = System.currentTimeMillis(); + if (hasResultSet) { + SQLQueryResult result = newRawResult(plan, resultStartedAt, queryFinishedAt, true); + try (ResultSet resultSet = statement.getResultSet()) { + result.setTruncated(this.readResultSet(resultSet, result, RAW_RESULT_ROW_LIMIT)); + } + result.setRowLimit(RAW_RESULT_ROW_LIMIT); + result.setTotal(result.getData().size()); + result.setFetchFinishedTime(new Time(System.currentTimeMillis())); + result.setEndTime(result.getFetchFinishedTime()); + this.handleDataMasking(result); + executionResult.getResults().add(result); + } else { + int updateCount = statement.getUpdateCount(); + if (updateCount == -1) { + break; + } + SQLQueryResult result = newRawResult(plan, resultStartedAt, queryFinishedAt, false); + result.setUpdateCount(updateCount); + result.setEndTime(new Time(System.currentTimeMillis())); + executionResult.getResults().add(result); + } + + resultStartedAt = System.currentTimeMillis(); + hasResultSet = statement.getMoreResults(Statement.CLOSE_CURRENT_RESULT); + } + + Time finishedAt = new Time(System.currentTimeMillis()); + executionResult.setQueryFinishedTime(finishedAt); + executionResult.setFetchFinishedTime(finishedAt); + executionResult.setEndTime(finishedAt); + } catch (SQLException e) { + throw e; + } catch (Exception e) { + throw new SQLException(e.getMessage(), e); } } - private void readResultSet(ResultSet resultSet, SQLQueryResult result) throws SQLException { + private SQLQueryResult newRawResult( + SQLExecutePlan plan, + long startedAt, + long queryFinishedAt, + boolean hasResultSet + ) { + SQLQueryResult result = new SQLQueryResult(plan.getTargetSQL()); + result.setAclResult(plan.getAclResult()); + result.setHasResultSet(hasResultSet); + result.setStartTime(new Time(startedAt)); + result.setQueryFinishedTime(new Time(queryFinishedAt)); + return result; + } + + private boolean readResultSet(ResultSet resultSet, SQLQueryResult result, int rowLimit) throws SQLException { var metaData = resultSet.getMetaData(); var columnCount = metaData.getColumnCount(); for (int index = 1; index <= columnCount; index++) { result.getFields().add(buildField(metaData, index)); } while (resultSet.next()) { + if (result.getData().size() >= rowLimit) { + return true; + } List row = new ArrayList<>(); for (int index = 1; index <= columnCount; index++) { try { @@ -217,6 +305,7 @@ private void readResultSet(ResultSet resultSet, SQLQueryResult result) throws SQ } result.getData().add(row); } + return false; } private void analyzeResultEditability(SQLExecutePlan plan, SQLQueryResult result) { @@ -545,10 +634,13 @@ private void handleDataMasking(SQLQueryResult result) { } private boolean matchField(Field field, String pattern) { - List names = List.of(field.getColumnName(), field.getLabel()); + String[] names = {field.getColumnName(), field.getLabel()}; String[] ps = pattern.split(","); for (String name : names) { + if (name == null) { + continue; + } for (String p : ps) { p = p.trim(); if (p.isEmpty()) continue; diff --git a/backend/framework/src/main/java/org/jumpserver/chen/framework/datasource/sql/SQLQueryResult.java b/backend/framework/src/main/java/org/jumpserver/chen/framework/datasource/sql/SQLQueryResult.java index 0177071..8193806 100644 --- a/backend/framework/src/main/java/org/jumpserver/chen/framework/datasource/sql/SQLQueryResult.java +++ b/backend/framework/src/main/java/org/jumpserver/chen/framework/datasource/sql/SQLQueryResult.java @@ -8,6 +8,7 @@ import java.sql.Time; import java.util.ArrayList; import java.util.List; +import java.util.StringJoiner; @Data @@ -15,10 +16,13 @@ public class SQLQueryResult { private String sql; private int total = -1; private boolean paged; + private boolean truncated; + private int rowLimit; private int updateCount; private boolean hasResultSet = true; private List fields = new ArrayList<>(); private List> data = new ArrayList<>(); + private List results = new ArrayList<>(); private Time startTime; private Time endTime; @@ -48,6 +52,14 @@ public SQLQueryResult(String sql) { } public String getOutput() { + if (!this.results.isEmpty()) { + StringJoiner output = new StringJoiner(System.lineSeparator()); + for (int index = 0; index < this.results.size(); index++) { + output.add("Result " + (index + 1) + ":"); + output.add(this.results.get(index).getOutput()); + } + return output.toString(); + } if (!this.hasResultSet) { return String.format("Query OK, %d rows affected", this.updateCount); }