JAVA-1068: Unwrap StatementWrappers when hashing the paging state.

djatnieks · olim7t · commit 63cbebddb15f · 2016-02-25T15:19:22.000+01:00
diff --git a/changelog/README.md b/changelog/README.md
@@ -18,6 +18,7 @@
 - [improvement] JAVA-979: Update javadoc for RegularStatement toString() and getQueryString() to indicate that consistency level and other parameters are not maintained in the query string.
 - [improvement] JAVA-1038: Fetch node info by rpc_address if its broadcast_address is not in system.peers.
 - [improvement] JAVA-974: Validate accessor parameter types against bound statement.
+- [bug] JAVA-1068: Unwrap StatementWrappers when hashing the paging state.
 
 Merged from 2.0 branch:
 
diff --git a/driver-core/src/main/java/com/datastax/driver/core/PagingState.java b/driver-core/src/main/java/com/datastax/driver/core/PagingState.java
@@ -74,6 +74,9 @@ private byte[] hash(Statement statement, ProtocolVersion protocolVersion) {
         byte[] digest;
         ByteBuffer[] values;
         MessageDigest md;
+        if (statement instanceof StatementWrapper) {
+            statement = ((StatementWrapper) statement).getWrappedStatement();
+        }
         assert !(statement instanceof BatchStatement);
         try {
             md = MessageDigest.getInstance("MD5");
diff --git a/driver-core/src/test/java/com/datastax/driver/core/PagingStateTest.java b/driver-core/src/test/java/com/datastax/driver/core/PagingStateTest.java
@@ -112,6 +112,27 @@ public void should_fail_if_paging_state_altered() {
         }
     }
 
+    /**
+     * Validates that {@link PagingState} can be reused with a wrapped Statement.
+     *
+     * @test_category paging
+     * @expected_result {@link ResultSet} from the query with the provided {@link PagingState} starts from the
+     * subsequent row from the first query.
+     */
+    @Test(groups = "short")
+    public void should_use_state_with_wrapped_statement() {
+        Statement st = new TestWrapper(new SimpleStatement(String.format("SELECT v FROM test WHERE k='%s'", KEY)));
+        ResultSet result = session().execute(st.setFetchSize(20));
+        int pageSize = result.getAvailableWithoutFetching();
+        String savedPagingStateString = result.getExecutionInfo().getPagingState().toString();
+
+        st = new TestWrapper(new SimpleStatement(String.format("SELECT v FROM test WHERE k='%s'", KEY)));
+        result = session().execute(st.setFetchSize(20).setPagingState(PagingState.fromString(savedPagingStateString)));
+
+        //We have the result starting from the next page we stopped
+        assertThat(result.one().getInt("v")).isEqualTo(pageSize);
+    }
+
     /**
      * Validates that {@link PagingState} can be reused with the same {@link BoundStatement}.
      *
@@ -224,5 +245,10 @@ public void should_complete_when_using_unsafe_paging_state() {
         assertThat(result.one().getInt("v")).isEqualTo(pageSize);
     }
 
+    static class TestWrapper extends StatementWrapper {
+        TestWrapper(Statement wrapped) {
+            super(wrapped);
+        }
+    }
 }
 
