From cca2ef9ed1f581f041a9f20694c456baa1265aec Mon Sep 17 00:00:00 2001 From: Bret McGuire Date: Tue, 11 Feb 2020 11:30:22 -0600 Subject: [PATCH] Re-adding support for per-request proxy auth --- .../plugin/model/DseCqlAttributes.scala | 2 ++ .../model/DseCqlAttributesBuilder.scala | 8 +++++ .../plugin/model/DseGraphAttributes.scala | 2 ++ .../model/DseGraphAttributesBuilder.scala | 35 +++++++++---------- .../plugin/request/CqlRequestAction.scala | 13 ++++--- .../plugin/request/GraphRequestAction.scala | 13 ++++--- .../model/CqlStatementBuildersSpec.scala | 7 ++-- 7 files changed, 52 insertions(+), 28 deletions(-) diff --git a/src/main/scala/com/datastax/gatling/plugin/model/DseCqlAttributes.scala b/src/main/scala/com/datastax/gatling/plugin/model/DseCqlAttributes.scala index 9e61b9d..d523c75 100644 --- a/src/main/scala/com/datastax/gatling/plugin/model/DseCqlAttributes.scala +++ b/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, diff --git a/src/main/scala/com/datastax/gatling/plugin/model/DseCqlAttributesBuilder.scala b/src/main/scala/com/datastax/gatling/plugin/model/DseCqlAttributesBuilder.scala index 956c463..b9231a4 100644 --- a/src/main/scala/com/datastax/gatling/plugin/model/DseCqlAttributesBuilder.scala +++ b/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 * diff --git a/src/main/scala/com/datastax/gatling/plugin/model/DseGraphAttributes.scala b/src/main/scala/com/datastax/gatling/plugin/model/DseGraphAttributes.scala index 3000689..4afbe1d 100644 --- a/src/main/scala/com/datastax/gatling/plugin/model/DseGraphAttributes.scala +++ b/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, diff --git a/src/main/scala/com/datastax/gatling/plugin/model/DseGraphAttributesBuilder.scala b/src/main/scala/com/datastax/gatling/plugin/model/DseGraphAttributesBuilder.scala index 7264756..e13f698 100644 --- a/src/main/scala/com/datastax/gatling/plugin/model/DseGraphAttributesBuilder.scala +++ b/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)) } diff --git a/src/main/scala/com/datastax/gatling/plugin/request/CqlRequestAction.scala b/src/main/scala/com/datastax/gatling/plugin/request/CqlRequestAction.scala index 8c13c03..fe98e93 100644 --- a/src/main/scala/com/datastax/gatling/plugin/request/CqlRequestAction.scala +++ b/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)) } } diff --git a/src/main/scala/com/datastax/gatling/plugin/request/GraphRequestAction.scala b/src/main/scala/com/datastax/gatling/plugin/request/GraphRequestAction.scala index 6f606ea..1bfca0e 100644 --- a/src/main/scala/com/datastax/gatling/plugin/request/GraphRequestAction.scala +++ b/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)) } } diff --git a/src/test/scala/com/datastax/gatling/plugin/model/CqlStatementBuildersSpec.scala b/src/test/scala/com/datastax/gatling/plugin/model/CqlStatementBuildersSpec.scala index b4ab03d..b3b6189 100644 --- a/src/test/scala/com/datastax/gatling/plugin/model/CqlStatementBuildersSpec.scala +++ b/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))