assume role option

Author: htorrence

build.sbt

@@ -2,7 +2,7 @@ organization := "com.audienceproject"
 
 name := "spark-dynamodb"
 
-version := "0.4.2"
+version := "0.5.0"
 
 description := "Plug-and-play implementation of an Apache Spark custom data source for AWS DynamoDB."
 
@@ -14,6 +14,7 @@ resolvers += "DynamoDBLocal" at "https://s3-us-west-2.amazonaws.com/dynamodb-loc
 
 libraryDependencies += "com.amazonaws" % "aws-java-sdk-dynamodb" % "1.11.466"
 libraryDependencies += "com.amazonaws" % "DynamoDBLocal" % "[1.11,2.0)" % "test" exclude("com.google.guava", "guava")
+libraryDependencies += "com.amazonaws" % "aws-java-sdk" % "1.11.49"
 
 libraryDependencies += "org.apache.spark" %% "spark-sql" % "2.4.0" % "provided"
 libraryDependencies += "com.google.guava" % "guava" % "14.0.1" % "provided"

src/main/scala/com/audienceproject/spark/dynamodb/connector/DynamoConnector.scala

@@ -20,21 +20,23 @@
   */
 package com.audienceproject.spark.dynamodb.connector
 
-import com.amazonaws.auth.DefaultAWSCredentialsProviderChain
+import com.amazonaws.auth.{AWSStaticCredentialsProvider, BasicSessionCredentials, DefaultAWSCredentialsProviderChain}
 import com.amazonaws.auth.profile.ProfileCredentialsProvider
 import com.amazonaws.client.builder.AwsClientBuilder.EndpointConfiguration
 import com.amazonaws.services.dynamodbv2.document.{DynamoDB, ItemCollection, ScanOutcome}
 import com.amazonaws.services.dynamodbv2.{AmazonDynamoDB, AmazonDynamoDBClientBuilder}
+import com.amazonaws.services.securitytoken.AWSSecurityTokenServiceClientBuilder
+import com.amazonaws.services.securitytoken.model.AssumeRoleRequest
 import org.apache.spark.sql.sources.Filter
 
 private[dynamodb] trait DynamoConnector {
 
-    def getDynamoDB(region: Option[String] = None): DynamoDB = {
-        val client: AmazonDynamoDB = getDynamoDBClient(region)
+    def getDynamoDB(region: Option[String] = None, assumedArn: Option[String] = None): DynamoDB = {
+        val client: AmazonDynamoDB = getDynamoDBClient(region, assumedArn)
         new DynamoDB(client)
     }
 
-    def getDynamoDBClient(region: Option[String] = None): AmazonDynamoDB = {
+    def getDynamoDBClient(region: Option[String] = None, assumedArn: Option[String] = None): AmazonDynamoDB = {
         val chosenRegion = region.getOrElse(sys.env.getOrElse("aws.dynamodb.region", "us-east-1"))
         Option(System.getProperty("aws.dynamodb.endpoint")).map(endpoint => {
             val credentials = Option(System.getProperty("aws.profile"))
@@ -44,7 +46,27 @@ private[dynamodb] trait DynamoConnector {
                 .withCredentials(credentials)
                 .withEndpointConfiguration(new EndpointConfiguration(endpoint, chosenRegion))
                 .build()
-        }).getOrElse(AmazonDynamoDBClientBuilder.standard().withRegion(chosenRegion).build())
+        }).getOrElse(
+            assumedArn.map(arn => {
+                val stsClient = AWSSecurityTokenServiceClientBuilder
+                    .standard()
+                    .withCredentials(new DefaultAWSCredentialsProviderChain)
+                    .withRegion(chosenRegion)
+                    .build()
+                val assumeRoleResult = stsClient.assumeRole(
+                    new AssumeRoleRequest()
+                    .withRoleSessionName("DynamoDBAssumed")
+                    .withRoleArn(arn)
+                )
+                val stsCredentials = assumeRoleResult.getCredentials
+                val assumeCreds = new BasicSessionCredentials(
+                    stsCredentials.getAccessKeyId,
+                    stsCredentials.getSecretAccessKey,
+                    stsCredentials.getSessionToken
+                )
+                AmazonDynamoDBClientBuilder.standard().withCredentials(new AWSStaticCredentialsProvider(assumeCreds)).build()
+            }).getOrElse(AmazonDynamoDBClientBuilder.standard().withRegion(chosenRegion).build())
+        )
     }
 
     val keySchema: KeySchema

src/main/scala/com/audienceproject/spark/dynamodb/connector/TableConnector.scala

@@ -41,9 +41,10 @@ private[dynamodb] class TableConnector(tableName: String, totalSegments: Int, pa
     private val consistentRead = parameters.getOrElse("stronglyConsistentReads", "false").toBoolean
     private val filterPushdown = parameters.getOrElse("filterPushdown", "true").toBoolean
     private val region = parameters.get("region")
+    private val assumedArn = parameters.get("assumedArn")
 
     override val (keySchema, readLimit, writeLimit, itemLimit, totalSizeInBytes) = {
-        val table = getDynamoDB(region).getTable(tableName)
+        val table = getDynamoDB(region, assumedArn).getTable(tableName)
         val desc = table.describe()
 
         // Key schema.
@@ -94,7 +95,7 @@ private[dynamodb] class TableConnector(tableName: String, totalSegments: Int, pa
             scanSpec.withExpressionSpec(xspec.buildForScan())
         }
 
-        getDynamoDB(region).getTable(tableName).scan(scanSpec)
+        getDynamoDB(region, assumedArn).getTable(tableName).scan(scanSpec)
     }
 
     override def putItems(schema: StructType, batchSize: Int)(items: Iterator[Row]): Unit = {
@@ -109,7 +110,7 @@ private[dynamodb] class TableConnector(tableName: String, totalSegments: Int, pa
         })
 
         val rateLimiter = RateLimiter.create(writeLimit max 1)
-        val client = getDynamoDB(region)
+        val client = getDynamoDB(region, assumedArn)
 
         // For each batch.
         items.grouped(batchSize).foreach(itemBatch => {
@@ -154,7 +155,7 @@ private[dynamodb] class TableConnector(tableName: String, totalSegments: Int, pa
         })
 
         val rateLimiter = RateLimiter.create(writeLimit max 1)
-        val client = getDynamoDB(region)
+        val client = getDynamoDB(region, assumedArn)
 
         // For each item.
         items.foreach(row => {

src/main/scala/com/audienceproject/spark/dynamodb/connector/TableIndexConnector.scala

@@ -34,9 +34,10 @@ private[dynamodb] class TableIndexConnector(tableName: String, indexName: String
     private val consistentRead = parameters.getOrElse("stronglyConsistentReads", "false").toBoolean
     private val filterPushdown = parameters.getOrElse("filterPushdown", "true").toBoolean
     private val region = parameters.get("region")
+    private val assumedArn = parameters.get("assumedArn")
 
     override val (keySchema, readLimit, itemLimit, totalSizeInBytes) = {
-        val table = getDynamoDB(region).getTable(tableName)
+        val table = getDynamoDB(region, assumedArn).getTable(tableName)
         val indexDesc = table.describe().getGlobalSecondaryIndexes.asScala.find(_.getIndexName == indexName).get
 
         // Key schema.
@@ -81,7 +82,7 @@ private[dynamodb] class TableIndexConnector(tableName: String, indexName: String
             scanSpec.withExpressionSpec(xspec.buildForScan())
         }
 
-        getDynamoDB(region).getTable(tableName).getIndex(indexName).scan(scanSpec)
+        getDynamoDB(region, assumedArn).getTable(tableName).getIndex(indexName).scan(scanSpec)
     }
 
 }