diff --git a/presto-jdbc/src/main/java/io/prestosql/jdbc/PrestoResultSet.java b/presto-jdbc/src/main/java/io/prestosql/jdbc/PrestoResultSet.java index 77ea0950ffe4..9cb9409e35bc 100644 --- a/presto-jdbc/src/main/java/io/prestosql/jdbc/PrestoResultSet.java +++ b/presto-jdbc/src/main/java/io/prestosql/jdbc/PrestoResultSet.java @@ -104,7 +104,6 @@ public class PrestoResultSet private final AtomicReference> row = new AtomicReference<>(); private final AtomicBoolean wasNull = new AtomicBoolean(); private final AtomicBoolean closed = new AtomicBoolean(); - private final WarningsManager warningsManager; PrestoResultSet(StatementClient client, long maxRows, Consumer progressCallback, WarningsManager warningsManager) throws SQLException @@ -119,7 +118,6 @@ public class PrestoResultSet this.fieldMap = getFieldMap(columns); this.columnInfoList = getColumnInfo(columns); this.resultSetMetaData = new PrestoResultSetMetaData(columnInfoList); - this.warningsManager = requireNonNull(warningsManager, "warningsManager is null"); this.results = flatten(new ResultsPageIterator(client, progressCallback, warningsManager), maxRows); } @@ -483,7 +481,7 @@ public SQLWarning getWarnings() throws SQLException { checkOpen(); - return warningsManager.getWarnings(); + return null; } @Override @@ -491,7 +489,6 @@ public void clearWarnings() throws SQLException { checkOpen(); - warningsManager.clearWarnings(); } @Override @@ -1761,35 +1758,17 @@ private static class ResultsPageIterator private final StatementClient client; private final Consumer progressCallback; private final WarningsManager warningsManager; - private final boolean isQuery; private ResultsPageIterator(StatementClient client, Consumer progressCallback, WarningsManager warningsManager) { this.client = requireNonNull(client, "client is null"); this.progressCallback = requireNonNull(progressCallback, "progressCallback is null"); this.warningsManager = requireNonNull(warningsManager, "warningsManager is null"); - this.isQuery = isQuery(client); - } - - private static boolean isQuery(StatementClient client) - { - String updateType; - if (client.isRunning()) { - updateType = client.currentStatusInfo().getUpdateType(); - } - else { - updateType = client.finalStatusInfo().getUpdateType(); - } - return updateType == null; } @Override protected Iterable> computeNext() { - if (isQuery) { - // Clear the warnings if this is a query, per ResultSet javadoc - warningsManager.clearWarnings(); - } while (client.isRunning()) { checkInterruption(null); diff --git a/presto-jdbc/src/test/java/io/prestosql/jdbc/TestJdbcWarnings.java b/presto-jdbc/src/test/java/io/prestosql/jdbc/TestJdbcWarnings.java index fe15c09ef507..0324bc800fb9 100644 --- a/presto-jdbc/src/test/java/io/prestosql/jdbc/TestJdbcWarnings.java +++ b/presto-jdbc/src/test/java/io/prestosql/jdbc/TestJdbcWarnings.java @@ -135,54 +135,28 @@ public void testStatementWarnings() @Test public void testLongRunningStatement() - throws SQLException, InterruptedException + throws Exception { Future future = executor.submit(() -> { statement.execute("CREATE TABLE test_long_running AS SELECT * FROM slow_table"); return null; }); - while (statement.getWarnings() == null) { - Thread.sleep(100); - } - int expectedWarnings = 10; - SQLWarning warning = statement.getWarnings(); - Set currentWarnings = new HashSet<>(); - assertTrue(currentWarnings.add(new WarningEntry(warning))); - for (int warnings = 1; !future.isDone() && warnings < expectedWarnings; warnings++) { - for (SQLWarning nextWarning = warning.getNextWarning(); nextWarning == null; nextWarning = warning.getNextWarning()) { - // Wait for new warnings - } - warning = warning.getNextWarning(); - assertTrue(currentWarnings.add(new WarningEntry(warning))); - Thread.sleep(100); - } - assertThat(currentWarnings).size().isGreaterThanOrEqualTo(expectedWarnings); + assertStatementWarnings(statement, future); + statement.execute("DROP TABLE test_long_running"); } @Test public void testLongRunningQuery() - throws SQLException, InterruptedException + throws Exception { Future future = executor.submit(() -> { - statement.execute("SELECT * FROM slow_table"); + ResultSet resultSet = statement.executeQuery("SELECT * FROM slow_table"); + while (resultSet.next()) { + // discard results + } return null; }); - while (statement.getResultSet() == null) { - Thread.sleep(100); - } - ResultSet resultSet = statement.getResultSet(); - Set currentWarnings = new HashSet<>(); - for (int rows = 0; !future.isDone() && rows < 10; ) { - if (resultSet.next()) { - for (SQLWarning warning = resultSet.getWarnings(); warning.getNextWarning() != null; warning = warning.getNextWarning()) { - assertTrue(currentWarnings.add(new WarningEntry(warning.getNextWarning()))); - } - } - else { - break; - } - Thread.sleep(100); - } + assertStatementWarnings(statement, future); } @Test @@ -195,7 +169,7 @@ public void testExecuteQueryWarnings() assertNull(rs.getWarnings()); Set currentWarnings = new HashSet<>(); while (rs.next()) { - assertWarnings(rs.getWarnings(), currentWarnings); + assertWarnings(statement.getWarnings(), currentWarnings); } TestingWarningCollectorConfig warningCollectorConfig = new TestingWarningCollectorConfig().setPreloadedWarnings(PRELOADED_WARNINGS).setAddWarnings(true); @@ -222,6 +196,43 @@ public void testSqlWarning() assertWarningsEqual(warning.getNextWarning().getNextWarning(), toPrestoSqlWarning(warnings.get(2))); } + private static void assertStatementWarnings(Statement statement, Future future) + throws Exception + { + // wait for initial warnings + while (!future.isDone() && statement.getWarnings() == null) { + Thread.sleep(100); + } + + Set warnings = new HashSet<>(); + SQLWarning warning = statement.getWarnings(); + + // collect initial set of warnings + assertTrue(warnings.add(new WarningEntry(warning))); + while (warning.getNextWarning() != null) { + warning = warning.getNextWarning(); + assertTrue(warnings.add(new WarningEntry(warning))); + } + + int initialSize = warnings.size(); + assertThat(initialSize).isGreaterThanOrEqualTo(PRELOADED_WARNINGS + 1); + + // collect additional warnings until query finish + while (!future.isDone()) { + if (warning.getNextWarning() == null) { + Thread.sleep(100); + continue; + } + warning = warning.getNextWarning(); + assertTrue(warnings.add(new WarningEntry(warning))); + } + + int finalSize = warnings.size(); + assertThat(finalSize).isGreaterThan(initialSize); + + future.get(); + } + private static SQLWarning fromPrestoWarnings(List warnings) { requireNonNull(warnings, "warnings is null");