Add ExecutionInfo to RequestTracker callbacks

Author: lukasz-antoniak

core/src/main/java/com/datastax/dse/driver/api/core/graph/GraphExecutionInfo.java

@@ -18,8 +18,10 @@
 package com.datastax.dse.driver.api.core.graph;
 
 import com.datastax.oss.driver.api.core.DefaultProtocolVersion;
+import com.datastax.oss.driver.api.core.config.DriverExecutionProfile;
 import com.datastax.oss.driver.api.core.metadata.Node;
 import com.datastax.oss.driver.api.core.specex.SpeculativeExecutionPolicy;
+import edu.umd.cs.findbugs.annotations.Nullable;
 import java.nio.ByteBuffer;
 import java.util.List;
 import java.util.Map;
@@ -37,6 +39,12 @@ public interface GraphExecutionInfo {
   /** The statement that was executed. */
   GraphStatement<?> getStatement();
 
+  /** @return Execution profile applied when executing given request. */
+  @Nullable
+  default DriverExecutionProfile getExecutionProfile() {
+    return null;
+  }
+
   /** The node that was used as a coordinator to successfully complete the query. */
   Node getCoordinator();
 

core/src/main/java/com/datastax/dse/driver/internal/core/cql/continuous/ContinuousRequestHandlerBase.java

@@ -834,7 +834,8 @@ public void onFailure(@NonNull Throwable error) {
     private void processResultResponse(@NonNull Result result, @Nullable Frame frame) {
       assert lock.isHeldByCurrentThread();
       try {
-        ExecutionInfo executionInfo = createExecutionInfo(result, frame);
+        ExecutionInfo executionInfo =
+            createExecutionInfo().withServerResponse(result, frame).build();
         if (result instanceof Rows) {
           DseRowsMetadata rowsMetadata = (DseRowsMetadata) ((Rows) result).getMetadata();
           if (columnDefinitions == null) {
@@ -1459,7 +1460,7 @@ private void trackNodeError(
                 latencyNanos,
                 executionProfile,
                 node,
-                createExecutionInfo(frame),
+                createExecutionInfo().withServerResponse(frame).build(),
                 logPrefix);
       }
     }
@@ -1577,7 +1578,7 @@ private void completeResultSetFuture(
 
           ExecutionInfo executionInfo = null;
           if (pageOrError instanceof AsyncPagingIterable) {
-            executionInfo = ((AsyncPagingIterable) pageOrError).getExecutionInfo();
+            executionInfo = ((AsyncPagingIterable<?, ?>) pageOrError).getExecutionInfo();
           } else if (pageOrError instanceof AsyncGraphResultSet) {
             executionInfo = ((AsyncGraphResultSet) pageOrError).getRequestExecutionInfo();
           }
@@ -1613,34 +1614,13 @@ private void completeResultSetFuture(
     }
 
     @NonNull
-    private ExecutionInfo createExecutionInfo(@NonNull Result result, @Nullable Frame response) {
-      ByteBuffer pagingState =
-          result instanceof Rows ? ((Rows) result).getMetadata().pagingState : null;
-      return new DefaultExecutionInfo(
+    private DefaultExecutionInfo.Builder createExecutionInfo() {
+      return DefaultExecutionInfo.builder(
           statement,
           node,
           startedSpeculativeExecutionsCount.get(),
           executionIndex,
           errors,
-          pagingState,
-          response,
-          true,
-          session,
-          context,
-          executionProfile);
-    }
-
-    @NonNull
-    private ExecutionInfo createExecutionInfo(@Nullable Frame response) {
-      return new DefaultExecutionInfo(
-          statement,
-          node,
-          startedSpeculativeExecutionsCount.get(),
-          executionIndex,
-          errors,
-          null,
-          response,
-          true,
           session,
           context,
           executionProfile);

core/src/main/java/com/datastax/dse/driver/internal/core/graph/GraphExecutionInfoConverter.java

@@ -18,6 +18,7 @@
 package com.datastax.dse.driver.internal.core.graph;
 
 import com.datastax.dse.driver.api.core.graph.GraphStatement;
+import com.datastax.oss.driver.api.core.config.DriverExecutionProfile;
 import com.datastax.oss.driver.api.core.cql.ExecutionInfo;
 import com.datastax.oss.driver.api.core.cql.QueryTrace;
 import com.datastax.oss.driver.api.core.cql.Statement;
@@ -62,6 +63,11 @@ public Statement<?> getStatement() {
         throw new ClassCastException("GraphStatement cannot be cast to Statement");
       }
 
+      @Override
+      public DriverExecutionProfile getExecutionProfile() {
+        return graphExecutionInfo.getExecutionProfile();
+      }
+
       @Nullable
       @Override
       public Node getCoordinator() {
@@ -146,6 +152,11 @@ public GraphStatement<?> getStatement() {
         return (GraphStatement<?>) executionInfo.getRequest();
       }
 
+      @Override
+      public DriverExecutionProfile getExecutionProfile() {
+        return executionInfo.getExecutionProfile();
+      }
+
       @Override
       public Node getCoordinator() {
         return executionInfo.getCoordinator();

core/src/main/java/com/datastax/dse/driver/internal/core/graph/GraphRequestHandler.java

@@ -330,7 +330,18 @@ private void cancelScheduledTasks() {
   private void setFinalResult(
       Result resultMessage, Frame responseFrame, NodeResponseCallback callback) {
     try {
-      ExecutionInfo executionInfo = buildExecutionInfo(callback, responseFrame);
+      ExecutionInfo executionInfo =
+          DefaultExecutionInfo.builder(
+                  callback.statement,
+                  callback.node,
+                  startedSpeculativeExecutionsCount.get(),
+                  callback.execution,
+                  errors,
+                  session,
+                  context,
+                  callback.executionProfile)
+              .withServerResponse(resultMessage, responseFrame)
+              .build();
       DriverExecutionProfile executionProfile =
           Conversions.resolveExecutionProfile(callback.statement, context);
       GraphProtocol subProtocol =
@@ -427,23 +438,6 @@ private void logServerWarnings(GraphStatement<?> statement, List<String> warning
             LOG.warn("Query '{}' generated server side warning(s): {}", statementString, warning));
   }
 
-  private ExecutionInfo buildExecutionInfo(NodeResponseCallback callback, Frame responseFrame) {
-    DriverExecutionProfile executionProfile =
-        Conversions.resolveExecutionProfile(callback.statement, context);
-    return new DefaultExecutionInfo(
-        callback.statement,
-        callback.node,
-        startedSpeculativeExecutionsCount.get(),
-        callback.execution,
-        errors,
-        null,
-        responseFrame,
-        true,
-        session,
-        context,
-        executionProfile);
-  }
-
   @Override
   public void onThrottleFailure(@NonNull RequestThrottlingException error) {
     DriverExecutionProfile executionProfile =
@@ -458,18 +452,16 @@ private void setFinalError(
     DriverExecutionProfile executionProfile =
         Conversions.resolveExecutionProfile(statement, context);
     ExecutionInfo executionInfo =
-        new DefaultExecutionInfo(
-            statement,
-            node,
-            startedSpeculativeExecutionsCount.get(),
-            execution,
-            errors,
-            null,
-            null,
-            true,
-            session,
-            context,
-            executionProfile);
+        DefaultExecutionInfo.builder(
+                statement,
+                node,
+                startedSpeculativeExecutionsCount.get(),
+                execution,
+                errors,
+                session,
+                context,
+                executionProfile)
+            .build();
     if (error instanceof DriverException) {
       ((DriverException) error).setExecutionInfo(executionInfo);
     }

core/src/main/java/com/datastax/oss/driver/api/core/cql/ExecutionInfo.java

@@ -21,6 +21,7 @@
 import com.datastax.oss.driver.api.core.DriverException;
 import com.datastax.oss.driver.api.core.RequestThrottlingException;
 import com.datastax.oss.driver.api.core.config.DefaultDriverOption;
+import com.datastax.oss.driver.api.core.config.DriverExecutionProfile;
 import com.datastax.oss.driver.api.core.detach.AttachmentPoint;
 import com.datastax.oss.driver.api.core.metadata.Node;
 import com.datastax.oss.driver.api.core.retry.RetryDecision;
@@ -66,6 +67,12 @@ default Request getRequest() {
   @Deprecated
   Statement<?> getStatement();
 
+  /** @return Execution profile applied when executing given request. */
+  @Nullable
+  default DriverExecutionProfile getExecutionProfile() {
+    return null;
+  }
+
   /**
    * The node that acted as a coordinator for the query.
    *