Re-adding support for per-request proxy auth

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

@@ -27,6 +27,7 @@ import com.datastax.oss.driver.api.core.time.TimestampGenerator
   * @param cqlChecks        Data-level checks to be run after response is returned
   * @param idempotent       Set request to be idempotent i.e. whether it can be applied multiple times
   * @param node             Set the node that should handle this query
+  * @param userOrRole       Set the user/role for this query if proxy authentication is used
   * @param customPayload    Custom payload for this request
   * @param enableTrace      Whether tracing should be enabled
   * @param pageSize         Set pageSize (formerly known as fetchSize)
@@ -50,6 +51,7 @@ case class DseCqlAttributes[T <: Statement[T], B <: StatementBuilder[B,T]]
    cl: Option[ConsistencyLevel] = None,
    idempotent: Option[Boolean] = None,
    node: Option[Node] = None,
+   userOrRole: Option[String] = None,
    /* CQL-specific attributes */
    customPayload: Option[Map[String, ByteBuffer]] = None,
    enableTrace: Option[Boolean] = None,

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

@@ -78,6 +78,14 @@ case class DseCqlAttributesBuilder[T <: Statement[T], B <: StatementBuilder[B,T]
   def withNode(node: Node):DseCqlAttributesBuilder[T, B] =
     DseCqlAttributesBuilder(attr.copy(node = Some(node)))
 
+  /**
+    * Set the user or role to use for proxy auth
+    * @param userOrRole String
+    * @return
+    */
+  def executeAs(userOrRole: String):DseCqlAttributesBuilder[T, B] =
+    DseCqlAttributesBuilder(attr.copy(userOrRole = Some(userOrRole)))
+
   /**
     * Enable CQL Tracing on the query
     *

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

@@ -24,6 +24,7 @@ import com.datastax.oss.driver.api.core.metadata.Node
   * @param graphChecks            Data-level checks to be run after response is returned
   * @param idempotent             Set request to be idempotent i.e. whether it can be applied multiple times
   * @param node                   Set the node that should handle this query
+  * @param userOrRole             Set the user/role for this query if proxy authentication is used
   * @param graphName              Name of the graph to use if different from the one used when connecting
   * @param readCL                 Consistency level to use for the read part of the query
   * @param subProtocol            Name of the graph protocol to use for encoding/decoding
@@ -40,6 +41,7 @@ case class DseGraphAttributes[T <: GraphStatement[T], B <: GraphStatementBuilder
    cl: Option[ConsistencyLevel] = None,
    idempotent: Option[Boolean] = None,
    node: Option[Node] = None,
+   userOrRole: Option[String] = None,
    /* Graph-specific attributes */
    graphName: Option[String] = None,
    readCL: Option[ConsistencyLevel] = None,

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

@@ -44,6 +44,23 @@ case class DseGraphAttributesBuilder[T <: GraphStatement[T], B <: GraphStatement
   def withIdempotency():DseGraphAttributesBuilder[T, B] =
     DseGraphAttributesBuilder(attr.copy(idempotent = Some(true)))
 
+  /**
+    * Set the node that should handle this query
+    * @param node Node
+    * @return
+    */
+  def withNode(node: Node):DseGraphAttributesBuilder[T, B] =
+    DseGraphAttributesBuilder(attr.copy(node = Some(node)))
+
+  /**
+    * Set the user or role to use for proxy auth
+    * @param userOrRole String
+    * @return
+    */
+  def executeAs(userOrRole: String):DseGraphAttributesBuilder[T, B] =
+    DseGraphAttributesBuilder(attr.copy(userOrRole = Some(userOrRole)))
+
+
   /**
     * Sets the graph name to use
     *
@@ -53,14 +70,6 @@ case class DseGraphAttributesBuilder[T <: GraphStatement[T], B <: GraphStatement
   def withName(name: String):DseGraphAttributesBuilder[T, B] =
     DseGraphAttributesBuilder(attr.copy(graphName = Some(name)))
 
-  /**
-    * Set the node that should handle this query
-    * @param node Node
-    * @return
-    */
-  def withNode(node: Node):DseGraphAttributesBuilder[T, B] =
-    DseGraphAttributesBuilder(attr.copy(node = Some(node)))
-
   /**
     * Define [[ConsistencyLevel]] to be used for read queries
     *
@@ -115,16 +124,6 @@ case class DseGraphAttributesBuilder[T <: GraphStatement[T], B <: GraphStatement
   def withWriteConsistency(writeCL: ConsistencyLevel):DseGraphAttributesBuilder[T, B] =
     DseGraphAttributesBuilder(attr.copy(writeCL = Some(writeCL)))
 
-  /**
-    * Backwards compatibility to set consistencyLevel
-    *
-    * @see [[DseGraphAttributesBuilder.withConsistencyLevel]]
-    * @param level Consistency Level to use
-    * @return
-    */
-  @deprecated("use withConsistencyLevel() instead, will be removed in future version")
-  def consistencyLevel(level: ConsistencyLevel):DseGraphAttributesBuilder[T, B] = withConsistencyLevel(level)
-
   def check(check: DseGraphCheck):DseGraphAttributesBuilder[T, B] =
     DseGraphAttributesBuilder(attr.copy(graphChecks = check :: attr.graphChecks))
 }

src/main/scala/com/datastax/gatling/plugin/request/CqlRequestAction.scala

@@ -12,6 +12,7 @@ import java.util.concurrent.ExecutorService
 import java.util.concurrent.TimeUnit.MICROSECONDS
 
 import akka.actor.ActorSystem
+import com.datastax.dse.driver.api.core.auth.ProxyAuthentication
 import com.datastax.gatling.plugin.DseProtocol
 import com.datastax.gatling.plugin.metrics.MetricsLogger
 import com.datastax.gatling.plugin.model.DseCqlAttributes
@@ -91,7 +92,11 @@ class CqlRequestAction[T <: Statement[T], B <: StatementBuilder[B,T]](val name:
 
   private def handleSuccess(session: Session, responseTimeBuilder: ResponseTimeBuilder)(builder:B): Unit = {
 
-    val stmt:T = buildStatement(builder)
+    val baseStmt:T = buildStatement(builder)
+    val stmt:T =
+      dseAttributes.userOrRole
+      .map(ProxyAuthentication.executeAs(_,baseStmt))
+      .getOrElse(baseStmt)
     val responseHandler =
       new CqlResponseHandler[T, B](
         next,
@@ -129,8 +134,8 @@ class CqlRequestAction[T <: Statement[T], B <: StatementBuilder[B,T]](val name:
       ThroughputVerifier.checkForGatlingOverloading(session, gatlingTimingSource)
       GatlingResponseTime.startedByGatling(session, gatlingTimingSource)
     }
-    val stmt = safely()(dseAttributes.statement.buildFromSession(session))
-    stmt.onFailure(handleFailure(session,responseTimeBuilder))
-    stmt.onSuccess(handleSuccess(session,responseTimeBuilder))
+    val stmtBuilder = safely()(dseAttributes.statement.buildFromSession(session))
+    stmtBuilder.onFailure(handleFailure(session,responseTimeBuilder))
+    stmtBuilder.onSuccess(handleSuccess(session,responseTimeBuilder))
   }
 }

src/main/scala/com/datastax/gatling/plugin/request/GraphRequestAction.scala

@@ -12,6 +12,7 @@ import java.util.concurrent.ExecutorService
 import java.util.concurrent.TimeUnit.MICROSECONDS
 
 import akka.actor.ActorSystem
+import com.datastax.dse.driver.api.core.auth.ProxyAuthentication
 import com.datastax.dse.driver.api.core.graph.{GraphStatement, GraphStatementBuilderBase}
 import com.datastax.gatling.plugin.DseProtocol
 import com.datastax.gatling.plugin.metrics.MetricsLogger
@@ -82,7 +83,11 @@ class GraphRequestAction[T <: GraphStatement[T], B <: GraphStatementBuilderBase[
 
   private def handleSuccess(session: Session, responseTimeBuilder: ResponseTimeBuilder)(builder:B): Unit = {
 
-    val stmt:T = buildStatement(builder)
+    val baseStmt:T = buildStatement(builder)
+    val stmt:T =
+      dseAttributes.userOrRole
+        .map(ProxyAuthentication.executeAs(_,baseStmt))
+        .getOrElse(baseStmt)
     val responseHandler =
       new GraphResponseHandler[T, B](
         next,
@@ -121,8 +126,8 @@ class GraphRequestAction[T <: GraphStatement[T], B <: GraphStatementBuilderBase[
       ThroughputVerifier.checkForGatlingOverloading(session, gatlingTimingSource)
       GatlingResponseTime.startedByGatling(session, gatlingTimingSource)
     }
-    val stmt = dseAttributes.statement.buildFromSession(session)
-    stmt.onFailure(handleFailure(session, responseTimeBuilder))
-    stmt.onSuccess(handleSuccess(session, responseTimeBuilder))
+    val stmtBuilder = dseAttributes.statement.buildFromSession(session)
+    stmtBuilder.onFailure(handleFailure(session, responseTimeBuilder))
+    stmtBuilder.onSuccess(handleSuccess(session, responseTimeBuilder))
   }
 }

src/test/scala/com/datastax/gatling/plugin/model/CqlStatementBuildersSpec.scala

@@ -19,7 +19,7 @@ class CqlStatementBuildersSpec extends FlatSpec with Matchers with EasyMockSugar
 
   it should "build statements from a CQL String" in {
     val statementAttributes: DseCqlAttributes[SimpleStatement,SimpleStatementBuilder] = cql("the-tag")
-      .executeCql("SELECT foo FROM bar.baz LIMIT 1")
+      .executeStatement("SELECT foo FROM bar.baz LIMIT 1")
       .build()
       .dseAttributes
     val statement = statementAttributes.statement
@@ -31,6 +31,7 @@ class CqlStatementBuildersSpec extends FlatSpec with Matchers with EasyMockSugar
 
   it should "forward all attributs to DseCqlAttributes" in {
     val node = mock[Node]
+    val userOrRole = "userOrRole"
     val customPayloadKey = "key"
     val customPayloadVal = mock[ByteBuffer]
     val pagingState = mock[ByteBuffer]
@@ -41,11 +42,12 @@ class CqlStatementBuildersSpec extends FlatSpec with Matchers with EasyMockSugar
     val timeout = Duration.ofHours(1)
     val cqlCheck = CqlChecks.resultSet.find.is(mock[AsyncResultSet].expressionSuccess).build
     val statementAttributes: DseCqlAttributes[_,_] = cql("the-session-tag")
-      .executeCql("FOO")
+      .executeStatement("FOO")
       .withConsistencyLevel(EACH_QUORUM)
       .addCustomPayload(customPayloadKey, customPayloadVal)
       .withIdempotency()
       .withNode(node)
+      .executeAs(userOrRole)
       .withTracingEnabled()
       .withPageSize(3)
       .withPagingState(pagingState)
@@ -64,6 +66,7 @@ class CqlStatementBuildersSpec extends FlatSpec with Matchers with EasyMockSugar
     statementAttributes.customPayload should be(Some(Map(customPayloadKey -> customPayloadVal)))
     statementAttributes.idempotent should be(Some(true))
     statementAttributes.node should be(Some(node))
+    statementAttributes.userOrRole should be(Some(userOrRole))
     statementAttributes.enableTrace should be(Some(true))
     statementAttributes.pageSize should be(Some(3))
     statementAttributes.pagingState should be(Some(pagingState))