Skip to content

Commit c08d71a

Browse files
emolssonemolsson
committed
JAVA-1470: Querying multiple pages overrides WrappedStatement
1 parent c7208ce commit c08d71a

3 files changed

Lines changed: 68 additions & 4 deletions

File tree

‎changelog/README.md‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,7 @@
11
## Changelog
22

33
### 3.3.0 (in progress)
4-
4+
- [bug] JAVA-1470: Querying multiple pages overrides WrappedStatement
55

66
### 3.2.0
77

‎driver-core/src/main/java/com/datastax/driver/core/ArrayBackedResultSet.java‎

Lines changed: 4 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -62,11 +62,12 @@ static ArrayBackedResultSet fromMessage(Responses.Result msg, SessionManager ses
6262

6363
ColumnDefinitions columnDefs;
6464
if (r.metadata.columns == null) {
65+
Statement actualStatement = statement;
6566
if (statement instanceof StatementWrapper) {
66-
statement = ((StatementWrapper) statement).getWrappedStatement();
67+
actualStatement = ((StatementWrapper) statement).getWrappedStatement();
6768
}
68-
assert statement instanceof BoundStatement;
69-
columnDefs = ((BoundStatement) statement).statement.getPreparedId().resultSetMetadata;
69+
assert actualStatement instanceof BoundStatement;
70+
columnDefs = ((BoundStatement) actualStatement).statement.getPreparedId().resultSetMetadata;
7071
assert columnDefs != null;
7172
} else {
7273
columnDefs = r.metadata.columns;

‎driver-core/src/test/java/com/datastax/driver/core/StatementWrapperTest.java‎

Lines changed: 63 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -30,13 +30,17 @@ public class StatementWrapperTest extends CCMTestsSupport {
3030
private static final String INSERT_QUERY = "insert into test (k, v) values (?, ?)";
3131
private static final String SELECT_QUERY = "select * from test where k = ?";
3232

33+
private static final String INSERT_MULTIPAGE_QUERY = "insert into test_multipage (k, v) values (?, ?)";
34+
private static final String SELECT_MULTIPAGE_QUERY = "select * from test_multipage where k = ?";
35+
3336
CustomLoadBalancingPolicy loadBalancingPolicy = new CustomLoadBalancingPolicy();
3437
CustomSpeculativeExecutionPolicy speculativeExecutionPolicy = new CustomSpeculativeExecutionPolicy();
3538
CustomRetryPolicy retryPolicy = new CustomRetryPolicy();
3639

3740
@Override
3841
public void onTestContextInitialized() {
3942
execute("create table test (k text primary key, v int)");
43+
execute("create table test_multipage (k text, v int, primary key (k, v))");
4044
}
4145

4246
@Override
@@ -59,6 +63,65 @@ public void should_pass_wrapped_statement_to_load_balancing_policy() {
5963
assertThat(loadBalancingPolicy.customStatementsHandled.get()).isEqualTo(1);
6064
}
6165

66+
@Test(groups = "short")
67+
public void should_reuse_wrapped_simple_statement_for_multipage_query() {
68+
loadBalancingPolicy.customStatementsHandled.set(0);
69+
70+
for (int v = 1; v <= 100; v++)
71+
session().execute(new SimpleStatement(INSERT_MULTIPAGE_QUERY, "key_simple_multipage", v));
72+
73+
SimpleStatement s = new SimpleStatement(SELECT_MULTIPAGE_QUERY, "key_simple_multipage");
74+
s.setFetchSize(1);
75+
76+
CustomStatement customStatement = new CustomStatement(s);
77+
78+
ResultSet rs = session().execute(customStatement);
79+
assertThat(loadBalancingPolicy.customStatementsHandled.get()).isEqualTo(1);
80+
81+
Iterator<Row> it = rs.iterator();
82+
83+
assertThat(it.hasNext()).isTrue();
84+
it.next();
85+
assertThat(rs.getExecutionInfo().getStatement()).isEqualTo(customStatement);
86+
assertThat(loadBalancingPolicy.customStatementsHandled.get()).isEqualTo(1);
87+
88+
assertThat(it.hasNext()).isTrue();
89+
it.next();
90+
assertThat(rs.getExecutionInfo().getStatement()).isEqualTo(customStatement);
91+
92+
assertThat(loadBalancingPolicy.customStatementsHandled.get()).isEqualTo(2);
93+
}
94+
95+
@Test(groups = "short")
96+
public void should_reuse_wrapped_bound_statement_for_multipage_query() {
97+
loadBalancingPolicy.customStatementsHandled.set(0);
98+
99+
for (int v = 1; v <= 100; v++)
100+
session().execute(new SimpleStatement(INSERT_MULTIPAGE_QUERY, "key_prepared_multipage", v));
101+
102+
PreparedStatement ps = session().prepare(SELECT_MULTIPAGE_QUERY);
103+
BoundStatement bs = ps.bind("key_prepared_multipage");
104+
bs.setFetchSize(1);
105+
106+
CustomStatement customStatement = new CustomStatement(bs);
107+
108+
ResultSet rs = session().execute(customStatement);
109+
assertThat(loadBalancingPolicy.customStatementsHandled.get()).isEqualTo(1);
110+
111+
Iterator<Row> it = rs.iterator();
112+
113+
assertThat(it.hasNext()).isTrue();
114+
it.next();
115+
assertThat(rs.getExecutionInfo().getStatement()).isEqualTo(customStatement);
116+
assertThat(loadBalancingPolicy.customStatementsHandled.get()).isEqualTo(1);
117+
118+
assertThat(it.hasNext()).isTrue();
119+
it.next();
120+
assertThat(rs.getExecutionInfo().getStatement()).isEqualTo(customStatement);
121+
122+
assertThat(loadBalancingPolicy.customStatementsHandled.get()).isEqualTo(2);
123+
}
124+
62125
@Test(groups = "short")
63126
public void should_pass_wrapped_statement_to_speculative_execution_policy() {
64127
speculativeExecutionPolicy.customStatementsHandled.set(0);

0 commit comments

Comments
 (0)