Merge pull request apache#665 from datastax/java1184

JAVA-1184: Unwrap StatementWrappers when extracting column definitions.

Author: olim7t

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/main/java/com/datastax/driver/core/QueryLogger.java

@@ -614,6 +614,9 @@ public void setMaxLoggedParameters(int maxLoggedParameters) {
  */
  @Override
  public void update(Host host, Statement statement, Exception exception, long newLatencyNanos) {
+ if (statement instanceof StatementWrapper)
+ statement = ((StatementWrapper) statement).getWrappedStatement();
+
  long latencyMs = NANOSECONDS.toMillis(newLatencyNanos);
  if (exception == null) {
  maybeLogNormalOrSlowQuery(host, statement, latencyMs);
@@ -745,9 +748,6 @@ protected String parameterValueAsString(ColumnDefinitions.Definition definition,
  }
 
  protected int append(Statement statement, StringBuilder buffer, int remaining) {
- if (statement instanceof StatementWrapper)
- statement = ((StatementWrapper) statement).getWrappedStatement();
-
  if (statement instanceof RegularStatement) {
  remaining = append(((RegularStatement) statement).getQueryString().trim(), buffer, remaining);
  } else if (statement instanceof BoundStatement) {

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

@@ -15,6 +15,7 @@
  */
 package com.datastax.driver.core;
 
+import com.datastax.driver.core.StatementWrapperTest.CustomStatement;
 import com.datastax.driver.core.exceptions.DriverException;
 import com.datastax.driver.core.utils.CassandraVersion;
 import com.google.common.base.Function;
@@ -713,6 +714,33 @@ public void should_log_all_parameters_when_max_unlimited() throws Exception {
  .contains("pk:42");
  }
 
+ @CassandraVersion(major=2.0)
+ @Test(groups = "short")
+ public void should_log_wrapped_bound_statement() throws Exception {
+ // given
+ normal.setLevel(TRACE);
+ queryLogger = QueryLogger.builder(cluster)
+ .withConstantThreshold(Long.MAX_VALUE)
+ .withMaxQueryStringLength(Integer.MAX_VALUE)
+ .build();
+ cluster.register(queryLogger);
+ // when
+ String query = "UPDATE test SET c_text = :param1 WHERE pk = :param2";
+ PreparedStatement ps = session.prepare(query);
+ BoundStatement bs = ps.bind();
+ bs.setString("param1", "foo");
+ bs.setInt("param2", 42);
+ session.execute(new CustomStatement(bs));
+ // then
+ String line = normalAppender.waitAndGet(10000);
+ assertThat(line)
+ .contains("Query completed normally")
+ .contains(ipOfNode(1))
+ .contains(query)
+ .contains("param2:42")
+ .contains("param1:'foo'");
+ }
+
  @Override
  protected List<String> getTableDefinitions() {
  return Lists.newArrayList("CREATE TABLE test (pk int PRIMARY KEY, " + definitions + ")");

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.
  */