Skip to content

Commit

Permalink
Move JDBC warnings to Statement
Browse files Browse the repository at this point in the history
  • Loading branch information
@electrum
electrum committed Oct 3, 2019
1 parent 25d91e6 commit 3bcda70
Show file tree
Hide file tree
Showing 2 changed files with 48 additions and 58 deletions.
23 changes: 1 addition & 22 deletions presto-jdbc/src/main/java/io/prestosql/jdbc/PrestoResultSet.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -104,7 +104,6 @@ public class PrestoResultSet
private final AtomicReference<List<Object>> 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<QueryStats> progressCallback, WarningsManager warningsManager)
throws SQLException
Expand All @@ -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);
}
Expand Down Expand Up @@ -483,15 +481,14 @@ public SQLWarning getWarnings()
throws SQLException
{
checkOpen();
return warningsManager.getWarnings();
return null;
}

@Override
public void clearWarnings()
throws SQLException
{
checkOpen();
warningsManager.clearWarnings();
}

@Override
Expand Down Expand Up @@ -1761,35 +1758,17 @@ private static class ResultsPageIterator
private final StatementClient client;
private final Consumer<QueryStats> progressCallback;
private final WarningsManager warningsManager;
private final boolean isQuery;

private ResultsPageIterator(StatementClient client, Consumer<QueryStats> 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<List<Object>> computeNext()
{
if (isQuery) {
// Clear the warnings if this is a query, per ResultSet javadoc
warningsManager.clearWarnings();
}
while (client.isRunning()) {
checkInterruption(null);

Expand Down
83 changes: 47 additions & 36 deletions presto-jdbc/src/test/java/io/prestosql/jdbc/TestJdbcWarnings.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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<WarningEntry> 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<WarningEntry> 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
Expand All @@ -195,7 +169,7 @@ public void testExecuteQueryWarnings()
assertNull(rs.getWarnings());
Set<WarningEntry> currentWarnings = new HashSet<>();
while (rs.next()) {
assertWarnings(rs.getWarnings(), currentWarnings);
assertWarnings(statement.getWarnings(), currentWarnings);
}

TestingWarningCollectorConfig warningCollectorConfig = new TestingWarningCollectorConfig().setPreloadedWarnings(PRELOADED_WARNINGS).setAddWarnings(true);
Expand All @@ -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<WarningEntry> 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<PrestoWarning> warnings)
{
requireNonNull(warnings, "warnings is null");
Expand Down

0 comments on commit 3bcda70

Please sign in to comment.