JAVA-1184: Unwrap StatementWrappers when extracting column definitions.

Author: Aulust

changelog/README.md

@@ -4,6 +4,7 @@
 
 - [bug] JAVA-1179: Request objects should be copied when executed.
 - [improvement] JAVA-1182: Throw error when synchronous call made on I/O thread.
+- [bug] JAVA-1184: Unwrap StatementWrappers when extracting column definitions.
 
 
 ### 2.0.12.1

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

@@ -59,6 +59,9 @@ static ArrayBackedResultSet fromMessage(Responses.Result msg, SessionManager ses
 
  ColumnDefinitions columnDefs;
  if (r.metadata.columns == null) {
+ if (statement instanceof StatementWrapper) {
+ statement = ((StatementWrapper) statement).getWrappedStatement();
+ }
  assert statement instanceof BoundStatement;
  columnDefs = ((BoundStatement) statement).statement.getPreparedId().resultSetMetadata;
  assert columnDefs != null;

driver-core/src/main/java/com/datastax/driver/core/BatchStatement.java

@@ -118,6 +118,9 @@ IdAndValues getIdAndValues() {
  * of statements for a BatchStatement allowed by the underlying protocol).
  */
  public BatchStatement add(Statement statement) {
+ if (statement instanceof StatementWrapper) {
+ statement = ((StatementWrapper) statement).getWrappedStatement();
+ }
 
  // We handle BatchStatement here (rather than in getIdAndValues) as it make it slightly
  // easier to avoid endless loop if the use mistakenly pass a batch that depends on this

driver-core/src/main/java/com/datastax/driver/core/DefaultResultSetFuture.java

@@ -143,7 +143,7 @@ public void run() {
  setException(new DriverInternalError(String.format("Got unexpected %s response from %s", response.type, connection.address)));
  break;
  }
- } catch (RuntimeException e) {
+ } catch (Throwable e) {
  // If we get a bug here, the client will not get it, so better forwarding the error
  setException(new DriverInternalError("Unexpected error while processing response from " + connection.address, e));
  }

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

@@ -16,6 +16,7 @@
 package com.datastax.driver.core;
 
 import com.datastax.driver.core.policies.*;
+import com.datastax.driver.core.utils.CassandraVersion;
 import com.google.common.collect.Lists;
 import org.testng.annotations.Test;
 
@@ -26,9 +27,15 @@
 import static org.assertj.core.api.Assertions.assertThat;
 
 public class StatementWrapperTest extends CCMBridge.PerClassSingleNodeCluster {
+
+ private static final String INSERT_QUERY = "insert into test (k, v) values (?, ?)";
+ private static final String SELECT_QUERY = "select * from test where k = ?";
+
  @Override
  protected Collection<String> getTableDefinitions() {
- return Lists.newArrayList();
+ return Lists.newArrayList(
+ "create table test (k text primary key, v int)"
+ );
  }
 
  CustomLoadBalancingPolicy loadBalancingPolicy = new CustomLoadBalancingPolicy();
@@ -83,6 +90,56 @@ public void should_pass_wrapped_statement_to_retry_policy() {
  assertThat(retryPolicy.customStatementsHandled.get()).isEqualTo(1);
  }
 
+ @CassandraVersion(major=2.0)
+ @Test(groups = "short")
+ public void should_execute_wrapped_simple_statement() {
+ session.execute(new CustomStatement(new SimpleStatement(INSERT_QUERY, "key_simple", 1)));
+
+ ResultSet rs = session.execute(new CustomStatement(new SimpleStatement(SELECT_QUERY, "key_simple")));
+ assertThat(rs.one().getInt("v")).isEqualTo(1);
+ }
+
+ @Test(groups = "short")
+ public void should_execute_wrapped_bound_statement() {
+ PreparedStatement preparedStatement = session.prepare(new SimpleStatement(INSERT_QUERY));
+ session.execute(new CustomStatement(preparedStatement.bind("key_bound", 1)));
+
+ preparedStatement = session.prepare(new SimpleStatement(SELECT_QUERY));
+ ResultSet rs = session.execute(new CustomStatement(preparedStatement.bind("key_bound")));
+ assertThat(rs.one().getInt("v")).isEqualTo(1);
+ }
+
+ @CassandraVersion(major=2.0)
+ @Test(groups = "short")
+ public void should_execute_wrapped_batch_statement() {
+ BatchStatement batchStatement = new BatchStatement();
+ batchStatement.add(new SimpleStatement(INSERT_QUERY, "key_batch", 1));
+
+ session.execute(new CustomStatement(batchStatement));
+
+ ResultSet rs = session.execute(SELECT_QUERY, "key_batch");
+ assertThat(rs.one().getInt("v")).isEqualTo(1);
+ }
+
+ @CassandraVersion(major=2.0)
+ @Test(groups = "short")
+ public void should_add_wrapped_batch_statement_to_batch_statement() {
+ BatchStatement batchStatementForWrapping = new BatchStatement();
+ batchStatementForWrapping.add(new SimpleStatement(INSERT_QUERY, "key1", 1));
+
+ BatchStatement batchStatement = new BatchStatement();
+ batchStatement.add(new CustomStatement(new SimpleStatement(INSERT_QUERY, "key2", 2)));
+ batchStatement.add(new CustomStatement(batchStatementForWrapping));
+
+ session.execute(batchStatement);
+
+ ResultSet rs = session.execute(SELECT_QUERY, "key1");
+ assertThat(rs.one().getInt("v")).isEqualTo(1);
+
+ rs = session.execute(SELECT_QUERY, "key2");
+ assertThat(rs.one().getInt("v")).isEqualTo(2);
+ }
+
  /**
  * A custom wrapper that's just used to mark statements.
  */