[wip] Upgrade java dse driver to 2.x

Author: landim

build.sbt

@@ -4,11 +4,13 @@ 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.datastax.dse"             %  "dse-java-driver-graph"    % "1.8.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

@@ -7,8 +7,9 @@
 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.gatling.plugin.model.{DseCqlAttributesBuilder, DseCqlStatementBuilder, DseGraphAttributesBuilder, DseGraphStatementBuilder}
 import com.datastax.gatling.plugin.request._
+import com.datastax.oss.driver.api.core.cql.SimpleStatement
 import io.gatling.core.action.builder.ActionBuilder
 
 import scala.language.implicitConversions
@@ -30,13 +31,13 @@ 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()
 }

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.{ResultSet, Row}
 import com.datastax.gatling.plugin.response.CqlResponse
 import io.gatling.commons.validation.{SuccessWrapper, Validation}
 import io.gatling.core.check._

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

@@ -6,14 +6,17 @@
 
 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
 
 /**
   * This class allows to execute checks on either CQL or Graph responses.
@@ -61,15 +64,15 @@ object GenericChecks {
       .toCheckBuilder
 
   val queriedHost =
-    new GenericResponseExtractor[Host](
+    new GenericResponseExtractor[Node](
       "queriedHost",
       r => r.queriedHost())
       .toCheckBuilder
 
   val achievedConsistencyLevel =
     new GenericResponseExtractor[ConsistencyLevel](
       "achievedConsistencyLevel",
-      r => r.achievedConsistencyLevel())
+      r => r.executionInfo().getStatement.getConsistencyLevel())
       .toCheckBuilder
 
   val speculativeExecutionsExtractor =
@@ -79,13 +82,13 @@ object GenericChecks {
       .toCheckBuilder
 
   val pagingState =
-    new GenericResponseExtractor[PagingState](
+    new GenericResponseExtractor[ByteBuffer](
       "pagingState",
       r => r.pagingState())
       .toCheckBuilder
 
   val triedHosts =
-    new GenericResponseExtractor[List[Host]](
+    new GenericResponseExtractor[List[Node]](
       "triedHost",
       r => r.triedHosts())
       .toCheckBuilder
@@ -123,7 +126,6 @@ object GenericChecks {
   val exhausted =
     new GenericResponseExtractor[Boolean](
       "exhausted",
-      r => r.exhausted())
+      r => r.applied())
       .toCheckBuilder
 }
-

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

@@ -6,12 +6,15 @@
 
 package com.datastax.gatling.plugin.checks
 
-import com.datastax.driver.dse.graph._
+import com.datastax.dse.driver.api.core.graph._
+import org.apache.tinkerpop.gremlin.structure.Property
 import com.datastax.gatling.plugin.response.GraphResponse
 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 org.apache.tinkerpop.gremlin.structure._
+import org.apache.tinkerpop.gremlin.process.traversal.Path
 
 import scala.collection.mutable
 
@@ -89,14 +92,14 @@ object GraphChecks {
       .toCheckBuilder
 
   def properties(column: String) =
-    new GraphResponseExtractor[Seq[Property]](
+    new GraphResponseExtractor[Seq[Property[String]]](
       "properties",
       r => r.getProperties(column))
       .toCheckBuilder
 
   def vertexProperties(column: String) =
-    new GraphResponseExtractor[Seq[Property]](
+    new GraphResponseExtractor[Seq[Property[String]]](
       "vertexProperties",
       r => r.getVertexProperties(column))
       .toCheckBuilder
-}
+}

src/main/scala/com/datastax/gatling/plugin/model/DseCqlAttributes.scala

@@ -8,8 +8,9 @@ package com.datastax.gatling.plugin.model
 
 import java.nio.ByteBuffer
 
-import com.datastax.driver.core.policies.RetryPolicy
-import com.datastax.driver.core.{ConsistencyLevel, PagingState, Statement}
+import com.datastax.oss.driver.api.core.retry.RetryPolicy
+import com.datastax.oss.driver.api.core.ConsistencyLevel
+import com.datastax.oss.driver.api.core.cql.Statement
 import com.datastax.gatling.plugin.response.{CqlResponse, DseResponse}
 import io.gatling.core.check.Check
 
@@ -36,8 +37,8 @@ import io.gatling.core.check.Check
   * @param cqlStatements    String version of the CQL statement that is sent
   *
   */
-case class DseCqlAttributes(tag: String,
-                            statement: DseStatement[Statement],
+case class DseCqlAttributes[T](tag: String,
+                            statement: DseCqlStatement[T],
                             cl: Option[ConsistencyLevel] = None,
                             cqlChecks: List[Check[CqlResponse]] = List.empty,
                             genericChecks: List[Check[DseResponse]] = List.empty,
@@ -50,5 +51,5 @@ case class DseCqlAttributes(tag: String,
                             serialCl: Option[ConsistencyLevel] = None,
                             fetchSize: Option[Int] = None,
                             retryPolicy: Option[RetryPolicy] = None,
-                            pagingState: Option[PagingState] = None,
+                            pagingState: Option[ByteBuffer] = None,
                             cqlStatements: Seq[String] = Seq.empty)

src/main/scala/com/datastax/gatling/plugin/model/DseCqlAttributesBuilder.scala

@@ -6,25 +6,26 @@
 
 package com.datastax.gatling.plugin.model
 
-import com.datastax.driver.core.policies.RetryPolicy
-import com.datastax.driver.core.{ConsistencyLevel, PagingState}
+import java.nio.ByteBuffer
+
 import com.datastax.gatling.plugin.checks.{DseCqlCheck, GenericCheck}
 import com.datastax.gatling.plugin.request.CqlRequestActionBuilder
-import io.gatling.core.action.builder.ActionBuilder
+import com.datastax.oss.driver.api.core.ConsistencyLevel
+import com.datastax.oss.driver.api.core.retry.RetryPolicy
 
 
 /**
   * Request Builder for CQL Requests
   *
   * @param attr Addition Attributes
   */
-case class DseCqlAttributesBuilder(attr: DseCqlAttributes) {
+case class DseCqlAttributesBuilder[T](attr: DseCqlAttributes[T]) {
   /**
     * Builds to final action to run
     *
     * @return
     */
-  def build(): CqlRequestActionBuilder = new CqlRequestActionBuilder(attr)
+  def build(): CqlRequestActionBuilder[T] = new CqlRequestActionBuilder(attr)
 
   /**
     * Set Consistency Level
@@ -39,8 +40,8 @@ case class DseCqlAttributesBuilder(attr: DseCqlAttributes) {
     *
     * This permission MUST be granted to the currently logged in user using the CQL statement: `GRANT PROXY.EXECUTE ON
     * ROLE someRole TO alice`.  The user MUST be logged in with
-    * [[com.datastax.driver.dse.auth.DsePlainTextAuthProvider]] or
-    * [[com.datastax.driver.dse.auth.DseGSSAPIAuthProvider]]
+    * [[com.datastax.dse.driver.internal.core.auth.DsePlainTextAuthProvider]] or
+    * [[com.datastax.dse.driver.internal.core.auth.DseGssApiAuthProvider]]
     *
     * @param userOrRole String
     * @return
@@ -83,7 +84,7 @@ case class DseCqlAttributesBuilder(attr: DseCqlAttributes) {
 
 
   /**
-    * Define the [[com.datastax.driver.core.policies.RetryPolicy]] to be used for query
+    * Define the [[com.datastax.oss.driver.api.core.retry.RetryPolicy]] to be used for query
     *
     * @param retryPolicy DataStax drivers retry policy
     * @return
@@ -113,7 +114,7 @@ case class DseCqlAttributesBuilder(attr: DseCqlAttributes) {
     * @param pagingState CQL Paging state
     * @return
     */
-  def withPagingState(pagingState: PagingState) = DseCqlAttributesBuilder(attr.copy(pagingState = Some(pagingState)))
+  def withPagingState(pagingState: ByteBuffer) = DseCqlAttributesBuilder(attr.copy(pagingState = Some(pagingState)))
 
 
   /**

src/main/scala/com/datastax/gatling/plugin/model/DseCqlStatementBuilders.scala

@@ -6,9 +6,9 @@
 
 package com.datastax.gatling.plugin.model
 
-import com.datastax.driver.core.{PreparedStatement, SimpleStatement}
 import com.datastax.gatling.plugin._
 import com.datastax.gatling.plugin.utils.CqlPreparedStatementUtil
+import com.datastax.oss.driver.api.core.cql.{BatchableStatement, BoundStatement, PreparedStatement, SimpleStatement}
 import io.gatling.core.session.Expression
 
 /**
@@ -20,7 +20,7 @@ import io.gatling.core.session.Expression
   *
   * @param tag Name of the CQL query to execute
   */
-case class DseCqlStatementBuilder(tag: String) {
+case class DseCqlStatementBuilder[T](tag: String) {
 
   /**
     * Execute a simple Statement built from a CQL string.
@@ -29,7 +29,7 @@ case class DseCqlStatementBuilder(tag: String) {
     * @return
     */
   @deprecated("Replaced by executeStatement(String)")
-  def executeCql(query: String): DseCqlAttributesBuilder =
+  def executeCql(query: String): DseCqlAttributesBuilder[SimpleStatement] =
     executeStatement(query)
 
   /**
@@ -38,21 +38,21 @@ case class DseCqlStatementBuilder(tag: String) {
     * @param query Simple string query
     * @return
     */
-  def executeStatement(query: String): DseCqlAttributesBuilder =
-    executeStatement(new SimpleStatement(query))
+  def executeStatement(query: String): DseCqlAttributesBuilder[SimpleStatement] =
+    executeStatement(SimpleStatement.newInstance(query))
 
   /**
     * Execute a Simple Statement
     *
     * @param statement SimpleStatement
     * @return
     */
-  def executeStatement(statement: SimpleStatement): DseCqlAttributesBuilder =
+  def executeStatement(statement: SimpleStatement): DseCqlAttributesBuilder[SimpleStatement] =
     DseCqlAttributesBuilder(
       DseCqlAttributes(
         tag,
         DseCqlSimpleStatement(statement),
-        cqlStatements = Seq(statement.getQueryString))
+        cqlStatements = Seq(statement.getQuery))
     )
 
   /**
@@ -90,7 +90,7 @@ case class DseCqlStatementBuilder(tag: String) {
     *
     * @param preparedStatement CQL Prepared statement with named parameters
     */
-  def executeNamed(preparedStatement: PreparedStatement): DseCqlAttributesBuilder =
+  def executeNamed(preparedStatement: PreparedStatement): DseCqlAttributesBuilder[BoundStatement] =
     DsePreparedCqlStatementBuilder(tag, preparedStatement).withSessionParams()
 
   /**
@@ -102,7 +102,7 @@ case class DseCqlStatementBuilder(tag: String) {
     DseCqlAttributes(
       tag,
       DseCqlBoundBatchStatement(CqlPreparedStatementUtil, preparedStatements),
-      cqlStatements = preparedStatements.map(_.getQueryString)
+      cqlStatements = preparedStatements.map(_.getQuery)
     )
   )
 
@@ -113,14 +113,14 @@ case class DseCqlStatementBuilder(tag: String) {
     * @param payloadSessionKey Session key of the payload from session/feed
     * @return
     */
-  def executeCustomPayload(statement: SimpleStatement, payloadSessionKey: String): DseCqlAttributesBuilder =
+  def executeCustomPayload(statement: SimpleStatement, payloadSessionKey: String): DseCqlAttributesBuilder[SimpleStatement] =
     DseCqlAttributesBuilder(
       DseCqlAttributes(
         tag,
         DseCqlCustomPayloadStatement(statement, payloadSessionKey),
-        cqlStatements = Seq(statement.getQueryString)))
+        cqlStatements = Seq(statement.getQuery)))
 
-  def executePreparedFromSession(key: String): DseCqlAttributesBuilder =
+  def executePreparedFromSession(key: String): DseCqlAttributesBuilder[SimpleStatement] =
     DseCqlAttributesBuilder(
       DseCqlAttributes(
         tag,
@@ -140,25 +140,25 @@ case class DsePreparedCqlStatementBuilder(tag: String, prepared: PreparedStateme
     *
     * @return
     */
-  def withSessionParams(): DseCqlAttributesBuilder =
+  def withSessionParams(): DseCqlAttributesBuilder[BoundStatement] =
     DseCqlAttributesBuilder(
       DseCqlAttributes(
         tag,
         DseCqlBoundStatementNamed(CqlPreparedStatementUtil, prepared),
-        cqlStatements = Seq(prepared.getQueryString)))
+        cqlStatements = Seq(prepared.getQuery)))
 
   /**
     * Bind Gatling Session Values to CQL Prepared Statement
     *
     * @param params Gatling Session variables
     * @return
     */
-  def withParams(params: Expression[AnyRef]*): DseCqlAttributesBuilder =
+  def withParams(params: Expression[AnyRef]*): DseCqlAttributesBuilder[BoundStatement] =
     DseCqlAttributesBuilder(
       DseCqlAttributes(
         tag,
         DseCqlBoundStatementWithPassedParams(CqlPreparedStatementUtil, prepared, params: _*),
-        cqlStatements = Seq(prepared.getQueryString))
+        cqlStatements = Seq(prepared.getQuery))
     )
 
   /**
@@ -172,6 +172,6 @@ case class DsePreparedCqlStatementBuilder(tag: String, prepared: PreparedStateme
       DseCqlAttributes(
         tag,
         DseCqlBoundStatementWithParamList(CqlPreparedStatementUtil, prepared, sessionKeys),
-        cqlStatements = Seq(prepared.getQueryString))
+        cqlStatements = Seq(prepared.getQuery))
     )
 }