[b/429112319] Support AWS session credentials (#921)

Add support for the temporary AWS session credentials. Additional flag iam-sessiontoken allows user to set the AWS IAM session token that is required if temporary IAM session is used. This should unblock users who use SSO or IAM role for authentication.

Author: kpiotrowski

dumper/app/src/main/java/com/google/edwmigration/dumper/application/dumper/ConnectorArguments.java

@@ -127,6 +127,7 @@ public class ConnectorArguments extends DefaultArguments {
   // redshift.
   public static final String OPT_IAM_ACCESSKEYID = "iam-accesskeyid";
   public static final String OPT_IAM_SECRETACCESSKEY = "iam-secretaccesskey";
+  public static final String OPT_IAM_SESSIONTOKEN = "iam-sessiontoken";
   public static final String OPT_IAM_PROFILE = "iam-profile";
 
   // Port 8020 is used by HDFS to communicate with the NameNode.
@@ -408,6 +409,8 @@ public class ConnectorArguments extends DefaultArguments {
       parser.accepts(OPT_IAM_ACCESSKEYID).withRequiredArg();
   private final OptionSpec<String> optionRedshiftIAMSecretAccessKey =
       parser.accepts(OPT_IAM_SECRETACCESSKEY).withRequiredArg();
+  private final OptionSpec<String> optionRedshiftIAMSessionToken =
+      parser.accepts(OPT_IAM_SESSIONTOKEN).withRequiredArg();
   private final OptionSpec<String> optionRedshiftIAMProfile =
       parser.accepts(OPT_IAM_PROFILE).withRequiredArg();
 
@@ -1068,6 +1071,11 @@ public String getIAMSecretAccessKey() {
     return getOptions().valueOf(optionRedshiftIAMSecretAccessKey);
   }
 
+  @CheckForNull
+  public String getIamSessionToken() {
+    return getOptions().valueOf(optionRedshiftIAMSessionToken);
+  }
+
   @CheckForNull
   public String getIAMProfile() {
     return getOptions().valueOf(optionRedshiftIAMProfile);

dumper/app/src/main/java/com/google/edwmigration/dumper/application/dumper/connector/redshift/AbstractAwsApiTask.java

@@ -19,6 +19,7 @@
 import com.amazonaws.auth.AWSCredentialsProvider;
 import com.amazonaws.auth.AWSStaticCredentialsProvider;
 import com.amazonaws.auth.BasicAWSCredentials;
+import com.amazonaws.auth.BasicSessionCredentials;
 import com.amazonaws.auth.profile.ProfileCredentialsProvider;
 import com.amazonaws.services.cloudwatch.AmazonCloudWatch;
 import com.amazonaws.services.cloudwatch.AmazonCloudWatchClient;
@@ -115,11 +116,17 @@ private static AWSCredentialsProvider doCreateProvider(ConnectorArguments argume
     }
     String accessKeyId = arguments.getIAMAccessKeyID();
     String secretAccessKey = arguments.getIAMSecretAccessKey();
-    if (accessKeyId != null && secretAccessKey != null) {
-      BasicAWSCredentials credentials = new BasicAWSCredentials(accessKeyId, secretAccessKey);
-      return new AWSStaticCredentialsProvider(credentials);
-    } else {
+    String sessionToken = arguments.getIamSessionToken();
+    if (accessKeyId == null || secretAccessKey == null) {
       return null;
     }
+
+    if (sessionToken != null) {
+      BasicSessionCredentials credentials =
+          new BasicSessionCredentials(accessKeyId, secretAccessKey, sessionToken);
+      return new AWSStaticCredentialsProvider(credentials);
+    }
+    BasicAWSCredentials credentials = new BasicAWSCredentials(accessKeyId, secretAccessKey);
+    return new AWSStaticCredentialsProvider(credentials);
   }
 }

dumper/app/src/main/java/com/google/edwmigration/dumper/application/dumper/connector/redshift/AbstractRedshiftConnector.java

@@ -64,6 +64,11 @@
     arg = ConnectorArguments.OPT_IAM_SECRETACCESSKEY,
     description = "The IAM Secret Access Key to use for authentication.",
     required = "If present, performs explicit IAM authentication")
+@RespectsInput(
+    order = 457,
+    arg = ConnectorArguments.OPT_IAM_SESSIONTOKEN,
+    description = "The IAM Session Token to use for authentication.",
+    required = "Required if temporary IAM session is created")
 @RespectsArgumentJDBCUri
 public abstract class AbstractRedshiftConnector extends AbstractJdbcConnector {
 

dumper/app/src/main/java/com/google/edwmigration/dumper/application/dumper/connector/redshift/RedshiftUrlUtil.java

@@ -63,7 +63,7 @@ static String makeJdbcUrlRedshiftSimple(ConnectorArguments arguments)
             .toJdbcPart();
   }
 
-  // TODO:  [cluster-id]:[region] syntax.
+  // TODO: [cluster-id]:[region] syntax.
   // either profile, or key+ secret
   @Nonnull
   static String makeJdbcUrlRedshiftIAM(ConnectorArguments arguments)
@@ -89,10 +89,14 @@ private static String makeIamProperties(ConnectorArguments arguments)
     }
     String keyId = arguments.getIAMAccessKeyID();
     String secretKey = arguments.getIAMSecretAccessKey();
+    String sessionToken = arguments.getIamSessionToken();
     if (keyId != null && secretKey != null) {
-      return new JdbcPropBuilder("?=&")
-          .prop("AccessKeyID", keyId)
-          .prop("SecretAccessKey", secretKey)
+      JdbcPropBuilder builder =
+          new JdbcPropBuilder("?=&").prop("AccessKeyID", keyId).prop("SecretAccessKey", secretKey);
+      if (sessionToken != null) {
+        builder.prop("SessionToken", sessionToken);
+      }
+      return builder
           .propOrError("DbUser", arguments.getUser(), "--user must be specified")
           .toJdbcPart();
     } else {