From a537565875805e77599aca5bf0d3001b7fd1fafd Mon Sep 17 00:00:00 2001 From: zipg Date: Mon, 3 Aug 2026 15:07:18 +0800 Subject: [PATCH] fix(dameng): handle DBMS_OUTPUT messages safely --- .../main/java/com/dbx/agent/JdbcExecutor.java | 106 +++++++++++-- .../java/com/dbx/agent/JdbcExecutorTest.java | 83 ++++++++++ .../com/dbx/agent/dameng/DamengAgent.java | 106 ++++++++++++- .../com/dbx/agent/dameng/DamengAgentTest.java | 148 ++++++++++++++++++ 4 files changed, 430 insertions(+), 13 deletions(-) diff --git a/agents/common/src/main/java/com/dbx/agent/JdbcExecutor.java b/agents/common/src/main/java/com/dbx/agent/JdbcExecutor.java index acd69e65f..d135da8a9 100644 --- a/agents/common/src/main/java/com/dbx/agent/JdbcExecutor.java +++ b/agents/common/src/main/java/com/dbx/agent/JdbcExecutor.java @@ -87,6 +87,21 @@ public final class JdbcExecutor { Integer fetchSize, int timeoutSecs, ResultValueReader valueReader + ) { + return execute(conn, sql, schema, setSchemaSql, resetSchemaSql, maxRows, fetchSize, timeoutSecs, valueReader, StatementMessageReader.NONE); + } + + public QueryResult execute( + Connection conn, + String sql, + String schema, + Function setSchemaSql, + Supplier resetSchemaSql, + int maxRows, + Integer fetchSize, + int timeoutSecs, + ResultValueReader valueReader, + StatementMessageReader statementMessageReader ) { return unchecked(() -> { String trimmedSql = trimSql(sql); @@ -122,7 +137,7 @@ public final class JdbcExecutor { false ); } - return withStatementWarnings(result, stmt); + return withStatementMessages(result, stmt, effectiveMaxRows, statementMessageReader); } finally { activeStatements.remove(stmt); } @@ -198,7 +213,19 @@ public final class JdbcExecutor { QueryPageOptions options, ResultValueReader valueReader ) { - return executePage(conn, sql, schema, setSchemaSql, () -> "", options, valueReader, sessions); + return executePage(conn, sql, schema, setSchemaSql, () -> "", options, valueReader, StatementMessageReader.NONE, sessions); + } + + public QueryPageResult executePage( + Connection conn, + String sql, + String schema, + Function setSchemaSql, + QueryPageOptions options, + ResultValueReader valueReader, + StatementMessageReader statementMessageReader + ) { + return executePage(conn, sql, schema, setSchemaSql, () -> "", options, valueReader, statementMessageReader, sessions); } public QueryPageResult executePage( @@ -210,7 +237,7 @@ public final class JdbcExecutor { QueryPageOptions options, ResultValueReader valueReader ) { - return executePage(conn, sql, schema, setSchemaSql, resetSchemaSql, options, valueReader, sessions); + return executePage(conn, sql, schema, setSchemaSql, resetSchemaSql, options, valueReader, StatementMessageReader.NONE, sessions); } public QueryPageResult startTableRead( @@ -221,7 +248,7 @@ public final class JdbcExecutor { QueryPageOptions options, ResultValueReader valueReader ) { - return executePage(conn, sql, schema, setSchemaSql, () -> "", options, valueReader, tableReadSessions); + return executePage(conn, sql, schema, setSchemaSql, () -> "", options, valueReader, StatementMessageReader.NONE, tableReadSessions); } public QueryPageResult startTableRead( @@ -233,7 +260,7 @@ public final class JdbcExecutor { QueryPageOptions options, ResultValueReader valueReader ) { - return executePage(conn, sql, schema, setSchemaSql, resetSchemaSql, options, valueReader, tableReadSessions); + return executePage(conn, sql, schema, setSchemaSql, resetSchemaSql, options, valueReader, StatementMessageReader.NONE, tableReadSessions); } private QueryPageResult executePage( @@ -244,6 +271,7 @@ public final class JdbcExecutor { Supplier resetSchemaSql, QueryPageOptions options, ResultValueReader valueReader, + StatementMessageReader statementMessageReader, ConcurrentHashMap targetSessions ) { return unchecked(() -> { @@ -267,13 +295,32 @@ public final class JdbcExecutor { long elapsed = System.currentTimeMillis() - start; if (!hasResultSet) { int updateCount = stmt.getUpdateCount(); - activeStatements.remove(stmt); - stmt.close(); - return new QueryPageResult( + QueryResult result = new QueryResult( Collections.emptyList(), Collections.emptyList(), updateCount >= 0 ? updateCount : 0, - elapsed + elapsed, + false + ); + if (statementMessageReader != StatementMessageReader.NONE) { + result = withStatementMessages( + result, + stmt, + Math.max(options.getMaxRows(), 1), + statementMessageReader + ); + } + activeStatements.remove(stmt); + stmt.close(); + return new QueryPageResult( + result.getColumns(), + result.getColumn_types(), + result.getRows(), + result.getAffected_rows(), + result.getExecution_time_ms(), + result.getTruncated(), + null, + false ); } @@ -700,17 +747,28 @@ public final class JdbcExecutor { } } - private static QueryResult withStatementWarnings(QueryResult result, Statement stmt) { + private static QueryResult withStatementMessages( + QueryResult result, + Statement stmt, + int maxRows, + StatementMessageReader statementMessageReader + ) { if (!result.getColumns().isEmpty() || !result.getRows().isEmpty()) { return result; } List> rows = new ArrayList<>(); + int effectiveMaxRows = Math.max(maxRows, 1); + boolean truncated = result.getTruncated(); try { Set seen = Collections.newSetFromMap(new IdentityHashMap<>()); for (SQLWarning warning = stmt.getWarnings(); warning != null && seen.add(warning); warning = warning.getNextWarning()) { String message = warning.getMessage(); if (message != null && !message.trim().isEmpty()) { + if (rows.size() >= effectiveMaxRows) { + truncated = true; + break; + } rows.add(Collections.singletonList(message)); } } @@ -720,6 +778,24 @@ public final class JdbcExecutor { // successfully executed statement into a query failure. } + try { + List messages = statementMessageReader.read(stmt); + if (messages != null) { + for (String message : messages) { + if (message == null) { + continue; + } + if (rows.size() >= effectiveMaxRows) { + truncated = true; + break; + } + rows.add(Collections.singletonList(message)); + } + } + } catch (Exception ignored) { + // Driver-specific informational output is advisory, like SQLWarning. + } + if (rows.isEmpty()) { return result; } @@ -729,7 +805,7 @@ public final class JdbcExecutor { rows, result.getAffected_rows(), result.getExecution_time_ms(), - result.getTruncated() + truncated ); } @@ -800,6 +876,14 @@ public final class JdbcExecutor { Object read(ResultSet rs, int index, int sqlType) throws SQLException; } + /** Reads driver-specific informational output that is not exposed as {@link SQLWarning}. */ + @FunctionalInterface + public interface StatementMessageReader { + StatementMessageReader NONE = statement -> Collections.emptyList(); + + List read(Statement statement) throws SQLException; + } + /** * Optional extension of {@link ResultValueReader} that exposes the JDBC * {@code getColumnTypeName} alongside the SQL type code, allowing per-driver diff --git a/agents/common/src/test/java/com/dbx/agent/JdbcExecutorTest.java b/agents/common/src/test/java/com/dbx/agent/JdbcExecutorTest.java index ca659389c..9244b6af4 100644 --- a/agents/common/src/test/java/com/dbx/agent/JdbcExecutorTest.java +++ b/agents/common/src/test/java/com/dbx/agent/JdbcExecutorTest.java @@ -20,8 +20,10 @@ import java.util.concurrent.atomic.AtomicInteger; import javax.sql.rowset.serial.SerialBlob; import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertFalse; import static org.junit.jupiter.api.Assertions.assertNull; import static org.junit.jupiter.api.Assertions.assertThrows; +import static org.junit.jupiter.api.Assertions.assertTrue; class JdbcExecutorTest { @Test @@ -145,6 +147,87 @@ class JdbcExecutorTest { assertEquals(3L, result.getAffected_rows()); } + @Test + void executeReturnsDriverMessagesForNoResultStatementsAndHonorsMaxRows() { + QueryResult result = JdbcExecutor.INSTANCE.execute( + executionConnection(false, -1, null, new AtomicInteger(), null, null), + "CALL LOG_ONLY_PROCEDURE()", + "", + schema -> "", + () -> "", + 2, + null, + 0, + JdbcExecutor.INSTANCE::defaultResultValue, + statement -> Arrays.asList("first", "second", "third") + ); + + assertEquals(Arrays.asList("Message"), result.getColumns()); + assertEquals(Arrays.asList(Arrays.asList("first"), Arrays.asList("second")), result.getRows()); + assertTrue(result.getTruncated()); + } + + @Test + void executeLimitsCombinedWarningsAndDriverMessages() { + SQLWarning first = new SQLWarning("first warning"); + first.setNextWarning(new SQLWarning("second warning")); + + QueryResult result = JdbcExecutor.INSTANCE.execute( + executionConnection(false, -1, first, new AtomicInteger(), null, null), + "CALL LOG_ONLY_PROCEDURE()", + "", + schema -> "", + () -> "", + 1, + null, + 0, + JdbcExecutor.INSTANCE::defaultResultValue, + statement -> Arrays.asList("driver message") + ); + + assertEquals(Arrays.asList("Message"), result.getColumns()); + assertEquals(Arrays.asList(Arrays.asList("first warning")), result.getRows()); + assertTrue(result.getTruncated()); + } + + @Test + void executePageReturnsDriverMessagesForNoResultStatements() { + QueryPageResult result = JdbcExecutor.INSTANCE.executePage( + executionConnection(false, -1, null, new AtomicInteger(), null, null), + "CALL LOG_ONLY_PROCEDURE()", + "", + schema -> "", + new QueryPageOptions(100, null, 100), + JdbcExecutor.INSTANCE::defaultResultValue, + statement -> Arrays.asList("first", "second") + ); + + assertEquals(Arrays.asList("Message"), result.getColumns()); + assertEquals(Arrays.asList(Arrays.asList("first"), Arrays.asList("second")), result.getRows()); + assertFalse(result.getHas_more()); + } + + @Test + void executePageKeepsWarningsHiddenWithoutADriverMessageReader() { + QueryPageResult result = JdbcExecutor.INSTANCE.executePage( + executionConnection( + false, + -1, + new SQLWarning("existing paged warning"), + new AtomicInteger(), + null, + null + ), + "CALL EXISTING_PROCEDURE()", + "", + schema -> "", + new QueryPageOptions() + ); + + assertEquals(Collections.emptyList(), result.getColumns()); + assertEquals(Collections.emptyList(), result.getRows()); + } + @Test void executeDoesNotReplaceOrdinaryResultSetsWithWarnings() { CountingResultSetFixture fixture = countingResultSet(new Object[][]{{1, "Ada"}}); diff --git a/agents/drivers/dameng/src/main/java/com/dbx/agent/dameng/DamengAgent.java b/agents/drivers/dameng/src/main/java/com/dbx/agent/dameng/DamengAgent.java index c38a4c114..e7eba9140 100644 --- a/agents/drivers/dameng/src/main/java/com/dbx/agent/dameng/DamengAgent.java +++ b/agents/drivers/dameng/src/main/java/com/dbx/agent/dameng/DamengAgent.java @@ -28,6 +28,11 @@ import java.sql.DriverManager; import java.sql.PreparedStatement; import java.sql.ResultSet; import java.sql.SQLException; +import java.sql.SQLFeatureNotSupportedException; +import java.sql.SQLNonTransientConnectionException; +import java.sql.SQLRecoverableException; +import java.sql.SQLSyntaxErrorException; +import java.sql.SQLTransientConnectionException; import java.sql.SQLXML; import java.sql.Statement; import java.sql.Types; @@ -102,6 +107,77 @@ public final class DamengAgent extends AbstractJdbcAgent { connectedUsername = params.getUsername(); } + @Override + protected void afterPhysicalConnect(ConnectParams params, Connection connection) throws SQLException { + try (Statement statement = connection.createStatement()) { + statement.execute("BEGIN DBMS_OUTPUT.ENABLE(1000000); END;"); + } catch (SQLException error) { + if (!isIgnorableDbmsOutputError(error)) { + throw error; + } + } + } + + private static boolean isIgnorableDbmsOutputError(SQLException error) { + for (Throwable current = error; current != null; current = current.getCause()) { + if (current instanceof SQLException sqlError) { + for (SQLException candidate = sqlError; candidate != null; candidate = candidate.getNextException()) { + if (isConnectionError(candidate)) { + return false; + } + } + } + } + for (Throwable current = error; current != null; current = current.getCause()) { + if (current instanceof SQLException sqlError) { + for (SQLException candidate = sqlError; candidate != null; candidate = candidate.getNextException()) { + if (candidate instanceof SQLFeatureNotSupportedException || candidate instanceof SQLSyntaxErrorException) { + return true; + } + String sqlState = candidate.getSQLState(); + if ("0A000".equalsIgnoreCase(sqlState) + || "42000".equalsIgnoreCase(sqlState) + || "42501".equalsIgnoreCase(sqlState)) { + return true; + } + String message = candidate.getMessage(); + if (message != null && isDbmsOutputUnavailableMessage(message.toLowerCase(Locale.ROOT))) { + return true; + } + } + } + } + return false; + } + + private static boolean isConnectionError(SQLException error) { + String sqlState = error.getSQLState(); + return error instanceof SQLNonTransientConnectionException + || error instanceof SQLRecoverableException + || error instanceof SQLTransientConnectionException + || (sqlState != null && sqlState.toUpperCase(Locale.ROOT).startsWith("08")); + } + + private static boolean isDbmsOutputUnavailableMessage(String message) { + if (!message.contains("dbms_output")) { + return false; + } + return message.contains("权限") + || message.contains("privilege") + || message.contains("permission") + || message.contains("access denied") + || message.contains("not authorized") + || message.contains("不支持") + || message.contains("unsupported") + || message.contains("not supported") + || message.contains("不存在") + || message.contains("not exist") + || message.contains("not found") + || message.contains("未找到") + || message.contains("undefined") + || message.contains("未定义"); + } + /** * The DM JDBC driver writes a banner to {@code System.out} during * {@code Class.forName} / driver initialization. This corrupts the @@ -1040,13 +1116,38 @@ public final class DamengAgent extends AbstractJdbcAgent { sql, schema, this::setSchemaSQL, + () -> "", options.getMaxRows(), options.getFetchSize(), options.getTimeoutSecs(), - this::resultValue + this::resultValue, + DamengAgent::statementPrintMessages ); } + static List statementPrintMessages(Statement statement) { + try { + Object target = statement; + Method method; + try { + method = statement.getClass().getMethod("getPrintMsg"); + } catch (NoSuchMethodException ignored) { + // Pooled connections expose a Hikari proxy rather than DmdbStatement directly. + Class damengStatementClass = Class.forName("dm.jdbc.driver.DmdbStatement"); + target = statement.unwrap(damengStatementClass); + method = damengStatementClass.getMethod("getPrintMsg"); + } + Object value = method.invoke(target); + if (!(value instanceof String)) { + return List.of(); + } + String message = (String) value; + return message.isEmpty() ? List.of() : message.lines().toList(); + } catch (Exception ignored) { + return List.of(); + } + } + private QueryResult executeExplainQuery(String sql, String schema, ExecuteQueryOptions options) { return explainQueryResult(sql, schema, options.getTimeoutSecs(), options.getMaxRows()); } @@ -1139,7 +1240,8 @@ public final class DamengAgent extends AbstractJdbcAgent { schema, this::setSchemaSQL, options, - this::resultValue + this::resultValue, + DamengAgent::statementPrintMessages ); } diff --git a/agents/drivers/dameng/src/test/java/com/dbx/agent/dameng/DamengAgentTest.java b/agents/drivers/dameng/src/test/java/com/dbx/agent/dameng/DamengAgentTest.java index fb80de5f4..13ff1c286 100644 --- a/agents/drivers/dameng/src/test/java/com/dbx/agent/dameng/DamengAgentTest.java +++ b/agents/drivers/dameng/src/test/java/com/dbx/agent/dameng/DamengAgentTest.java @@ -12,11 +12,22 @@ import com.dbx.agent.test.JdbcAgentFake; import com.dbx.agent.test.TestSupport; import org.junit.jupiter.api.Test; +import java.lang.reflect.InvocationHandler; +import java.lang.reflect.Method; +import java.lang.reflect.Proxy; +import java.sql.Connection; +import java.sql.SQLException; +import java.sql.SQLFeatureNotSupportedException; +import java.sql.SQLTransientConnectionException; +import java.sql.Statement; +import java.util.ArrayList; import java.util.List; +import static org.junit.jupiter.api.Assertions.assertDoesNotThrow; import static org.junit.jupiter.api.Assertions.assertEquals; import static org.junit.jupiter.api.Assertions.assertFalse; import static org.junit.jupiter.api.Assertions.assertNull; +import static org.junit.jupiter.api.Assertions.assertSame; import static org.junit.jupiter.api.Assertions.assertThrows; import static org.junit.jupiter.api.Assertions.assertTrue; @@ -47,6 +58,80 @@ class DamengAgentTest extends JdbcFakeExecutionBehaviorTest { assertEquals(List.of("executeQuery"), JdbcAgentFake.calls); } + @Test + void physicalConnectionsEnableDbmsOutputWithoutChangingUserSql() throws Exception { + List executedSql = new ArrayList<>(); + DamengAgent agent = new DamengAgent(); + + agent.afterPhysicalConnect(null, printMessageConnection(null, executedSql)); + + assertEquals(List.of("BEGIN DBMS_OUTPUT.ENABLE(1000000); END;"), executedSql); + } + + @Test + void physicalConnectionsIgnoreUnsupportedOrRestrictedDbmsOutput() { + DamengAgent agent = new DamengAgent(); + + assertDoesNotThrow(() -> agent.afterPhysicalConnect( + null, + failingDbmsOutputConnection(new SQLFeatureNotSupportedException("unsupported", "0A000")) + )); + assertDoesNotThrow(() -> agent.afterPhysicalConnect( + null, + failingDbmsOutputConnection(new SQLException("permission denied", "42000")) + )); + } + + @Test + void physicalConnectionsPropagateConnectionFailures() { + DamengAgent agent = new DamengAgent(); + SQLException transientFailure = new SQLTransientConnectionException("connection closed"); + SQLException sqlStateFailure = new SQLException("connection failure", "08006"); + SQLException wrappedFailure = new SQLException("permission denied", "42000"); + wrappedFailure.initCause(new SQLTransientConnectionException("connection closed")); + + assertSame(transientFailure, assertThrows( + SQLException.class, + () -> agent.afterPhysicalConnect(null, failingDbmsOutputConnection(transientFailure)) + )); + assertSame(sqlStateFailure, assertThrows( + SQLException.class, + () -> agent.afterPhysicalConnect(null, failingDbmsOutputConnection(sqlStateFailure)) + )); + assertSame(wrappedFailure, assertThrows( + SQLException.class, + () -> agent.afterPhysicalConnect(null, failingDbmsOutputConnection(wrappedFailure)) + )); + } + + @Test + void physicalConnectionsPropagateUnrelatedSetupFailures() { + DamengAgent agent = new DamengAgent(); + SQLException failure = new SQLException("resource busy", "HY000"); + + assertSame(failure, assertThrows( + SQLException.class, + () -> agent.afterPhysicalConnect(null, failingDbmsOutputConnection(failure)) + )); + } + + @Test + void executeQueryReturnsDamengPrintMessagesForLogOnlyProcedures() { + List executedSql = new ArrayList<>(); + DamengAgent agent = new DamengAgent(); + TestSupport.setPrivateConnection(agent, printMessageConnection("first\n中文日志\n", executedSql)); + + QueryResult result = agent.executeQuery( + "CALL LOG_ONLY_PROCEDURE('input')", + null, + new ExecuteQueryOptions() + ); + + assertEquals(List.of("Message"), result.getColumns()); + assertEquals(List.of(List.of("first"), List.of("中文日志")), result.getRows()); + assertEquals(List.of("CALL LOG_ONLY_PROCEDURE('input')"), executedSql); + } + @Test void executeQueryPageReturnsPlanRowsForExplainStatements() { DamengAgent agent = new DamengAgent(); @@ -274,4 +359,67 @@ class DamengAgentTest extends JdbcFakeExecutionBehaviorTest { assertTrue(query.sql().endsWith("LIMIT ? OFFSET ?")); assertEquals(List.of("REPORTING", "VIEW", "MATERIALIZED_VIEW", "%S%A%L%E%S%", 10, 30), query.args()); } + + private static Connection printMessageConnection(String printMessage, List executedSql) { + return statementConnection(printMessage, executedSql, null); + } + + private static Connection failingDbmsOutputConnection(SQLException failure) { + return statementConnection(null, new ArrayList<>(), failure); + } + + private static Connection statementConnection( + String printMessage, + List executedSql, + SQLException executeFailure + ) { + InvocationHandler statementHandler = (Object unused, Method method, Object[] args) -> { + switch (method.getName()) { + case "execute": + if (executeFailure != null) { + throw executeFailure; + } + executedSql.add((String) args[0]); + return false; + case "getPrintMsg": + return printMessage; + case "getUpdateCount": + return -1; + default: + return defaultValue(method.getReturnType()); + } + }; + Statement statement = (Statement) Proxy.newProxyInstance( + DamengAgentTest.class.getClassLoader(), + new Class[]{Statement.class, PrintMessageStatement.class}, + statementHandler + ); + InvocationHandler connectionHandler = (Object unused, Method method, Object[] args) -> { + if (method.getName().equals("createStatement")) { + return statement; + } + return defaultValue(method.getReturnType()); + }; + return (Connection) Proxy.newProxyInstance( + DamengAgentTest.class.getClassLoader(), + new Class[]{Connection.class}, + connectionHandler + ); + } + + private static Object defaultValue(Class type) { + if (type == Boolean.TYPE) return false; + if (type == Byte.TYPE) return (byte) 0; + if (type == Short.TYPE) return (short) 0; + if (type == Integer.TYPE) return 0; + if (type == Long.TYPE) return 0L; + if (type == Float.TYPE) return 0f; + if (type == Double.TYPE) return 0.0d; + if (type == Character.TYPE) return '\0'; + return null; + } + + public interface PrintMessageStatement { + String getPrintMsg(); + } }