apache · adutra · Dec 4, 2019 · Nov 28, 2019 · Dec 2, 2019 · Dec 2, 2019
diff --git a/changelog/README.md b/changelog/README.md
@@ -4,6 +4,7 @@
 
 ### 4.4.0 (in progress)
 
+- [bug] JAVA-2556: Make ExecutionInfo compatible with any Request type
 - [new feature] JAVA-2532: Add BoundStatement ReturnType for insert, update, and delete DAO methods
 - [improvement] JAVA-2107: Add XML formatting plugin
 - [bug] JAVA-2527: Allow AllNodesFailedException to accept more than one error per node

diff --git a/core/src/main/java/com/datastax/oss/driver/api/core/cql/ExecutionInfo.java b/core/src/main/java/com/datastax/oss/driver/api/core/cql/ExecutionInfo.java
@@ -22,6 +22,7 @@
 import com.datastax.oss.driver.api.core.metadata.Node;
 import com.datastax.oss.driver.api.core.retry.RetryDecision;
 import com.datastax.oss.driver.api.core.servererrors.CoordinatorException;
+import com.datastax.oss.driver.api.core.session.Request;
 import com.datastax.oss.driver.api.core.session.Session;
 import com.datastax.oss.driver.api.core.specex.SpeculativeExecutionPolicy;
 import com.datastax.oss.driver.internal.core.util.concurrent.BlockingOperation;
@@ -45,8 +46,20 @@
  */
 public interface ExecutionInfo {
 
-  /** The statement that was executed. */
+  /** @return The {@link Request} that was executed. */
   @NonNull
+  default Request getRequest() {
+    return getStatement();
+  }
+
+  /**
+   * @return The {@link Request} that was executed, if it can be cast to {@link Statement}.
+   * @deprecated Use {@link #getRequest()} instead.
+   * @throws ClassCastException If the request that was executed cannot be cast to {@link
+   *     Statement}.
+   */
+  @NonNull
+  @Deprecated
   Statement<?> getStatement();
 
   /**

diff --git a/core/src/main/java/com/datastax/oss/driver/internal/core/cql/Conversions.java b/core/src/main/java/com/datastax/oss/driver/internal/core/cql/Conversions.java
@@ -323,7 +323,7 @@ public static AsyncResultSet toResultSet(
       InternalDriverContext context) {
     if (result instanceof Rows) {
       Rows rows = (Rows) result;
-      Statement<?> statement = executionInfo.getStatement();
+      Statement<?> statement = (Statement<?>) executionInfo.getRequest();
       ColumnDefinitions columnDefinitions = getResultDefinitions(rows, statement, context);
       return new DefaultAsyncResultSet(
           columnDefinitions, executionInfo, rows.getData(), session, context);

diff --git a/core/src/main/java/com/datastax/oss/driver/internal/core/cql/DefaultAsyncResultSet.java b/core/src/main/java/com/datastax/oss/driver/internal/core/cql/DefaultAsyncResultSet.java
@@ -101,7 +101,7 @@ public CompletionStage<AsyncResultSet> fetchNextPage() throws IllegalStateExcept
       throw new IllegalStateException(
           "No next page. Use #hasMorePages before calling this method to avoid this error.");
     }
-    Statement<?> statement = executionInfo.getStatement();
+    Statement<?> statement = (Statement<?>) executionInfo.getRequest();
     LOG.trace("Fetching next page for {}", statement);
     Statement<?> nextStatement = statement.copy(nextState);
     return session.executeAsync(nextStatement);

diff --git a/core/src/main/java/com/datastax/oss/driver/internal/core/cql/DefaultExecutionInfo.java b/core/src/main/java/com/datastax/oss/driver/internal/core/cql/DefaultExecutionInfo.java
@@ -20,6 +20,7 @@
 import com.datastax.oss.driver.api.core.cql.QueryTrace;
 import com.datastax.oss.driver.api.core.cql.Statement;
 import com.datastax.oss.driver.api.core.metadata.Node;
+import com.datastax.oss.driver.api.core.session.Request;
 import com.datastax.oss.driver.internal.core.context.InternalDriverContext;
 import com.datastax.oss.driver.internal.core.session.DefaultSession;
 import com.datastax.oss.driver.internal.core.util.concurrent.CompletableFutures;
@@ -37,7 +38,7 @@
 @Immutable
 public class DefaultExecutionInfo implements ExecutionInfo {
 
-  private final Statement<?> statement;
+  private final Request request;
   private final Node coordinator;
   private final int speculativeExecutionCount;
   private final int successfulExecutionIndex;
@@ -54,7 +55,7 @@ public class DefaultExecutionInfo implements ExecutionInfo {
   private final DriverExecutionProfile executionProfile;
 
   public DefaultExecutionInfo(
-      Statement<?> statement,
+      Request request,
       Node coordinator,
       int speculativeExecutionCount,
       int successfulExecutionIndex,
@@ -65,7 +66,8 @@ public DefaultExecutionInfo(
       DefaultSession session,
       InternalDriverContext context,
       DriverExecutionProfile executionProfile) {
-    this.statement = statement;
+
+    this.request = request;
     this.coordinator = coordinator;
     this.speculativeExecutionCount = speculativeExecutionCount;
     this.successfulExecutionIndex = successfulExecutionIndex;
@@ -86,8 +88,15 @@ public DefaultExecutionInfo(
 
   @NonNull
   @Override
+  @Deprecated
   public Statement<?> getStatement() {
-    return statement;
+    return (Statement<?>) request;
+  }
+
+  @NonNull
+  @Override
+  public Request getRequest() {
+    return request;
   }
 
   @Nullable

diff --git a/core/src/test/java/com/datastax/oss/driver/internal/core/AsyncPagingIterableWrapperTest.java b/core/src/test/java/com/datastax/oss/driver/internal/core/AsyncPagingIterableWrapperTest.java
@@ -127,7 +127,7 @@ public void should_share_iteration_progress_with_wrapped_result_set() {
 
   private ExecutionInfo mockExecutionInfo() {
     ExecutionInfo executionInfo = mock(ExecutionInfo.class);
-    when(executionInfo.getStatement()).thenAnswer(invocation -> statement);
+    when(executionInfo.getRequest()).thenAnswer(invocation -> statement);
     return executionInfo;
   }
 

diff --git a/core/src/test/java/com/datastax/oss/driver/internal/core/cql/DefaultAsyncResultSetTest.java b/core/src/test/java/com/datastax/oss/driver/internal/core/cql/DefaultAsyncResultSetTest.java
@@ -58,7 +58,7 @@ public class DefaultAsyncResultSetTest {
   public void setup() {
     MockitoAnnotations.initMocks(this);
 
-    when(executionInfo.getStatement()).thenAnswer(invocation -> statement);
+    when(executionInfo.getRequest()).thenAnswer(invocation -> statement);
     when(context.getCodecRegistry()).thenReturn(CodecRegistry.DEFAULT);
     when(context.getProtocolVersion()).thenReturn(DefaultProtocolVersion.DEFAULT);
   }

diff --git a/integration-tests/src/test/java/com/datastax/oss/driver/core/session/ExceptionIT.java b/integration-tests/src/test/java/com/datastax/oss/driver/core/session/ExceptionIT.java
@@ -96,8 +96,7 @@ public void should_expose_execution_info_on_exceptions() {
               assertThat(info).isNotNull();
               assertThat(info.getCoordinator().getEndPoint().resolve())
                   .isEqualTo(SIMULACRON_RULE.cluster().node(1).inetSocketAddress());
-              assertThat(((SimpleStatement) info.getStatement()).getQuery())
-                  .isEqualTo(QUERY_STRING);
+              assertThat(((SimpleStatement) info.getRequest()).getQuery()).isEqualTo(QUERY_STRING);
 
               // specex disabled => the initial execution completed the response
               assertThat(info.getSpeculativeExecutionCount()).isEqualTo(0);