Skip to content

Commit e5ccd4a

Browse files
committed
Introduce a wrapped result set
This allows registering a callback to correctly log the time for queries that iterate through a result set.
1 parent 8ccca43 commit e5ccd4a

18 files changed

Lines changed: 103 additions & 47 deletions

src/sqlancer/ComparatorHelper.java

Lines changed: 1 addition & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,6 @@
11
package sqlancer;
22

33
import java.io.IOException;
4-
import java.sql.ResultSet;
54
import java.sql.SQLException;
65
import java.util.ArrayList;
76
import java.util.HashSet;
@@ -46,7 +45,7 @@ public static List<String> getResultSetFirstColumnAsString(String queryString, S
4645
}
4746
QueryAdapter q = new QueryAdapter(queryString, errors);
4847
List<String> resultSet = new ArrayList<>();
49-
ResultSet result = null;
48+
SQLancerResultSet result = null;
5049
try {
5150
result = q.executeAndGet(state);
5251
if (result == null) {
@@ -55,7 +54,6 @@ public static List<String> getResultSetFirstColumnAsString(String queryString, S
5554
while (result.next()) {
5655
resultSet.add(result.getString(1));
5756
}
58-
result.getStatement().close();
5957
} catch (Exception e) {
6058
if (e instanceof IgnoreMeException) {
6159
throw e;
@@ -75,7 +73,6 @@ public static List<String> getResultSetFirstColumnAsString(String queryString, S
7573
throw new AssertionError(queryString, e);
7674
} finally {
7775
if (result != null && !result.isClosed()) {
78-
result.getStatement().close();
7976
result.close();
8077
}
8178
}

src/sqlancer/GlobalState.java

Lines changed: 11 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,6 @@
11
package sqlancer;
22

33
import java.sql.Connection;
4-
import java.sql.ResultSet;
54
import java.sql.SQLException;
65

76
import sqlancer.Main.QueryManager;
@@ -129,11 +128,19 @@ public boolean executeStatement(Query q, String... fills) throws SQLException {
129128
return success;
130129
}
131130

132-
public ResultSet executeStatementAndGet(Query q, String... fills) throws SQLException {
131+
public SQLancerResultSet executeStatementAndGet(Query q, String... fills) throws SQLException {
133132
ExecutionTimer timer = executePrologue(q);
134-
ResultSet result = manager.executeAndGet(q, fills);
133+
SQLancerResultSet result = manager.executeAndGet(q, fills);
135134
boolean success = result != null;
136-
executeEpilogue(q, success, timer);
135+
if (success) {
136+
result.registerEpilogue(() -> {
137+
try {
138+
executeEpilogue(q, success, timer);
139+
} catch (SQLException e) {
140+
throw new AssertionError(e);
141+
}
142+
});
143+
}
137144
return result;
138145
}
139146

src/sqlancer/Main.java

Lines changed: 2 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -8,7 +8,6 @@
88
import java.io.Writer;
99
import java.nio.file.Files;
1010
import java.sql.Connection;
11-
import java.sql.ResultSet;
1211
import java.sql.SQLException;
1312
import java.sql.SQLFeatureNotSupportedException;
1413
import java.text.DateFormat;
@@ -260,9 +259,9 @@ public boolean execute(Query q, String... fills) throws SQLException {
260259
return success;
261260
}
262261

263-
public ResultSet executeAndGet(Query q, String... fills) throws SQLException {
262+
public SQLancerResultSet executeAndGet(Query q, String... fills) throws SQLException {
264263
globalState.getState().logStatement(q);
265-
ResultSet result;
264+
SQLancerResultSet result;
266265
result = q.executeAndGet(globalState, fills);
267266
Main.nrSuccessfulActions.addAndGet(1);
268267
return result;

src/sqlancer/Query.java

Lines changed: 2 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,5 @@
11
package sqlancer;
22

3-
import java.sql.ResultSet;
43
import java.sql.SQLException;
54
import java.util.Collection;
65

@@ -32,7 +31,7 @@ public String toString() {
3231
return getQueryString();
3332
}
3433

35-
public ResultSet executeAndGet(GlobalState<?, ?> globalState, String... fills) throws SQLException {
34+
public SQLancerResultSet executeAndGet(GlobalState<?, ?> globalState, String... fills) throws SQLException {
3635
throw new AssertionError();
3736
}
3837

@@ -41,7 +40,7 @@ public boolean executeLogged(GlobalState<?, ?> globalState) throws SQLException
4140
return execute(globalState);
4241
}
4342

44-
public ResultSet executeAndGetLogged(GlobalState<?, ?> globalState) throws SQLException {
43+
public SQLancerResultSet executeAndGetLogged(GlobalState<?, ?> globalState) throws SQLException {
4544
logQueryString(globalState);
4645
return executeAndGet(globalState);
4746
}

src/sqlancer/QueryAdapter.java

Lines changed: 5 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -97,7 +97,7 @@ public void checkException(Exception e) throws AssertionError {
9797
}
9898

9999
@Override
100-
public ResultSet executeAndGet(GlobalState<?, ?> globalState, String... fills) throws SQLException {
100+
public SQLancerResultSet executeAndGet(GlobalState<?, ?> globalState, String... fills) throws SQLException {
101101
Statement s;
102102
if (fills.length > 0) {
103103
s = globalState.getConnection().prepareStatement(getQueryString());
@@ -115,7 +115,10 @@ public ResultSet executeAndGet(GlobalState<?, ?> globalState, String... fills) t
115115
result = s.executeQuery(query);
116116
}
117117
Main.nrSuccessfulActions.addAndGet(1);
118-
return result;
118+
if (result == null) {
119+
return null;
120+
}
121+
return new SQLancerResultSet(result);
119122
} catch (Exception e) {
120123
s.close();
121124
Main.nrUnsuccessfulActions.addAndGet(1);
Lines changed: 53 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,53 @@
1+
package sqlancer;
2+
3+
import java.io.Closeable;
4+
import java.sql.ResultSet;
5+
import java.sql.SQLException;
6+
7+
public class SQLancerResultSet implements Closeable {
8+
9+
ResultSet rs;
10+
private Runnable runnableEpilogue;
11+
12+
public SQLancerResultSet(ResultSet rs) {
13+
this.rs = rs;
14+
}
15+
16+
@Override
17+
public void close() {
18+
try {
19+
if (runnableEpilogue != null) {
20+
runnableEpilogue.run();
21+
}
22+
rs.getStatement().close();
23+
rs.close();
24+
} catch (SQLException e) {
25+
throw new AssertionError(e);
26+
}
27+
}
28+
29+
public boolean next() throws SQLException {
30+
return rs.next();
31+
}
32+
33+
public int getInt(int i) throws SQLException {
34+
return rs.getInt(i);
35+
}
36+
37+
public String getString(int i) throws SQLException {
38+
return rs.getString(i);
39+
}
40+
41+
public boolean isClosed() throws SQLException {
42+
return rs.isClosed();
43+
}
44+
45+
public long getLong(int i) throws SQLException {
46+
return rs.getLong(i);
47+
}
48+
49+
public void registerEpilogue(Runnable runnableEpilogue) {
50+
this.runnableEpilogue = runnableEpilogue;
51+
}
52+
53+
}

src/sqlancer/clickhouse/oracle/tlp/ClickHouseTLPAggregateOracle.java

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,5 @@
11
package sqlancer.clickhouse.oracle.tlp;
22

3-
import java.sql.ResultSet;
43
import java.sql.SQLException;
54
import java.util.Arrays;
65
import java.util.List;
@@ -10,6 +9,7 @@
109
import sqlancer.IgnoreMeException;
1110
import sqlancer.QueryAdapter;
1211
import sqlancer.Randomly;
12+
import sqlancer.SQLancerResultSet;
1313
import sqlancer.TestOracle;
1414
import sqlancer.clickhouse.ClickHouseProvider;
1515
import sqlancer.clickhouse.ClickHouseSchema;
@@ -69,7 +69,7 @@ public void check() throws SQLException {
6969
String firstResult;
7070
String secondResult;
7171
QueryAdapter q = new QueryAdapter(originalQuery);
72-
try (ResultSet result = q.executeAndGet(state)) {
72+
try (SQLancerResultSet result = q.executeAndGet(state)) {
7373
if (result == null) {
7474
throw new IgnoreMeException();
7575
}
@@ -80,7 +80,7 @@ public void check() throws SQLException {
8080
}
8181

8282
QueryAdapter q2 = new QueryAdapter(metamorphicText);
83-
try (ResultSet result = q2.executeAndGet(state)) {
83+
try (SQLancerResultSet result = q2.executeAndGet(state)) {
8484
if (result == null) {
8585
throw new IgnoreMeException();
8686
}

src/sqlancer/cockroachdb/oracle/CockroachDBNoRECOracle.java

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,5 @@
11
package sqlancer.cockroachdb.oracle;
22

3-
import java.sql.ResultSet;
43
import java.sql.SQLException;
54
import java.util.ArrayList;
65
import java.util.Arrays;
@@ -14,6 +13,7 @@
1413
import sqlancer.Query;
1514
import sqlancer.QueryAdapter;
1615
import sqlancer.Randomly;
16+
import sqlancer.SQLancerResultSet;
1717
import sqlancer.TestOracle;
1818
import sqlancer.cockroachdb.CockroachDBCommon;
1919
import sqlancer.cockroachdb.CockroachDBErrors;
@@ -137,7 +137,7 @@ private int getNonOptimizedResult(CockroachDBExpression whereCondition, List<Coc
137137

138138
private int getCount(GlobalState<?, ?> globalState, Query q) throws AssertionError {
139139
int count = 0;
140-
try (ResultSet rs = q.executeAndGet(globalState)) {
140+
try (SQLancerResultSet rs = q.executeAndGet(globalState)) {
141141
if (rs == null) {
142142
return -1;
143143
}

src/sqlancer/cockroachdb/oracle/tlp/CockroachDBTLPAggregateOracle.java

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,5 @@
11
package sqlancer.cockroachdb.oracle.tlp;
22

3-
import java.sql.ResultSet;
43
import java.sql.SQLException;
54
import java.util.ArrayList;
65
import java.util.Arrays;
@@ -15,6 +14,7 @@
1514
import sqlancer.IgnoreMeException;
1615
import sqlancer.QueryAdapter;
1716
import sqlancer.Randomly;
17+
import sqlancer.SQLancerResultSet;
1818
import sqlancer.TestOracle;
1919
import sqlancer.cockroachdb.CockroachDBCommon;
2020
import sqlancer.cockroachdb.CockroachDBErrors;
@@ -118,7 +118,7 @@ private String createMetamorphicUnionQuery(CockroachDBSelect select, CockroachDB
118118
private String getAggregateResult(String queryString) throws SQLException {
119119
String resultString;
120120
QueryAdapter q = new QueryAdapter(queryString, errors);
121-
try (ResultSet result = q.executeAndGet(state)) {
121+
try (SQLancerResultSet result = q.executeAndGet(state)) {
122122
if (result == null) {
123123
throw new IgnoreMeException();
124124
}

src/sqlancer/duckdb/test/DuckDBNoRECOracle.java

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -13,6 +13,7 @@
1313
import sqlancer.Query;
1414
import sqlancer.QueryAdapter;
1515
import sqlancer.Randomly;
16+
import sqlancer.SQLancerResultSet;
1617
import sqlancer.TestOracle;
1718
import sqlancer.ast.newast.ColumnReferenceNode;
1819
import sqlancer.ast.newast.NewPostfixTextNode;
@@ -83,7 +84,7 @@ private int getSecondQuery(List<Node<DuckDBExpression>> tableList, Node<DuckDBEx
8384
unoptimizedQueryString = "SELECT SUM(count) FROM (" + DuckDBToStringVisitor.asString(select) + ") as res";
8485
errors.add("canceling statement due to statement timeout");
8586
Query q = new QueryAdapter(unoptimizedQueryString, errors);
86-
ResultSet rs;
87+
SQLancerResultSet rs;
8788
try {
8889
rs = q.executeAndGetLogged(state);
8990
} catch (Exception e) {

0 commit comments

Comments
 (0)