[wip] Upgrade java dse driver to 2.x

build.sbt

@@ -4,11 +4,12 @@ val gatlingVersion = "2.3.0"
 
 scalacOptions += "-target:jvm-1.8"
 
-libraryDependencies += "com.datastax.dse"             %  "dse-java-driver-core"     % "1.6.8"
-libraryDependencies += "com.datastax.dse"             %  "dse-java-driver-graph"    % "1.6.8"
+libraryDependencies += "com.datastax.dse"             %  "dse-java-driver-core"     % "2.0.1"
 libraryDependencies += "com.github.nscala-time"       %% "nscala-time"              % "2.18.0"
 libraryDependencies += "com.fasterxml.jackson.module" %% "jackson-module-scala"     % "2.9.1"
 libraryDependencies += "org.hdrhistogram"             %  "HdrHistogram"             % "2.1.10"
+libraryDependencies += "com.google.guava" % "guava" % "23.0"
+
 
 libraryDependencies += "io.gatling.highcharts"        % "gatling-charts-highcharts" % gatlingVersion % Provided
 

src/main/scala/com/datastax/gatling/plugin/DseProtocol.scala

@@ -11,7 +11,7 @@ import java.util.concurrent.atomic.AtomicLong
 
 import akka.Done
 import akka.actor.ActorSystem
-import com.datastax.driver.dse.DseSession
+import com.datastax.dse.driver.api.core.DseSession
 import com.datastax.gatling.plugin.metrics.MetricsLogger
 import com.datastax.gatling.plugin.request.{CqlRequestActionBuilder, GraphRequestActionBuilder}
 import com.datastax.gatling.plugin.utils.GatlingTimingSource
@@ -132,4 +132,3 @@ object DseProtocolBuilder {
 case class DseProtocolBuilder(session: DseSession) {
   def build = DseProtocol(session)
 }
-

src/main/scala/com/datastax/gatling/plugin/Predef.scala

@@ -6,9 +6,11 @@
 
 package com.datastax.gatling.plugin
 
-import com.datastax.gatling.plugin.checks.DseCheckSupport
-import com.datastax.gatling.plugin.model.{DseCqlAttributesBuilder, DseCqlStatementBuilder, DseGraphStatementBuilder, DseGraphAttributesBuilder}
+import com.datastax.dse.driver.api.core.graph.{GraphExecutionInfo, ScriptGraphStatement}
+import com.datastax.gatling.plugin.checks.{CqlChecks, CqlGenericChecks, DseCheckSupport, GraphGenericChecks}
+import com.datastax.gatling.plugin.model.{DseCqlAttributesBuilder, DseCqlStatementBuilder, DseGraphAttributesBuilder, DseGraphStatementBuilder}
 import com.datastax.gatling.plugin.request._
+import com.datastax.oss.driver.api.core.cql.{ExecutionInfo, SimpleStatement}
 import io.gatling.core.action.builder.ActionBuilder
 
 import scala.language.implicitConversions
@@ -30,32 +32,55 @@ trait DsePredefBase extends DseCheckSupport {
   @deprecated("use dseProtocolBuilder instead, will be removed in future versions")
   val cql: DseProtocolBuilder.type = dseProtocolBuilder
 
-  def cql(tag: String): DseCqlStatementBuilder = DseCqlStatementBuilder(tag)
+  def cql(tag: String): DseCqlStatementBuilder[SimpleStatement] = DseCqlStatementBuilder(tag)
 
   def graph(tag: String): DseGraphStatementBuilder = DseGraphStatementBuilder(tag)
 
   implicit def protocolBuilder2DseProtocol(builder: DseProtocolBuilder): DseProtocol = builder.build
 
-  implicit def cqlRequestAttributes2ActionBuilder(builder: DseCqlAttributesBuilder): ActionBuilder = builder.build()
+  implicit def cqlRequestAttributes2ActionBuilder(builder: DseCqlAttributesBuilder[SimpleStatement]): ActionBuilder = builder.build()
 
-  implicit def graphRequestAttributes2ActionBuilder(builder: DseGraphAttributesBuilder): ActionBuilder = builder.build()
+  implicit def graphRequestAttributes2ActionBuilder(builder: DseGraphAttributesBuilder[ScriptGraphStatement]): ActionBuilder = builder.build()
 }
 
 /**
   * DsePredef which should be used for both
   */
 object DsePredef extends DsePredefBase {}
 
-/**
-  * Present for backwards compatibility
-  */
-@deprecated("use DsePredef instead, will be removed in future versions")
-object CqlPredef extends DsePredefBase {}
+object CqlPredef extends DsePredefBase {
+  // start global checks
+  lazy val exhausted = CqlGenericChecks.exhausted
+  lazy val applied = CqlGenericChecks.applied
+  lazy val rowCount = CqlGenericChecks.rowCount
+
+  // execution info and subsets
+  lazy val executionInfo = CqlGenericChecks.executionInfo
+  lazy val achievedCL = CqlGenericChecks.achievedConsistencyLevel
+  lazy val pagingState = CqlGenericChecks.pagingState
+  lazy val queriedHost = CqlGenericChecks.queriedHost
+  lazy val schemaAgreement = CqlGenericChecks.schemaInAgreement
+  lazy val successfulExecutionIndex = CqlGenericChecks.successfulExecutionIndex
+  lazy val triedHosts = CqlGenericChecks.triedHosts
+  lazy val warnings = CqlGenericChecks.warnings
+}
 
-/**
-  * Present for backwards compatibility
-  */
-@deprecated("use DsePredef instead, will be removed in future versions")
-object GraphPredef extends DsePredefBase {}
+object GraphPredef extends DsePredefBase {
+  // start global checks
+  lazy val exhausted = GraphGenericChecks.exhausted
+  lazy val applied = GraphGenericChecks.applied
+  lazy val rowCount = GraphGenericChecks.rowCount
+
+  // execution info and subsets
+  lazy val executionInfo = GraphGenericChecks.executionInfo
+  lazy val achievedCL = GraphGenericChecks.achievedConsistencyLevel
+  lazy val pagingState = GraphGenericChecks.pagingState
+  lazy val queriedHost = GraphGenericChecks.queriedHost
+  lazy val schemaAgreement = GraphGenericChecks.schemaInAgreement
+  lazy val successfulExecutionIndex = GraphGenericChecks.successfulExecutionIndex
+  lazy val triedHosts = GraphGenericChecks.triedHosts
+  lazy val warnings = GraphGenericChecks.warnings
+
+}
 
 

src/main/scala/com/datastax/gatling/plugin/checks/CqlChecks.scala

@@ -6,7 +6,7 @@
 
 package com.datastax.gatling.plugin.checks
 
-import com.datastax.driver.core.{ResultSet, Row}
+import com.datastax.oss.driver.api.core.cql.{ExecutionInfo, ResultSet, Row}
 import com.datastax.gatling.plugin.response.CqlResponse
 import io.gatling.commons.validation.{SuccessWrapper, Validation}
 import io.gatling.core.check._
@@ -21,28 +21,28 @@ import scala.collection.mutable
   *
   * @param wrapped the underlying check
   */
-case class DseCqlCheck(wrapped: Check[CqlResponse]) extends Check[CqlResponse] {
-  override def check(response: CqlResponse, session: Session)(implicit cache: mutable.Map[Any, Any]): Validation[CheckResult] = {
+case class DseCqlCheck(wrapped: Check[CqlResponse[ExecutionInfo]]) extends Check[CqlResponse[ExecutionInfo]] {
+  override def check(response: CqlResponse[ExecutionInfo], session: Session)(implicit cache: mutable.Map[Any, Any]): Validation[CheckResult] = {
     wrapped.check(response, session)
   }
 }
 
-class CqlCheckBuilder[X](extractor: Expression[Extractor[CqlResponse, X]])
-  extends FindCheckBuilder[DseCqlCheck, CqlResponse, CqlResponse, X] {
+class CqlCheckBuilder[X](extractor: Expression[Extractor[CqlResponse[ExecutionInfo], X]])
+  extends FindCheckBuilder[DseCqlCheck, CqlResponse[ExecutionInfo], CqlResponse[ExecutionInfo], X] {
 
-  private val cqlResponseExtender: Extender[DseCqlCheck, CqlResponse] =
+  private val cqlResponseExtender: Extender[DseCqlCheck, CqlResponse[ExecutionInfo]] =
     wrapped => DseCqlCheck(wrapped)
 
-  def find: ValidatorCheckBuilder[DseCqlCheck, CqlResponse, CqlResponse, X] = {
+  def find: ValidatorCheckBuilder[DseCqlCheck, CqlResponse[ExecutionInfo], CqlResponse[ExecutionInfo], X] = {
     ValidatorCheckBuilder(cqlResponseExtender, x => x.success, extractor)
   }
 }
 
 private class CqlResponseExtractor[X](val name: String,
-                                      val extractor: CqlResponse => X)
-  extends Extractor[CqlResponse, X] with SingleArity {
+                                      val extractor: CqlResponse[ExecutionInfo] => X)
+  extends Extractor[CqlResponse[ExecutionInfo], X] with SingleArity {
 
-  override def apply(response: CqlResponse): Validation[Option[X]] = {
+  override def apply(response: CqlResponse[ExecutionInfo]): Validation[Option[X]] = {
     Some(extractor.apply(response)).success
   }
 
@@ -51,22 +51,22 @@ private class CqlResponseExtractor[X](val name: String,
   }
 }
 
-private abstract class ColumnValueExtractor[X] extends CriterionExtractor[CqlResponse, Any, X] {
+private abstract class ColumnValueExtractor[X] extends CriterionExtractor[CqlResponse[ExecutionInfo], Any, X] {
   val criterionName = "columnValue"
 }
 
 private class SingleColumnValueExtractor(val criterion: String, val occurrence: Int) extends ColumnValueExtractor[Any] with FindArity {
-  def extract(response: CqlResponse): Validation[Option[Any]] =
+  def extract(response: CqlResponse[ExecutionInfo]): Validation[Option[Any]] =
     response.getCqlResultColumnValues(criterion).lift(occurrence).success
 }
 
 private class MultipleColumnValueExtractor(val criterion: String) extends ColumnValueExtractor[Seq[Any]] with FindAllArity {
-  def extract(response: CqlResponse): Validation[Option[Seq[Any]]] =
+  def extract(response: CqlResponse[ExecutionInfo]): Validation[Option[Seq[Any]]] =
     response.getCqlResultColumnValues(criterion).liftSeqOption.success
 }
 
 private class CountColumnValueExtractor(val criterion: String) extends ColumnValueExtractor[Int] with CountArity {
-  def extract(response: CqlResponse): Validation[Option[Int]] =
+  def extract(response: CqlResponse[ExecutionInfo]): Validation[Option[Int]] =
     response.getCqlResultColumnValues(criterion).liftSeqOption.map(_.size).success
 }
 
@@ -90,8 +90,8 @@ object CqlChecks {
       .toCheckBuilder
 
   def columnValue(columnName: Expression[String]) = {
-    val cqlResponseExtender: Extender[DseCqlCheck, CqlResponse] = wrapped => DseCqlCheck(wrapped)
-    new DefaultMultipleFindCheckBuilder[DseCqlCheck, CqlResponse, CqlResponse, Any](cqlResponseExtender, x => x.success) {
+    val cqlResponseExtender: Extender[DseCqlCheck, CqlResponse[ExecutionInfo]] = wrapped => DseCqlCheck(wrapped)
+    new DefaultMultipleFindCheckBuilder[DseCqlCheck, CqlResponse[ExecutionInfo], CqlResponse[ExecutionInfo], Any](cqlResponseExtender, x => x.success) {
       def findExtractor(occurrence: Int) = columnName.map(new SingleColumnValueExtractor(_, occurrence))
       def findAllExtractor = columnName.map(new MultipleColumnValueExtractor(_))
       def countExtractor = columnName.map(new CountColumnValueExtractor(_))

src/main/scala/com/datastax/gatling/plugin/checks/DseCheckSupport.scala

@@ -9,22 +9,6 @@ package com.datastax.gatling.plugin.checks
 import io.gatling.core.session.ExpressionSuccessWrapper
 
 trait DseCheckSupport {
-
-  // start global checks
-  lazy val exhausted = GenericChecks.exhausted
-  lazy val applied = GenericChecks.applied
-  lazy val rowCount = GenericChecks.rowCount
-
-  // execution info and subsets
-  lazy val executionInfo = GenericChecks.executionInfo
-  lazy val achievedCL = GenericChecks.achievedConsistencyLevel
-  lazy val pagingState = GenericChecks.pagingState
-  lazy val queriedHost = GenericChecks.queriedHost
-  lazy val schemaAgreement = GenericChecks.schemaInAgreement
-  lazy val successfulExecutionIndex = GenericChecks.successfulExecutionIndex
-  lazy val triedHosts = GenericChecks.triedHosts
-  lazy val warnings = GenericChecks.warnings
-
   // start CQL only checks
   lazy val resultSet = CqlChecks.resultSet
   lazy val allRows = CqlChecks.allRows

src/main/scala/com/datastax/gatling/plugin/checks/GenericChecks.scala

@@ -6,14 +6,19 @@
 
 package com.datastax.gatling.plugin.checks
 
-import com.datastax.driver.core._
+import com.datastax.oss.driver.api.core.cql._
+import com.datastax.oss.driver.api.core._
+import com.datastax.oss.driver.api.core.metadata._
 import com.datastax.gatling.plugin.response.DseResponse
 import io.gatling.commons.validation.{SuccessWrapper, Validation}
 import io.gatling.core.check.extractor.{Extractor, SingleArity}
 import io.gatling.core.check._
 import io.gatling.core.session.{Expression, ExpressionSuccessWrapper, Session}
 
 import scala.collection.mutable
+import java.nio.ByteBuffer
+
+import com.datastax.dse.driver.api.core.graph.GraphExecutionInfo
 
 /**
   * This class allows to execute checks on either CQL or Graph responses.
@@ -23,107 +28,109 @@ import scala.collection.mutable
   * would make it possible to execute CQL checks on Graph responses, or
   * vice-versa.
   */
-case class GenericCheck(wrapped: Check[DseResponse]) extends Check[DseResponse] {
-  override def check(response: DseResponse, session: Session)(implicit cache: mutable.Map[Any, Any]): Validation[CheckResult] = {
+case class GenericCheck[E](wrapped: Check[DseResponse[E]]) extends Check[DseResponse[E]] {
+  override def check(response: DseResponse[E], session: Session)(implicit cache: mutable.Map[Any, Any]): Validation[CheckResult] = {
     wrapped.check(response, session)
   }
 }
 
-class GenericCheckBuilder[X](extractor: Expression[Extractor[DseResponse, X]])
-  extends FindCheckBuilder[GenericCheck, DseResponse, DseResponse, X] {
+class GenericCheckBuilder[X, E](extractor: Expression[Extractor[DseResponse[E], X]])
+  extends FindCheckBuilder[GenericCheck[E], DseResponse[E], DseResponse[E], X] {
 
-  private val dseResponseExtender: Extender[GenericCheck, DseResponse] =
+  private val dseResponseExtender: Extender[GenericCheck[E], DseResponse[E]] =
     wrapped => GenericCheck(wrapped)
 
-  def find: ValidatorCheckBuilder[GenericCheck, DseResponse, DseResponse, X] = {
+  def find: ValidatorCheckBuilder[GenericCheck[E], DseResponse[E], DseResponse[E], X] = {
     ValidatorCheckBuilder(dseResponseExtender, x => x.success, extractor)
   }
 }
 
-private class GenericResponseExtractor[X](val name: String,
-                                          val extractor: DseResponse => X)
-  extends Extractor[DseResponse, X] with SingleArity {
+private class GenericResponseExtractor[X, E](val name: String,
+                                          val extractor: DseResponse[E] => X)
+  extends Extractor[DseResponse[E], X] with SingleArity {
 
-  override def apply(response: DseResponse): Validation[Option[X]] = {
+  override def apply(response: DseResponse[E]): Validation[Option[X]] = {
     Some(extractor.apply(response)).success
   }
 
-  def toCheckBuilder: GenericCheckBuilder[X] = {
-    new GenericCheckBuilder[X](this.expressionSuccess)
+  def toCheckBuilder: GenericCheckBuilder[X, E] = {
+    new GenericCheckBuilder[X, E](this.expressionSuccess)
   }
 }
 
-object GenericChecks {
+class GenericChecks[T] {
   val executionInfo =
-    new GenericResponseExtractor[ExecutionInfo](
+    new GenericResponseExtractor[T, T](
       "executionInfo",
       r => r.executionInfo())
       .toCheckBuilder
 
   val queriedHost =
-    new GenericResponseExtractor[Host](
+    new GenericResponseExtractor[Node, T](
       "queriedHost",
       r => r.queriedHost())
       .toCheckBuilder
 
   val achievedConsistencyLevel =
-    new GenericResponseExtractor[ConsistencyLevel](
+    new GenericResponseExtractor[ConsistencyLevel, T](
       "achievedConsistencyLevel",
-      r => r.achievedConsistencyLevel())
+      r => r.getConsistencyLevel())
       .toCheckBuilder
 
   val speculativeExecutionsExtractor =
-    new GenericResponseExtractor[Int](
+    new GenericResponseExtractor[Int, T](
       "speculativeExecutions",
       r => r.speculativeExecutions())
       .toCheckBuilder
 
   val pagingState =
-    new GenericResponseExtractor[PagingState](
+    new GenericResponseExtractor[ByteBuffer, T](
       "pagingState",
       r => r.pagingState())
       .toCheckBuilder
 
   val triedHosts =
-    new GenericResponseExtractor[List[Host]](
+    new GenericResponseExtractor[List[Node], T](
       "triedHost",
       r => r.triedHosts())
       .toCheckBuilder
 
   val warnings =
-    new GenericResponseExtractor[List[String]](
+    new GenericResponseExtractor[List[String], T](
       "warnings",
       r => r.warnings())
       .toCheckBuilder
 
   val successfulExecutionIndex =
-    new GenericResponseExtractor[Int](
+    new GenericResponseExtractor[Int, T](
       "successfulExecutionIndex",
       r => r.successFullExecutionIndex())
       .toCheckBuilder
 
   val schemaInAgreement =
-    new GenericResponseExtractor[Boolean](
+    new GenericResponseExtractor[Boolean, T](
       "schemaInAgreement",
       r => r.schemaInAgreement())
       .toCheckBuilder
 
   val rowCount =
-    new GenericResponseExtractor[Int](
+    new GenericResponseExtractor[Int, T](
       "rowCount",
       r => r.rowCount())
       .toCheckBuilder
 
   val applied =
-    new GenericResponseExtractor[Boolean](
+    new GenericResponseExtractor[Boolean, T](
       "applied",
       r => r.applied())
       .toCheckBuilder
 
   val exhausted =
-    new GenericResponseExtractor[Boolean](
+    new GenericResponseExtractor[Boolean, T](
       "exhausted",
-      r => r.exhausted())
+      r => r.applied())
       .toCheckBuilder
 }
 
+object CqlGenericChecks extends GenericChecks[ExecutionInfo];
+object GraphGenericChecks extends GenericChecks[GraphExecutionInfo];