Merge branch 'htorrence-assume_role_option'

Author: jacobfi

README.md

@@ -49,6 +49,7 @@ Spark is used in the library as a "provided" dependency, which means Spark has t
 ## Parameters
 The following parameters can be set as options on the Spark reader and writer object before loading/saving.
 - `region` sets the region where the dynamodb table. Default is environment specific.
+- `assumedArn` sets an IAM role to assume. This allows for access to a DynamoDB in a different account than the spark cluster. Defaults to the standard role configuration.
 
 
 The following parameters can be set as options on the Spark reader object before loading.

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,31 +20,63 @@
   */
 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 = {
+    /**
+      * Get credentials from a passed in arn or from profile or return the default credential provider
+      * */
+    private def getCredentials(chosenRegion: String, assumedArn: Option[String]) = {
+        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
+            )
+            new AWSStaticCredentialsProvider(assumeCreds)
+        }).orElse(Option(System.getProperty("aws.profile")).map(new ProfileCredentialsProvider(_)))
+          .getOrElse(new DefaultAWSCredentialsProviderChain)
+    }
+
+    def getDynamoDBClient(region: Option[String] = None, assumedArn: Option[String] = None): AmazonDynamoDB = {
         val chosenRegion = region.getOrElse(sys.env.getOrElse("aws.dynamodb.region", "us-east-1"))
+        val credentials = getCredentials(chosenRegion, assumedArn)
+
         Option(System.getProperty("aws.dynamodb.endpoint")).map(endpoint => {
-            val credentials = Option(System.getProperty("aws.profile"))
-                .map(new ProfileCredentialsProvider(_))
-                .getOrElse(new DefaultAWSCredentialsProviderChain)
             AmazonDynamoDBClientBuilder.standard()
                 .withCredentials(credentials)
                 .withEndpointConfiguration(new EndpointConfiguration(endpoint, chosenRegion))
                 .build()
-        }).getOrElse(AmazonDynamoDBClientBuilder.standard().withRegion(chosenRegion).build())
+        }).getOrElse(
+            AmazonDynamoDBClientBuilder.standard()
+                .withCredentials(credentials)
+                .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)
     }
 
 }