Trying to fix as much of the base CQL stuff as possible.

Big chunks of this are based on/copied wholesale from previous work done by Arthur Landim (#17)
datastax · Dec 3, 2019 · 46afb63 · 46afb63
1 parent 4a25388
commit 46afb63
Show file tree

Hide file tree

Showing 2 changed files with 98 additions and 99 deletions.
diff --git a/src/main/scala/com/datastax/gatling/plugin/model/DseCqlStatements.scala b/src/main/scala/com/datastax/gatling/plugin/model/DseCqlStatements.scala
@@ -9,27 +9,25 @@ package com.datastax.gatling.plugin.model
 import java.nio.ByteBuffer
 
 import com.datastax.oss.driver.api.core.cql._
-import com.datastax.oss.driver.api.core.`type`.DataType
 import com.datastax.gatling.plugin.exceptions.DseCqlStatementException
 import com.datastax.gatling.plugin.utils.CqlPreparedStatementUtil
 import io.gatling.commons.validation._
-import io.gatling.core.session.{Session, _}
+import io.gatling.core.session._
 
 import scala.collection.JavaConverters._
-import scala.collection.mutable.ArrayBuffer
 import scala.util.{Try, Failure => TryFailure, Success => TrySuccess}
 
 
-trait DseCqlStatement extends DseStatement[Statement] {
-  def buildFromSession(session: Session): Validation[Statement]
+trait DseCqlStatement[T] extends DseStatement[T] {
+  def buildFromSession(session: Session): Validation[T]
 }
 
 /**
   * Simple CQL Statement from the java driver
   *
   * @param statement the statement to execute
   */
-case class DseCqlSimpleStatement(statement: SimpleStatement) extends DseCqlStatement {
+case class DseCqlSimpleStatement(statement: SimpleStatement) extends DseCqlStatement[SimpleStatement] {
   def buildFromSession(gatlingSession: Session): Validation[SimpleStatement] = {
     statement.success
   }
@@ -41,7 +39,7 @@ case class DseCqlSimpleStatement(statement: SimpleStatement) extends DseCqlState
   * @param preparedStatement the prepared statement on which to bind parameters
   */
 case class DseCqlBoundStatementNamed(cqlTypes: CqlPreparedStatementUtil, preparedStatement: PreparedStatement)
-    extends DseCqlStatement {
+    extends DseCqlStatement[BoundStatement] {
 
   def buildFromSession(gatlingSession: Session): Validation[BoundStatement] =
     bindParams(
@@ -58,7 +56,7 @@ case class DseCqlBoundStatementNamed(cqlTypes: CqlPreparedStatementUtil, prepare
     * @return
     */
   protected def bindParams(gatlingSession: Session, boundStatement: BoundStatement,
-                           queryParams: Map[String, DataType.Name]): BoundStatement = {
+                           queryParams: Map[String, Int]): BoundStatement = {
     queryParams.foreach {
       case (gatlingSessionKey, valType) =>
         cqlTypes.bindParamByName(gatlingSession, boundStatement, valType, gatlingSessionKey)
@@ -75,7 +73,7 @@ case class DseCqlBoundStatementNamed(cqlTypes: CqlPreparedStatementUtil, prepare
   */
 case class DseCqlBoundStatementWithPassedParams(cqlTypes: CqlPreparedStatementUtil,
                                                 preparedStatement: PreparedStatement,
-                                                params: Expression[AnyRef]*) extends DseCqlStatement {
+                                                params: Expression[AnyRef]*) extends DseCqlStatement[BoundStatement] {
 
   def buildFromSession(gatlingSession: Session): Validation[BoundStatement] = {
     val parsedParams: Seq[Validation[AnyRef]] = params.map(param => param(gatlingSession))
@@ -99,7 +97,7 @@ case class DseCqlBoundStatementWithPassedParams(cqlTypes: CqlPreparedStatementUt
   */
 case class DseCqlBoundStatementWithParamList(cqlTypes: CqlPreparedStatementUtil,
                                              preparedStatement: PreparedStatement,
-                                             sessionKeys: Seq[String]) extends DseCqlStatement {
+                                             sessionKeys: Seq[String]) extends DseCqlStatement[BoundStatement] {
 
   /**
     * Apply the Gatling session params to the Prepared statement
@@ -142,16 +140,16 @@ case class DseCqlBoundStatementWithParamList(cqlTypes: CqlPreparedStatementUtil,
   * @param statements CQL Prepared Statements
   */
 case class DseCqlBoundBatchStatement(cqlTypes: CqlPreparedStatementUtil, statements: Seq[PreparedStatement])
-    extends DseCqlStatement {
+    extends DseCqlStatement[BatchStatement] {
 
   def buildFromSession(gatlingSession: Session): Validation[BatchStatement] = {
 
-    val batch = new BatchStatement()
+    val batch = BatchStatement.builder(DefaultBatchType.LOGGED)
 
     statements.foreach(s =>
-      batch.add(bindParams(gatlingSession, s, cqlTypes.getParamsMap(s))))
+      batch.addStatement(bindParams(gatlingSession, s, cqlTypes.getParamsMap(s))))
 
-    batch.success
+    batch.build().success
   }
 
 
@@ -164,7 +162,7 @@ case class DseCqlBoundBatchStatement(cqlTypes: CqlPreparedStatementUtil, stateme
     * @return
     */
   protected def bindParams(gatlingSession: Session, statement: PreparedStatement,
-                           queryParams: Map[String, DataType.Name]): BoundStatement = {
+                           queryParams: Map[String, Int]): BoundStatement = {
 
     val boundStatement = statement.bind()
 
@@ -185,17 +183,18 @@ case class DseCqlBoundBatchStatement(cqlTypes: CqlPreparedStatementUtil, stateme
   * @param statement  SimpleStaten
   * @param payloadRef session variable for custom payload
   */
-case class DseCqlCustomPayloadStatement(statement: SimpleStatement, payloadRef: String) extends DseCqlStatement {
+case class DseCqlCustomPayloadStatement(statement: SimpleStatement, payloadRef: String)
+  extends DseCqlStatement[SimpleStatement] {
 
-  def buildFromSession(gatlingSession: Session): Validation[Statement] = {
+  def buildFromSession(gatlingSession: Session): Validation[SimpleStatement] = {
 
     if (!gatlingSession.contains(payloadRef)) {
       throw new DseCqlStatementException(s"Passed sessionKey: {$payloadRef} does not exist in Session.")
     }
 
     Try {
       val payload = gatlingSession(payloadRef).as[Map[String, ByteBuffer]].asJava
-      statement.setOutgoingPayload(payload)
+      statement.setCustomPayload(payload)
     } match {
       case TrySuccess(stmt) => stmt.success
       case TryFailure(error) => error.getMessage.failure
@@ -210,7 +209,8 @@ case class DseCqlCustomPayloadStatement(statement: SimpleStatement, payloadRef:
   *
   * @param sessionKey the session key which is associated to a PreparedStatement
   */
-case class DseCqlBoundStatementNamedFromSession(cqlTypes: CqlPreparedStatementUtil, sessionKey: String) extends DseCqlStatement {
+case class DseCqlBoundStatementNamedFromSession(cqlTypes: CqlPreparedStatementUtil, sessionKey: String)
+  extends DseCqlStatement[BoundStatement] {
 
   def buildFromSession(gatlingSession: Session): Validation[BoundStatement] = {
     if (!gatlingSession.contains(sessionKey)) {