diff --git a/agents/common/src/main/java/com/dbx/agent/AbstractJdbcAgent.java b/agents/common/src/main/java/com/dbx/agent/AbstractJdbcAgent.java index e6b64132a9..ccdb74b91e 100644 --- a/agents/common/src/main/java/com/dbx/agent/AbstractJdbcAgent.java +++ b/agents/common/src/main/java/com/dbx/agent/AbstractJdbcAgent.java @@ -202,7 +202,7 @@ public abstract class AbstractJdbcAgent extends BaseDatabaseAgent { options.getFetchSize(), options.getTimeoutSecs(), resultValueReader(), - JdbcExecutor.StatementMessageReader.NONE, + statementMessageReader(), advancePastUpdateCounts() ); } @@ -221,6 +221,7 @@ public abstract class AbstractJdbcAgent extends BaseDatabaseAgent { this::resetSchemaSQL, options, resultValueReader(), + statementMessageReader(), advancePastUpdateCounts() ); } @@ -602,6 +603,10 @@ public abstract class AbstractJdbcAgent extends BaseDatabaseAgent { protected void beforeQueryExecution(Connection connection, int timeoutSecs) throws Exception { } + protected JdbcExecutor.StatementMessageReader statementMessageReader() { + return JdbcExecutor.StatementMessageReader.NONE; + } + /** Called for each session RPC before its SQL can replace the connection's last trace ID. */ protected void beforeAgentMethod(String method, String querySessionId) { } 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 dfb75839bd..31345d7231 100644 --- a/agents/common/src/main/java/com/dbx/agent/JdbcExecutor.java +++ b/agents/common/src/main/java/com/dbx/agent/JdbcExecutor.java @@ -15,6 +15,7 @@ import java.util.Collections; import java.util.IdentityHashMap; import java.util.List; import java.util.Locale; +import java.util.Map; import java.util.Set; import java.util.UUID; import java.util.concurrent.ConcurrentHashMap; @@ -300,6 +301,7 @@ public final class JdbcExecutor { try (ResultSet rs = stmt.getResultSet()) { result = readResultSet(rs, elapsed, effectiveMaxRows, valueReader); } + appendStatementReaderMessages(result, stmt, statementMessageReader); // `elapsed` stops at `stmt.execute()`. Retrieving and // converting the rows above is frequently the dominant // cost of a statement, so report the whole duration the @@ -316,7 +318,9 @@ public final class JdbcExecutor { false ); } - return withStatementMessages(result, stmt, effectiveMaxRows, statementMessageReader); + QueryResult completed = withStatementMessages(result, stmt, effectiveMaxRows, statementMessageReader); + completed.setExecution_time_ms(System.currentTimeMillis() - start); + return completed; } finally { activeStatements.remove(stmt); } @@ -411,6 +415,21 @@ public final class JdbcExecutor { return executePage(conn, sql, schema, setSchemaSql, () -> "", options, valueReader, statementMessageReader, sessions); } + public QueryPageResult executePage( + Connection conn, + String sql, + String schema, + Function setSchemaSql, + Supplier resetSchemaSql, + QueryPageOptions options, + ResultValueReader valueReader, + StatementMessageReader statementMessageReader, + boolean advancePastUpdateCounts + ) { + return executePage(conn, sql, schema, setSchemaSql, resetSchemaSql, options, valueReader, + statementMessageReader, sessions, advancePastUpdateCounts); + } + public QueryPageResult executePage( Connection conn, String sql, @@ -471,6 +490,20 @@ public final class JdbcExecutor { StatementMessageReader.NONE, sessions, false, true); } + public QueryPageResult executeBoundedPage( + Connection conn, + String sql, + String schema, + Function setSchemaSql, + Supplier resetSchemaSql, + QueryPageOptions options, + ResultValueReader valueReader, + StatementMessageReader statementMessageReader + ) { + return executePage(conn, sql, schema, setSchemaSql, resetSchemaSql, options, valueReader, + statementMessageReader, sessions, false, true); + } + public QueryPageResult startTableRead( Connection conn, String sql, @@ -600,6 +633,7 @@ public final class JdbcExecutor { } else { appendStatementWarnings(result, stmt); } + result.setExecution_time_ms(System.currentTimeMillis() - start); activeStatements.remove(stmt); stmt.close(); QueryPageResult page = new QueryPageResult( @@ -642,7 +676,8 @@ public final class JdbcExecutor { sqlTypeByIndex, typeNameByIndex, Math.max(options.getMaxRows(), 1), - valueReader + valueReader, + statementMessageReader ); createdSession = session; targetSessions.put(sessionId, session); @@ -939,7 +974,7 @@ public final class JdbcExecutor { while (rows.size() < effectivePageSize && session.rowsRead < session.maxRows) { if (!session.resultSet.next()) { QueryPageResult page = sessionPageResult(session, rows, executionTimeMs, false, null, false); - closeSession(targetSessions, session.id); + appendStatementReaderMessages(page, closeSessionMessages(targetSessions, session.id)); return page; } rows.add(rowValues(session.resultSet, session.valueReader, session.sqlTypeByIndex, session.typeNameByIndex)); @@ -949,14 +984,14 @@ public final class JdbcExecutor { if (session.rowsRead >= session.maxRows) { boolean truncated = session.resultSet.next(); QueryPageResult page = sessionPageResult(session, rows, executionTimeMs, truncated, null, false); - closeSession(targetSessions, session.id); + appendStatementReaderMessages(page, closeSessionMessages(targetSessions, session.id)); return page; } boolean hasMore = session.resultSet.next(); if (!hasMore) { QueryPageResult page = sessionPageResult(session, rows, executionTimeMs, false, null, false); - closeSession(targetSessions, session.id); + appendStatementReaderMessages(page, closeSessionMessages(targetSessions, session.id)); return page; } @@ -990,9 +1025,13 @@ public final class JdbcExecutor { } private boolean closeSession(ConcurrentHashMap targetSessions, String sessionId) { + return closeSessionMessages(targetSessions, sessionId) != null; + } + + private List closeSessionMessages(ConcurrentHashMap targetSessions, String sessionId) { QuerySession session = targetSessions.remove(sessionId); if (session == null) { - return false; + return null; } synchronized (session) { activeStatements.remove(session.statement); @@ -1000,12 +1039,17 @@ public final class JdbcExecutor { session.resultSet.close(); } catch (Exception ignored) { } + List messages = Collections.emptyList(); + try { + messages = session.statementMessageReader.close(session.statement); + } catch (Exception ignored) { + } try { session.statement.close(); } catch (Exception ignored) { } + return messages == null ? Collections.emptyList() : messages; } - return true; } static String safeColumnTypeName(ResultSetMetaData meta, int columnIndex) { @@ -1144,6 +1188,32 @@ public final class JdbcExecutor { return messageResult; } + private static void appendStatementReaderMessages( + QueryResult result, + Statement statement, + StatementMessageReader reader + ) { + try { + List messages = reader.close(statement); + if (messages != null) { + for (String message : messages) { + if (message != null) result.addInformationalMessage(message, null); + } + } + } catch (Exception ignored) { + // Driver-specific informational output is advisory. + } + } + + private static void appendStatementReaderMessages(QueryPageResult page, List messages) { + if (messages == null || messages.isEmpty()) return; + List> result = new ArrayList<>(page.getMessages()); + for (String message : messages) { + if (message != null) result.add(Map.of("severity", "INFO", "message", message)); + } + page.setMessages(result); + } + private QueryResult emptyQueryResult(long start) { return new QueryResult( Collections.emptyList(), @@ -1217,6 +1287,10 @@ public final class JdbcExecutor { StatementMessageReader NONE = statement -> Collections.emptyList(); List read(Statement statement) throws SQLException; + + default List close(Statement statement) throws SQLException { + return Collections.emptyList(); + } } /** @@ -1262,6 +1336,7 @@ public final class JdbcExecutor { private final String[] typeNameByIndex; private final int maxRows; private final ResultValueReader valueReader; + private final StatementMessageReader statementMessageReader; private int rowsRead; private List pendingRow; private long lastAccessedAtMillis; @@ -1275,7 +1350,8 @@ public final class JdbcExecutor { int[] sqlTypeByIndex, String[] typeNameByIndex, int maxRows, - ResultValueReader valueReader + ResultValueReader valueReader, + StatementMessageReader statementMessageReader ) { this.id = id; this.statement = statement; @@ -1286,6 +1362,7 @@ public final class JdbcExecutor { this.typeNameByIndex = typeNameByIndex; this.maxRows = maxRows; this.valueReader = valueReader; + this.statementMessageReader = statementMessageReader; this.lastAccessedAtMillis = System.currentTimeMillis(); } } 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 8ff05ac81b..79024d3d69 100644 --- a/agents/common/src/test/java/com/dbx/agent/JdbcExecutorTest.java +++ b/agents/common/src/test/java/com/dbx/agent/JdbcExecutorTest.java @@ -203,6 +203,39 @@ class JdbcExecutorTest { ); } + @Test + void executeCountsDriverMessageReadingInExecutionTime() { + JdbcExecutor.StatementMessageReader reader = statement -> { + try { + Thread.sleep(30L); + } catch (InterruptedException error) { + Thread.currentThread().interrupt(); + throw new SQLException(error); + } + return Collections.singletonList("procedure output"); + }; + + QueryResult result = JdbcExecutor.INSTANCE.execute( + executionConnection(false, -1, null, new AtomicInteger(), null, null), + "BEGIN OUTPUT_TEXT(); END;", "", schema -> "", () -> "", 100, null, 0, + JdbcExecutor.INSTANCE::defaultResultValue, reader + ); + + assertTrue(result.getExecution_time_ms() >= 25L, + "driver message reading must count toward execution_time_ms but was " + result.getExecution_time_ms() + "ms"); + for (boolean bounded : List.of(false, true)) { + Connection connection = executionConnection(false, -1, null, new AtomicInteger(), null, null); + QueryPageResult page = bounded + ? JdbcExecutor.INSTANCE.executeBoundedPage(connection, "BEGIN OUTPUT_TEXT(); END;", "", + schema -> "", () -> "", new QueryPageOptions(), JdbcExecutor.INSTANCE::defaultResultValue, reader) + : JdbcExecutor.INSTANCE.executePage(connection, "BEGIN OUTPUT_TEXT(); END;", "", + schema -> "", () -> "", new QueryPageOptions(), JdbcExecutor.INSTANCE::defaultResultValue, reader, false); + assertEquals(List.of(List.of("procedure output")), page.getRows()); + assertTrue(page.getExecution_time_ms() >= 25L, + "paged message reading must count toward execution_time_ms but was " + page.getExecution_time_ms() + "ms"); + } + } + @Test void executeReturnsMultipleStatementWarningsForNoResultStatements() { SQLWarning first = new SQLWarning("identity value is 443", "S0003", 7998); @@ -336,6 +369,42 @@ class JdbcExecutorTest { assertEquals("informational", result.getMessages().get(0).get("message")); } + @Test + void executeRetainsRowsAndAddsDriverMessagesWhenResultSetCompletes() { + CountingResultSetFixture fixture = countingResultSet(new Object[][]{{1, "Ada"}}); + boolean[] closed = {false}; + JdbcExecutor.StatementMessageReader reader = new JdbcExecutor.StatementMessageReader() { + @Override + public List read(Statement statement) { + return Collections.emptyList(); + } + + @Override + public List close(Statement statement) { + closed[0] = true; + return Arrays.asList("procedure output"); + } + }; + + QueryResult result = JdbcExecutor.INSTANCE.execute( + executionConnection(true, -1, null, new AtomicInteger(), fixture.resultSet(), null), + "BEGIN OPEN_RESULT(); END;", + "", + schema -> "", + () -> "", + 100, + null, + 0, + JdbcExecutor.INSTANCE::defaultResultValue, + reader + ); + + assertTrue(closed[0]); + assertEquals(Arrays.asList("id", "name"), result.getColumns()); + assertEquals(Arrays.asList(Arrays.asList(1, "Ada")), result.getRows()); + assertEquals("procedure output", result.getMessages().get(0).get("message")); + } + @Test void executePageRetainsWarningsWithOrdinaryRows() { CountingResultSetFixture fixture = countingResultSet(new Object[][]{{1, "Ada"}}); diff --git a/agents/drivers/oceanbase-oracle/src/main/java/com/dbx/agent/oceanbaseoracle/OceanBaseOracleAgent.java b/agents/drivers/oceanbase-oracle/src/main/java/com/dbx/agent/oceanbaseoracle/OceanBaseOracleAgent.java index d3be303534..5dd092897d 100644 --- a/agents/drivers/oceanbase-oracle/src/main/java/com/dbx/agent/oceanbaseoracle/OceanBaseOracleAgent.java +++ b/agents/drivers/oceanbase-oracle/src/main/java/com/dbx/agent/oceanbaseoracle/OceanBaseOracleAgent.java @@ -35,10 +35,13 @@ import com.dbx.agent.SpatialColumn; import com.dbx.agent.TableInfo; import com.dbx.agent.TriggerInfo; +import java.sql.CallableStatement; import java.sql.Connection; import java.sql.PreparedStatement; import java.sql.ResultSet; import java.sql.SQLException; +import java.sql.SQLTimeoutException; +import java.sql.Statement; import java.sql.Types; import java.util.ArrayDeque; import java.util.ArrayList; @@ -73,6 +76,9 @@ public final class OceanBaseOracleAgent extends ConfiguredJdbcAgent { private boolean queryTimeoutChanged; private final OceanBaseLobValues lobValues = new OceanBaseLobValues(); private boolean deferCharacterLobs; + private boolean captureDbmsOutput; + private boolean dbmsOutputEnabled; + private String dbmsOutputWarning; public static final JdbcAgentProfile OCEANBASE_ORACLE_PROFILE = new JdbcAgentProfile( "com.oceanbase.jdbc.Driver", @@ -108,8 +114,11 @@ public final class OceanBaseOracleAgent extends ConfiguredJdbcAgent { @Override public QueryResult executeQuery(String sql, String schema, ExecuteQueryOptions options) { var lobCheckpoint = lobValues.checkpoint(); + Connection outputConnection = null; try (QueryTiming timing = QueryTiming.begin()) { prepareLobExecution(sql, options.getDeferLobs()); + captureDbmsOutput = isDbmsOutputStatement(sql); + outputConnection = requireConnected(); QueryResult result = super.executeQuery(sql, schema, options); var marked = OceanBaseLobValues.mark(result.getColumns(), result.getColumn_types(), result.getRows()); if (marked != null) { @@ -124,9 +133,13 @@ public final class OceanBaseOracleAgent extends ConfiguredJdbcAgent { result.setQuery_timings_ms(timing.finish()); return result; } catch (RuntimeException error) { + if (dbmsOutputEnabled && outputConnection != null) cleanupDbmsOutputAfterFailure(outputConnection, error); lobValues.releaseSince(lobCheckpoint); throw error; } finally { + captureDbmsOutput = false; + dbmsOutputEnabled = false; + dbmsOutputWarning = null; deferCharacterLobs = false; } } @@ -134,22 +147,29 @@ public final class OceanBaseOracleAgent extends ConfiguredJdbcAgent { @Override public QueryPageResult executeQueryPage(String sql, String schema, QueryPageOptions options) { var lobCheckpoint = lobValues.checkpoint(); + Connection outputConnection = null; try (QueryTiming timing = QueryTiming.begin()) { prepareLobExecution(sql, options.getDeferLobs()); + captureDbmsOutput = isDbmsOutputStatement(sql); long prepareStarted = System.nanoTime(); - Connection connection = requireConnected(); + outputConnection = requireConnected(); + Connection connection = outputConnection; uncheckedVoid(() -> beforeQueryExecution(connection, options.getTimeoutSecs())); QueryTiming.record("session_prepare", prepareStarted); QueryPageResult result = JdbcExecutor.current().executeBoundedPage( - connection, sql, schema, this::setSchemaSQL, this::resetSchemaSQL, options, resultValueReader() + connection, sql, schema, this::setSchemaSQL, this::resetSchemaSQL, options, resultValueReader(), statementMessageReader() ); markLobPage(result); result.setQuery_timings_ms(timing.finish()); return result; } catch (RuntimeException error) { + if (dbmsOutputEnabled && outputConnection != null) cleanupDbmsOutputAfterFailure(outputConnection, error); lobValues.releaseSince(lobCheckpoint); throw error; } finally { + captureDbmsOutput = false; + dbmsOutputEnabled = false; + dbmsOutputWarning = null; deferCharacterLobs = false; } } @@ -342,13 +362,173 @@ public final class OceanBaseOracleAgent extends ConfiguredJdbcAgent { try (var stmt = connection.createStatement()) { try { stmt.execute(queryTimeoutSql(timeoutSecs)); + queryTimeoutChanged = true; } catch (SQLException error) { - if (isReadOnlyTransactionError(error)) { - return; + if (!isReadOnlyTransactionError(error)) throw error; + } + } + dbmsOutputEnabled = captureDbmsOutput; + dbmsOutputWarning = captureDbmsOutput ? enableDbmsOutput(connection) : null; + } + + @Override + protected JdbcExecutor.StatementMessageReader statementMessageReader() { + if (!captureDbmsOutput) return JdbcExecutor.StatementMessageReader.NONE; + String setupWarning = dbmsOutputWarning; + boolean outputEnabled = dbmsOutputEnabled; + return new JdbcExecutor.StatementMessageReader() { + private boolean cleaned; + + @Override + public List read(Statement statement) throws SQLException { + List lines = new ArrayList<>(); + if (!outputEnabled) { + if (setupWarning != null) lines.add(setupWarning); + return lines; } - throw error; + int characters = 0; + boolean truncated = false; + try (CallableStatement call = statement.getConnection().prepareCall("{call DBMS_OUTPUT.GET_LINE(?, ?)}")) { + call.setQueryTimeout(2); + call.registerOutParameter(1, Types.VARCHAR); + call.registerOutParameter(2, Types.INTEGER); + long deadline = System.nanoTime() + 5_000_000_000L; + while (lines.size() < 256 && characters < 65_536) { + if (System.nanoTime() >= deadline) { + lines.add("DBMS_OUTPUT truncated after 5 seconds"); + truncated = true; + break; + } + call.execute(); + if (call.getInt(2) != 0) break; + String line = call.getString(1); + if (line == null) line = ""; + int remaining = 65_536 - characters; + if (line.length() > remaining) { + int end = remaining; + if (end > 0 && Character.isHighSurrogate(line.charAt(end - 1)) + && Character.isLowSurrogate(line.charAt(end))) { + end--; + } + lines.add(line.substring(0, end)); + lines.add("DBMS_OUTPUT truncated at 65536 characters"); + truncated = true; + break; + } + lines.add(line); + characters += line.length(); + } + if (!truncated && (lines.size() == 256 || characters >= 65_536)) { + if (System.nanoTime() >= deadline) { + lines.add("DBMS_OUTPUT truncated after 5 seconds"); + } else { + call.execute(); + if (call.getInt(2) == 0) lines.add("DBMS_OUTPUT truncated at the output limit"); + } + } + } catch (SQLException ignored) { + lines.add("DBMS_OUTPUT could not be read after statement execution"); + } finally { + cleanup(statement, lines); + } + return lines; + } + + @Override + public List close(Statement statement) throws SQLException { + if (!outputEnabled) return setupWarning == null ? List.of() : List.of(setupWarning); + return read(statement); + } + + private void cleanup(Statement statement, List messages) throws SQLException { + if (cleaned) return; + cleaned = true; + appendDbmsOutputCleanupWarning(statement.getConnection(), messages); + } + }; + } + + private static boolean isDbmsOutputStatement(String sql) { + if (sql == null) return false; + String leading = sql.replaceFirst("(?s)^(?:\\s|--[^\\r\\n]*(?:\\r?\\n|$)|/\\*.*?\\*/)*", ""); + return leading.matches("(?is)^(?:<<\\s*[\\p{L}_][\\p{L}\\p{N}_$#]*\\s*>>\\s*)?(?:BEGIN|DECLARE|CALL|EXEC|EXECUTE)\\b.*") + || leading.matches("(?is)^\\{\\s*(?:\\?\\s*=\\s*)?CALL\\b.*\\}"); + } + + private String enableDbmsOutput(Connection connection) { + try (var statement = connection.createStatement()) { + statement.setQueryTimeout(2); + statement.execute("BEGIN DBMS_OUTPUT.DISABLE; DBMS_OUTPUT.ENABLE(NULL); END;"); + return null; + } catch (SQLException error) { + if (error instanceof SQLTimeoutException || hasConnectionFailure(error)) { + try { + discardConnection(connection); + } catch (SQLException discardError) { + error.addSuppressed(discardError); + } + captureDbmsOutput = false; + dbmsOutputEnabled = false; + throw new RuntimeException("DBMS_OUTPUT setup failed; the connection was discarded", error); + } + if (isExpectedDbmsOutputUnavailable(error)) { + dbmsOutputEnabled = false; + return "DBMS_OUTPUT could not be enabled; statement execution continued"; + } + throw new RuntimeException("DBMS_OUTPUT setup failed", error); + } + } + + private static boolean hasConnectionFailure(SQLException error) { + for (SQLException current = error; current != null; current = current.getNextException()) { + String state = current.getSQLState(); + if (state != null && state.startsWith("08")) return true; + } + return false; + } + + private static boolean isExpectedDbmsOutputUnavailable(SQLException error) { + for (SQLException current = error; current != null; current = current.getNextException()) { + int code = Math.abs(current.getErrorCode()); + String message = current.getMessage(); + String lower = message == null ? "" : message.toLowerCase(Locale.ROOT); + if (code == 1031 || code == 5543 || code == 5546 || lower.contains("insufficient privilege") + || (lower.contains("not declared") || lower.contains("must be declared")) && lower.contains("dbms_output") + || lower.contains("dbms_output") && (lower.contains("not supported") || lower.contains("does not exist"))) { + return true; + } + } + return false; + } + + private void appendDbmsOutputCleanupWarning(Connection connection, List messages) { + try (var statement = connection.createStatement()) { + statement.setQueryTimeout(2); + statement.execute("BEGIN DBMS_OUTPUT.DISABLE; END;"); + } catch (SQLException cleanupError) { + try { + discardConnection(connection); + messages.add("DBMS_OUTPUT cleanup failed; the connection was discarded"); + } catch (SQLException discardError) { + cleanupError.addSuppressed(discardError); + messages.add("DBMS_OUTPUT cleanup and connection disposal failed"); + } + } + } + + private void cleanupDbmsOutputAfterFailure(Connection connection, RuntimeException error) { + try { + try (var statement = connection.createStatement()) { + statement.setQueryTimeout(2); + statement.execute("BEGIN DBMS_OUTPUT.DISABLE; END;"); + } + } catch (SQLException cleanupError) { + error.addSuppressed(cleanupError); + try { + discardConnection(connection); + } catch (SQLException discardError) { + error.addSuppressed(discardError); } - queryTimeoutChanged = true; } } diff --git a/agents/drivers/oceanbase-oracle/src/test/java/com/dbx/agent/oceanbaseoracle/OceanBaseOracleAgentTest.java b/agents/drivers/oceanbase-oracle/src/test/java/com/dbx/agent/oceanbaseoracle/OceanBaseOracleAgentTest.java index 405e7c6274..45b0915d7c 100644 --- a/agents/drivers/oceanbase-oracle/src/test/java/com/dbx/agent/oceanbaseoracle/OceanBaseOracleAgentTest.java +++ b/agents/drivers/oceanbase-oracle/src/test/java/com/dbx/agent/oceanbaseoracle/OceanBaseOracleAgentTest.java @@ -22,22 +22,404 @@ import org.junit.jupiter.params.ParameterizedTest; import org.junit.jupiter.params.provider.CsvSource; import java.lang.reflect.InvocationHandler; +import java.lang.reflect.Field; import java.lang.reflect.Method; import java.lang.reflect.Proxy; +import java.sql.CallableStatement; import java.sql.Connection; import java.sql.PreparedStatement; import java.sql.ResultSet; import java.sql.ResultSetMetaData; import java.sql.SQLException; import java.sql.SQLFeatureNotSupportedException; +import java.sql.SQLTimeoutException; import java.sql.Statement; import java.sql.Types; import java.util.ArrayList; import java.util.Arrays; import java.util.List; import java.util.Locale; +import java.util.concurrent.atomic.AtomicInteger; class OceanBaseOracleAgentTest { + @Test + void returnsDbmsOutputThroughRegularAndBoundedExecutionPaths() throws Exception { + assertDbmsOutputExecution(false); + } + + @Test + void returnsDbmsOutputWhenSessionTimeoutIsRejectedAsReadOnly() throws Exception { + assertDbmsOutputExecution(true); + } + + private void assertDbmsOutputExecution(boolean readOnlyTimeout) throws Exception { + String block = "BEGIN DBMS_OUTPUT.PUT_LINE('hello world'); END;"; + for (boolean paged : List.of(false, true)) { + AtomicInteger prepareCalls = new AtomicInteger(); + AtomicInteger getLineCalls = new AtomicInteger(); + List executedSql = new ArrayList<>(); + Connection[] connectionRef = new Connection[1]; + CallableStatement output = proxy(CallableStatement.class, (method, args) -> { + if ("execute".equals(method.getName())) { + getLineCalls.incrementAndGet(); + return false; + } + if ("getInt".equals(method.getName())) return getLineCalls.get() == 1 ? 0 : 1; + if ("getString".equals(method.getName())) return "hello world"; + return defaultValue(method.getReturnType()); + }); + Statement statement = proxy(Statement.class, (method, args) -> { + if ("execute".equals(method.getName())) { + executedSql.add(String.valueOf(args[0])); + if (readOnlyTimeout && String.valueOf(args[0]).startsWith("ALTER SESSION SET ob_query_timeout")) { + throw new SQLException("read only", "25006"); + } + return false; + } + if ("getConnection".equals(method.getName())) return connectionRef[0]; + if ("getUpdateCount".equals(method.getName())) return -1; + if ("getWarnings".equals(method.getName())) return null; + return defaultValue(method.getReturnType()); + }); + Connection connection = proxy(Connection.class, (method, args) -> { + if ("createStatement".equals(method.getName())) return statement; + if ("prepareCall".equals(method.getName())) { + prepareCalls.incrementAndGet(); + return output; + } + return defaultValue(method.getReturnType()); + }); + connectionRef[0] = connection; + OceanBaseOracleAgent agent = new OceanBaseOracleAgent(); + TestSupport.setPrivateConnection(agent, connection); + + List columns; + List> rows; + if (paged) { + QueryPageResult result = agent.executeQueryPage(block, null, new QueryPageOptions()); + columns = result.getColumns(); + rows = result.getRows(); + } else { + QueryResult result = agent.executeQuery(block, null, new ExecuteQueryOptions()); + columns = result.getColumns(); + rows = result.getRows(); + } + + Assertions.assertEquals(List.of("Message"), columns); + Assertions.assertEquals(List.of(List.of("hello world")), rows); + Assertions.assertEquals(1, prepareCalls.get()); + Assertions.assertEquals(2, getLineCalls.get()); + Assertions.assertEquals(1, executedSql.stream().filter(block::equals).count()); + Assertions.assertTrue(executedSql.contains("BEGIN DBMS_OUTPUT.DISABLE; DBMS_OUTPUT.ENABLE(NULL); END;")); + Assertions.assertTrue(executedSql.indexOf("BEGIN DBMS_OUTPUT.DISABLE; DBMS_OUTPUT.ENABLE(NULL); END;") + > executedSql.indexOf("ALTER SESSION SET ob_query_timeout")); + Assertions.assertEquals("BEGIN DBMS_OUTPUT.DISABLE; END;", executedSql.get(executedSql.size() - 1)); + + if (paged) { + agent.executeQueryPage("SELECT 1 FROM DUAL", null, new QueryPageOptions()); + } else { + agent.executeQuery("SELECT 1 FROM DUAL", null, new ExecuteQueryOptions()); + } + Assertions.assertEquals(1, prepareCalls.get(), "ordinary SQL must not read DBMS_OUTPUT"); + } + } + + @Test + void doesNotInstallDbmsOutputReaderForOrdinaryStatements() throws Exception { + OceanBaseOracleAgent agent = new OceanBaseOracleAgent(); + Method readerMethod = OceanBaseOracleAgent.class.getDeclaredMethod("statementMessageReader"); + readerMethod.setAccessible(true); + + Assertions.assertSame( + com.dbx.agent.JdbcExecutor.StatementMessageReader.NONE, + readerMethod.invoke(agent) + ); + } + + @Test + void recognizesLabeledAnonymousBlocksForDbmsOutput() throws Exception { + Method classify = OceanBaseOracleAgent.class.getDeclaredMethod("isDbmsOutputStatement", String.class); + classify.setAccessible(true); + + Assertions.assertEquals(true, classify.invoke(null, "<> BEGIN DBMS_OUTPUT.PUT_LINE('x'); END;")); + Assertions.assertEquals(true, classify.invoke(null, "{call app.print_output() }")); + Assertions.assertEquals(true, classify.invoke(null, "{ ? = call app.get_output() }")); + Assertions.assertEquals(false, classify.invoke(null, "SELECT 1 FROM DUAL")); + } + + @Test + void continuesWhenOceanBaseReportsDbmsOutputAsUndeclared() throws Exception { + String block = "BEGIN NULL; END;"; + AtomicInteger userStatementCalls = new AtomicInteger(); + AtomicInteger closedConnections = new AtomicInteger(); + Statement statement = proxy(Statement.class, (method, args) -> { + if ("execute".equals(method.getName())) { + String sql = String.valueOf(args[0]); + if (sql.startsWith("BEGIN DBMS_OUTPUT.DISABLE")) { + throw new SQLException("PLS-00201: identifier 'DBMS_OUTPUT' must be declared", "42000", 5543); + } + if (block.equals(sql)) userStatementCalls.incrementAndGet(); + return false; + } + if ("getUpdateCount".equals(method.getName())) return -1; + if ("getWarnings".equals(method.getName())) return null; + return defaultValue(method.getReturnType()); + }); + Connection connection = proxy(Connection.class, (method, args) -> { + if ("createStatement".equals(method.getName())) return statement; + if ("close".equals(method.getName())) { + closedConnections.incrementAndGet(); + return null; + } + return defaultValue(method.getReturnType()); + }); + OceanBaseOracleAgent agent = new OceanBaseOracleAgent(); + TestSupport.setPrivateConnection(agent, connection); + + QueryResult result = agent.executeQuery(block, null, new ExecuteQueryOptions()); + + Assertions.assertEquals(1, userStatementCalls.get()); + Assertions.assertEquals(0, closedConnections.get()); + Assertions.assertEquals(List.of("Message"), result.getColumns()); + Assertions.assertTrue(result.getRows().stream().flatMap(List::stream) + .anyMatch(value -> String.valueOf(value).contains("could not be enabled"))); + } + + @Test + void doesNotExecuteUserSqlAfterDbmsOutputSetupTimeout() throws Exception { + String block = "BEGIN DBMS_OUTPUT.PUT_LINE('hello world'); END;"; + AtomicInteger userStatementCalls = new AtomicInteger(); + AtomicInteger closedConnections = new AtomicInteger(); + Statement statement = proxy(Statement.class, (method, args) -> { + if ("execute".equals(method.getName())) { + String sql = String.valueOf(args[0]); + if (block.equals(sql)) userStatementCalls.incrementAndGet(); + if (sql.startsWith("ALTER SESSION SET ob_query_timeout")) return false; + throw new SQLTimeoutException("setup timed out", "HYT00"); + } + return defaultValue(method.getReturnType()); + }); + Connection connection = proxy(Connection.class, (method, args) -> { + if ("createStatement".equals(method.getName())) return statement; + if ("close".equals(method.getName())) { + closedConnections.incrementAndGet(); + return null; + } + return defaultValue(method.getReturnType()); + }); + OceanBaseOracleAgent agent = new OceanBaseOracleAgent(); + TestSupport.setPrivateConnection(agent, connection); + + Assertions.assertThrows(RuntimeException.class, + () -> agent.executeQuery(block, null, new ExecuteQueryOptions())); + + Assertions.assertEquals(0, userStatementCalls.get()); + Assertions.assertEquals(1, closedConnections.get()); + } + + @Test + void doesNotReportTruncationWhenOutputEndsExactlyAtLimits() throws Exception { + assertOutputLimit(List.of()); + assertOutputLimit(java.util.Collections.nCopies(256, "x")); + assertOutputLimit(List.of("x".repeat(32_767), "x".repeat(32_767), "x".repeat(2))); + } + + @Test + void doesNotSplitSupplementaryUnicodeWhenTruncatingOutput() throws Exception { + List messages = readOutputLines(List.of("x".repeat(65_535), "😀")); + + Assertions.assertEquals(65_535, messages.get(0).length()); + Assertions.assertFalse(messages.get(0).contains("\uD800")); + Assertions.assertEquals("", messages.get(1)); + Assertions.assertTrue(messages.get(2).contains("truncated")); + } + + @Test + void reportsTruncationOnlyWhenAnotherLineExistsPastLineLimit() throws Exception { + List outputLines = new ArrayList<>(java.util.Collections.nCopies(257, "x")); + AtomicInteger callCount = new AtomicInteger(); + Connection[] connectionRef = new Connection[1]; + CallableStatement output = proxy(CallableStatement.class, (method, args) -> { + if ("execute".equals(method.getName())) { + callCount.incrementAndGet(); + return false; + } + if ("getInt".equals(method.getName())) return callCount.get() <= outputLines.size() ? 0 : 1; + if ("getString".equals(method.getName())) return outputLines.get(callCount.get() - 1); + return defaultValue(method.getReturnType()); + }); + Statement statement = proxy(Statement.class, (method, args) -> + "getConnection".equals(method.getName()) ? connectionRef[0] : defaultValue(method.getReturnType())); + Connection connection = proxy(Connection.class, (method, args) -> { + if ("prepareCall".equals(method.getName())) return output; + if ("createStatement".equals(method.getName())) return proxy(Statement.class, (m, a) -> defaultValue(m.getReturnType())); + return defaultValue(method.getReturnType()); + }); + connectionRef[0] = connection; + OceanBaseOracleAgent agent = new OceanBaseOracleAgent(); + setField(agent, "captureDbmsOutput", true); + setField(agent, "dbmsOutputEnabled", true); + Method readerMethod = OceanBaseOracleAgent.class.getDeclaredMethod("statementMessageReader"); + readerMethod.setAccessible(true); + + @SuppressWarnings("unchecked") + List messages = ((com.dbx.agent.JdbcExecutor.StatementMessageReader) readerMethod.invoke(agent)).read(statement); + + Assertions.assertEquals(257, callCount.get()); + Assertions.assertEquals(257, messages.size()); + Assertions.assertEquals("DBMS_OUTPUT truncated at the output limit", messages.get(256)); + } + + private static void assertOutputLimit(List outputLines) throws Exception { + List messages = readOutputLines(outputLines); + Assertions.assertEquals(outputLines, messages); + Assertions.assertFalse(messages.stream().anyMatch(message -> message.contains("truncated"))); + } + + private static List readOutputLines(List outputLines) throws Exception { + AtomicInteger callCount = new AtomicInteger(); + Connection[] connectionRef = new Connection[1]; + CallableStatement output = proxy(CallableStatement.class, (method, args) -> { + if ("execute".equals(method.getName())) { + callCount.incrementAndGet(); + return false; + } + if ("getInt".equals(method.getName())) return callCount.get() <= outputLines.size() ? 0 : 1; + if ("getString".equals(method.getName())) return outputLines.get(callCount.get() - 1); + return defaultValue(method.getReturnType()); + }); + Statement statement = proxy(Statement.class, (method, args) -> + "getConnection".equals(method.getName()) ? connectionRef[0] : defaultValue(method.getReturnType())); + Connection connection = proxy(Connection.class, (method, args) -> { + if ("prepareCall".equals(method.getName())) return output; + if ("createStatement".equals(method.getName())) return proxy(Statement.class, (m, a) -> defaultValue(m.getReturnType())); + return defaultValue(method.getReturnType()); + }); + connectionRef[0] = connection; + OceanBaseOracleAgent agent = new OceanBaseOracleAgent(); + setField(agent, "captureDbmsOutput", true); + setField(agent, "dbmsOutputEnabled", true); + Method readerMethod = OceanBaseOracleAgent.class.getDeclaredMethod("statementMessageReader"); + readerMethod.setAccessible(true); + + @SuppressWarnings("unchecked") + List messages = ((com.dbx.agent.JdbcExecutor.StatementMessageReader) readerMethod.invoke(agent)).read(statement); + + int characters = outputLines.stream().mapToInt(String::length).sum(); + int expectedCalls = outputLines.size() > 256 || characters > 65_536 + ? Math.min(outputLines.size(), 256) + : outputLines.size() + 1; + Assertions.assertEquals(expectedCalls, callCount.get()); + return messages; + } + + @Test + void discardsConnectionWhenDbmsOutputCleanupFails() throws Exception { + String block = "BEGIN DBMS_OUTPUT.PUT_LINE('hello world'); END;"; + AtomicInteger outputSetupCalls = new AtomicInteger(); + AtomicInteger userStatementCalls = new AtomicInteger(); + AtomicInteger closedConnections = new AtomicInteger(); + Connection[] connectionRef = new Connection[1]; + CallableStatement output = proxy(CallableStatement.class, (method, args) -> { + if ("execute".equals(method.getName())) throw new SQLException("GET_LINE failed"); + return defaultValue(method.getReturnType()); + }); + Statement statement = proxy(Statement.class, (method, args) -> { + if ("execute".equals(method.getName())) { + String sql = String.valueOf(args[0]); + if (sql.startsWith("BEGIN DBMS_OUTPUT.DISABLE")) { + if (outputSetupCalls.incrementAndGet() > 1) throw new SQLException("cleanup failed"); + } + if (block.equals(sql)) userStatementCalls.incrementAndGet(); + return false; + } + if ("getConnection".equals(method.getName())) return connectionRef[0]; + if ("getUpdateCount".equals(method.getName())) return -1; + if ("getWarnings".equals(method.getName())) return null; + return defaultValue(method.getReturnType()); + }); + Connection connection = proxy(Connection.class, (method, args) -> { + if ("createStatement".equals(method.getName())) return statement; + if ("prepareCall".equals(method.getName())) return output; + if ("close".equals(method.getName())) { + closedConnections.incrementAndGet(); + return null; + } + return defaultValue(method.getReturnType()); + }); + connectionRef[0] = connection; + OceanBaseOracleAgent agent = new OceanBaseOracleAgent(); + TestSupport.setPrivateConnection(agent, connection); + + QueryResult result = agent.executeQuery(block, null, new ExecuteQueryOptions()); + + Assertions.assertEquals(List.of("Message"), result.getColumns()); + Assertions.assertTrue(result.getRows().stream().flatMap(List::stream) + .anyMatch(value -> String.valueOf(value).contains("DBMS_OUTPUT cleanup failed"))); + Assertions.assertEquals(1, closedConnections.get(), "a connection with unknown DBMS_OUTPUT state must be discarded"); + Assertions.assertEquals(2, outputSetupCalls.get(), "one initial enable and one failed reset are expected"); + Assertions.assertEquals(1, userStatementCalls.get(), "cleanup failure must not retry the user statement"); + } + + @Test + void cleansDbmsOutputWhenPagedPlsqlReturnsAResultSet() throws Exception { + String block = "BEGIN OPEN :cursor FOR SELECT 1 FROM DUAL; END;"; + List executedSql = new ArrayList<>(); + int[] row = {-1}; + AtomicInteger outputCalls = new AtomicInteger(); + ResultSetMetaData meta = proxy(ResultSetMetaData.class, (method, args) -> { + if ("getColumnCount".equals(method.getName())) return 1; + if ("getColumnLabel".equals(method.getName())) return "N"; + if ("getColumnType".equals(method.getName())) return Types.INTEGER; + if ("getColumnTypeName".equals(method.getName())) return "NUMBER"; + return defaultValue(method.getReturnType()); + }); + ResultSet cursor = proxy(ResultSet.class, (method, args) -> { + if ("next".equals(method.getName())) return ++row[0] == 0; + if ("getMetaData".equals(method.getName())) return meta; + if ("getObject".equals(method.getName()) || "getInt".equals(method.getName())) return 1; + return defaultValue(method.getReturnType()); + }); + CallableStatement output = proxy(CallableStatement.class, (method, args) -> { + if ("execute".equals(method.getName())) { + outputCalls.incrementAndGet(); + return false; + } + if ("getInt".equals(method.getName())) return outputCalls.get() == 1 ? 0 : 1; + if ("getString".equals(method.getName())) return "cursor output"; + return defaultValue(method.getReturnType()); + }); + Connection[] connectionRef = new Connection[1]; + Statement statement = proxy(Statement.class, (method, args) -> { + if ("execute".equals(method.getName())) { + String sql = String.valueOf(args[0]); + executedSql.add(sql); + return block.equals(sql); + } + if ("getResultSet".equals(method.getName())) return cursor; + if ("getConnection".equals(method.getName())) return connectionRef[0]; + if ("getUpdateCount".equals(method.getName())) return -1; + return defaultValue(method.getReturnType()); + }); + Connection connection = proxy(Connection.class, (method, args) -> { + if ("createStatement".equals(method.getName())) return statement; + if ("prepareCall".equals(method.getName())) return output; + return defaultValue(method.getReturnType()); + }); + connectionRef[0] = connection; + OceanBaseOracleAgent agent = new OceanBaseOracleAgent(); + TestSupport.setPrivateConnection(agent, connection); + + QueryPageResult result = agent.executeQueryPage(block, null, new QueryPageOptions(10, null, 10, 0)); + + Assertions.assertEquals(List.of("N"), result.getColumns()); + Assertions.assertEquals(List.of(List.of(1)), result.getRows()); + Assertions.assertTrue(result.getMessages().stream().anyMatch(message -> "cursor output".equals(message.get("message")))); + Assertions.assertEquals(2, outputCalls.get()); + Assertions.assertTrue(executedSql.contains("BEGIN DBMS_OUTPUT.DISABLE; END;"), + "closing the completed result-set session must disable DBMS_OUTPUT"); + } + @Test void buildsOceanBaseJdbcUrl() { ConnectParams params = new ConnectParams();