Merge PR htorrence:assume_role_option

Author: jacobfi

README.md

@@ -49,7 +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.
+- `roleArn` 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.5.0"
+version := "0.4.3"
 
 description := "Plug-and-play implementation of an Apache Spark custom data source for AWS DynamoDB."
 
@@ -12,9 +12,9 @@ crossScalaVersions := Seq("2.11.12", "2.12.7")
 
 resolvers += "DynamoDBLocal" at "https://s3-us-west-2.amazonaws.com/dynamodb-local/release"
 
-libraryDependencies += "com.amazonaws" % "aws-java-sdk-dynamodb" % "1.11.466"
+libraryDependencies += "com.amazonaws" % "aws-java-sdk-sts" % "1.11.571"
+libraryDependencies += "com.amazonaws" % "aws-java-sdk-dynamodb" % "1.11.571"
 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,8 +20,8 @@
   */
 package com.audienceproject.spark.dynamodb.connector
 
-import com.amazonaws.auth.{AWSStaticCredentialsProvider, BasicSessionCredentials, DefaultAWSCredentialsProviderChain}
 import com.amazonaws.auth.profile.ProfileCredentialsProvider
+import com.amazonaws.auth.{AWSStaticCredentialsProvider, BasicSessionCredentials, DefaultAWSCredentialsProviderChain}
 import com.amazonaws.client.builder.AwsClientBuilder.EndpointConfiguration
 import com.amazonaws.services.dynamodbv2.document.{DynamoDB, ItemCollection, ScanOutcome}
 import com.amazonaws.services.dynamodbv2.{AmazonDynamoDB, AmazonDynamoDBClientBuilder}
@@ -31,16 +31,33 @@ import org.apache.spark.sql.sources.Filter
 
 private[dynamodb] trait DynamoConnector {
 
-    def getDynamoDB(region: Option[String] = None, assumedArn: Option[String] = None): DynamoDB = {
-        val client: AmazonDynamoDB = getDynamoDBClient(region, assumedArn)
+    def getDynamoDB(region: Option[String] = None, roleArn: Option[String] = None): DynamoDB = {
+        val client: AmazonDynamoDB = getDynamoDBClient(region, roleArn)
         new DynamoDB(client)
     }
 
+    private def getDynamoDBClient(region: Option[String] = None, roleArn: Option[String] = None): AmazonDynamoDB = {
+        val chosenRegion = region.getOrElse(sys.env.getOrElse("aws.dynamodb.region", "us-east-1"))
+        val credentials = getCredentials(chosenRegion, roleArn)
+
+        Option(System.getProperty("aws.dynamodb.endpoint")).map(endpoint => {
+            AmazonDynamoDBClientBuilder.standard()
+                .withCredentials(credentials)
+                .withEndpointConfiguration(new EndpointConfiguration(endpoint, chosenRegion))
+                .build()
+        }).getOrElse(
+            AmazonDynamoDBClientBuilder.standard()
+                .withCredentials(credentials)
+                .withRegion(chosenRegion)
+                .build()
+        )
+    }
+
     /**
       * 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 => {
+      **/
+    private def getCredentials(chosenRegion: String, roleArn: Option[String]) = {
+        roleArn.map(arn => {
             val stsClient = AWSSecurityTokenServiceClientBuilder
                 .standard()
                 .withCredentials(new DefaultAWSCredentialsProviderChain)
@@ -59,24 +76,7 @@ private[dynamodb] trait DynamoConnector {
             )
             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 => {
-            AmazonDynamoDBClientBuilder.standard()
-                .withCredentials(credentials)
-                .withEndpointConfiguration(new EndpointConfiguration(endpoint, chosenRegion))
-                .build()
-        }).getOrElse(
-            AmazonDynamoDBClientBuilder.standard()
-                .withCredentials(credentials)
-                .withRegion(chosenRegion)
-                .build()
-        )
+            .getOrElse(new DefaultAWSCredentialsProviderChain)
     }
 
     val keySchema: KeySchema

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

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

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

@@ -34,10 +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")
+    private val roleArn = parameters.get("roleArn")
 
     override val (keySchema, readLimit, itemLimit, totalSizeInBytes) = {
-        val table = getDynamoDB(region, assumedArn).getTable(tableName)
+        val table = getDynamoDB(region, roleArn).getTable(tableName)
         val indexDesc = table.describe().getGlobalSecondaryIndexes.asScala.find(_.getIndexName == indexName).get
 
         // Key schema.
@@ -82,7 +82,7 @@ private[dynamodb] class TableIndexConnector(tableName: String, indexName: String
             scanSpec.withExpressionSpec(xspec.buildForScan())
         }
 
-        getDynamoDB(region, assumedArn).getTable(tableName).getIndex(indexName).scan(scanSpec)
+        getDynamoDB(region, roleArn).getTable(tableName).getIndex(indexName).scan(scanSpec)
     }
 
 }