Allow PrestoS3FileSystem and GlueHiveMetastore to assume an aws role

Extracted from: #10864

Co-authored-by: albert-franzi <albert.franzi@schibsted.com>

pom.xml

@@ -1128,6 +1128,10 @@
                         <groupId>joda-time</groupId>
                         <artifactId>joda-time</artifactId>
                     </exclusion>
+                    <exclusion>
+                        <groupId>org.apache.httpcomponents</groupId>
+                        <artifactId>httpclient</artifactId>
+                    </exclusion>
                 </exclusions>
             </dependency>
 
@@ -1425,6 +1429,22 @@
                 </exclusions>
             </dependency>
 
+            <dependency>
+                <groupId>com.amazonaws</groupId>
+                <artifactId>aws-java-sdk-sts</artifactId>
+                <version>${dep.aws-sdk.version}</version>
+                <exclusions>
+                    <exclusion>
+                        <groupId>commons-logging</groupId>
+                        <artifactId>commons-logging</artifactId>
+                    </exclusion>
+                    <exclusion>
+                        <groupId>joda-time</groupId>
+                        <artifactId>joda-time</artifactId>
+                    </exclusion>
+                </exclusions>
+            </dependency>
+
             <dependency>
                 <groupId>com.facebook.presto</groupId>
                 <artifactId>testing-mysql-server-5</artifactId>

presto-docs/src/main/sphinx/connector/hive.rst

@@ -211,6 +211,8 @@ Property Name                                Description
 
 ``hive.s3.aws-secret-key``                   Default AWS secret key to use.
 
+``hive.s3.iam-role``                         IAM role to assume.
+
 ``hive.s3.endpoint``                         The S3 storage endpoint server. This can be used to
                                              connect to an S3-compatible storage system instead
                                              of AWS. When using v4 signatures, it is recommended to
@@ -270,10 +272,11 @@ it is highly recommended that you set ``hive.s3.use-instance-credentials``
 to ``true`` and use IAM Roles for EC2 to govern access to S3. If this is
 the case, your EC2 instances will need to be assigned an IAM Role which
 grants appropriate access to the data stored in the S3 bucket(s) you wish
-to use.  This is much cleaner than setting AWS access and secret keys in
-the ``hive.s3.aws-access-key`` and ``hive.s3.aws-secret-key`` settings, and also
-allows EC2 to automatically rotate credentials on a regular basis without
-any additional work on your part.
+to use. It's also possible to configure an IAM role with ``hive.s3.iam-role``
+that will be assumed for accessing any S3 bucket. This is much cleaner than
+setting AWS access and secret keys in the ``hive.s3.aws-access-key``
+and ``hive.s3.aws-secret-key`` settings, and also allows EC2 to automatically
+rotate credentials on a regular basis without any additional work on your part.
 
 Custom S3 Credentials Provider
 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^

presto-hive-metastore/pom.xml

@@ -96,6 +96,11 @@
             <artifactId>aws-java-sdk-glue</artifactId>
         </dependency>
 
+        <dependency>
+            <groupId>com.amazonaws</groupId>
+            <artifactId>aws-java-sdk-sts</artifactId>
+        </dependency>
+
         <dependency>
             <groupId>com.google.guava</groupId>
             <artifactId>guava</artifactId>

presto-hive-metastore/src/main/java/com/facebook/presto/hive/metastore/glue/GlueHiveMetastore.java

@@ -15,6 +15,8 @@
 
 import com.amazonaws.AmazonServiceException;
 import com.amazonaws.ClientConfiguration;
+import com.amazonaws.auth.AWSCredentialsProvider;
+import com.amazonaws.auth.STSAssumeRoleSessionCredentialsProvider;
 import com.amazonaws.regions.Region;
 import com.amazonaws.regions.Regions;
 import com.amazonaws.services.glue.AWSGlueAsync;
@@ -191,6 +193,13 @@ else if (config.getPinGlueClientToCurrentRegion()) {
             }
         }
 
+        if (config.getIamRole().isPresent()) {
+            AWSCredentialsProvider credentialsProvider = new STSAssumeRoleSessionCredentialsProvider
+                    .Builder(config.getIamRole().get(), "presto-session")
+                    .build();
+            asyncGlueClientBuilder.setCredentials(credentialsProvider);
+        }
+
         return asyncGlueClientBuilder.build();
     }
 

presto-hive-metastore/src/main/java/com/facebook/presto/hive/metastore/glue/GlueHiveMetastoreConfig.java

@@ -30,6 +30,7 @@ public class GlueHiveMetastoreConfig
     private Optional<String> catalogId = Optional.empty();
     private int partitionSegments = 5;
     private int getPartitionThreads = 20;
+    private Optional<String> iamRole = Optional.empty();
 
     public Optional<String> getGlueRegion()
     {
@@ -125,4 +126,17 @@ public GlueHiveMetastoreConfig setGetPartitionThreads(int getPartitionThreads)
         this.getPartitionThreads = getPartitionThreads;
         return this;
     }
+
+    public Optional<String> getIamRole()
+    {
+        return iamRole;
+    }
+
+    @Config("hive.metastore.glue.iam-role")
+    @ConfigDescription("IAM role to assume when connecting to the Hive Glue metastore")
+    public GlueHiveMetastoreConfig setIamRole(String iamRole)
+    {
+        this.iamRole = Optional.ofNullable(iamRole);
+        return this;
+    }
 }

presto-hive-metastore/src/test/java/com/facebook/presto/hive/metastore/glue/TestGlueHiveMetastoreConfig.java

@@ -34,7 +34,8 @@ public void testDefaults()
                 .setDefaultWarehouseDir(null)
                 .setCatalogId(null)
                 .setPartitionSegments(5)
-                .setGetPartitionThreads(20));
+                .setGetPartitionThreads(20)
+                .setIamRole(null));
     }
 
     @Test
@@ -48,6 +49,7 @@ public void testExplicitPropertyMapping()
                 .put("hive.metastore.glue.catalogid", "0123456789")
                 .put("hive.metastore.glue.partitions-segments", "10")
                 .put("hive.metastore.glue.get-partition-threads", "42")
+                .put("hive.metastore.glue.iam-role", "role")
                 .build();
 
         GlueHiveMetastoreConfig expected = new GlueHiveMetastoreConfig()
@@ -57,7 +59,8 @@ public void testExplicitPropertyMapping()
                 .setDefaultWarehouseDir("/location")
                 .setCatalogId("0123456789")
                 .setPartitionSegments(10)
-                .setGetPartitionThreads(42);
+                .setGetPartitionThreads(42)
+                .setIamRole("role");
 
         assertFullMapping(properties, expected);
     }

presto-hive/pom.xml

@@ -243,6 +243,11 @@
             <artifactId>gcs-connector</artifactId>
         </dependency>
 
+        <dependency>
+            <groupId>com.amazonaws</groupId>
+            <artifactId>aws-java-sdk-sts</artifactId>
+        </dependency>
+
         <dependency>
             <groupId>org.xerial.snappy</groupId>
             <artifactId>snappy-java</artifactId>

presto-hive/src/main/java/com/facebook/presto/hive/s3/HiveS3Config.java

@@ -38,6 +38,7 @@ public class HiveS3Config
     private PrestoS3SignerType s3SignerType;
     private boolean s3PathStyleAccess;
     private boolean s3UseInstanceCredentials = true;
+    private String s3IamRole;
     private boolean s3SslEnabled = true;
     private boolean s3SseEnabled;
     private PrestoS3SseType s3SseType = PrestoS3SseType.S3;
@@ -133,6 +134,18 @@ public HiveS3Config setS3UseInstanceCredentials(boolean s3UseInstanceCredentials
         return this;
     }
 
+    public String getS3IamRole()
+    {
+        return s3IamRole;
+    }
+
+    @Config("hive.s3.iam-role")
+    public HiveS3Config setS3IamRole(String s3IamRole)
+    {
+        this.s3IamRole = s3IamRole;
+        return this;
+    }
+
     public boolean isS3SslEnabled()
     {
         return s3SslEnabled;

presto-hive/src/main/java/com/facebook/presto/hive/s3/PrestoS3ConfigurationUpdater.java

@@ -30,6 +30,7 @@ public class PrestoS3ConfigurationUpdater
     private final PrestoS3SignerType signerType;
     private final boolean pathStyleAccess;
     private final boolean useInstanceCredentials;
+    private String s3IamRole;
     private final boolean sslEnabled;
     private final boolean sseEnabled;
     private final PrestoS3SseType sseType;
@@ -60,6 +61,7 @@ public PrestoS3ConfigurationUpdater(HiveS3Config config)
         this.signerType = config.getS3SignerType();
         this.pathStyleAccess = config.isS3PathStyleAccess();
         this.useInstanceCredentials = config.isS3UseInstanceCredentials();
+        this.s3IamRole = config.getS3IamRole();
         this.sslEnabled = config.isS3SslEnabled();
         this.sseEnabled = config.isS3SseEnabled();
         this.sseType = config.getS3SseType();
@@ -104,6 +106,9 @@ public void updateConfiguration(Configuration config)
         }
         config.setBoolean(S3_PATH_STYLE_ACCESS, pathStyleAccess);
         config.setBoolean(S3_USE_INSTANCE_CREDENTIALS, useInstanceCredentials);
+        if (s3IamRole != null) {
+            config.set(S3_IAM_ROLE, s3IamRole);
+        }
         config.setBoolean(S3_SSL_ENABLED, sslEnabled);
         config.setBoolean(S3_SSE_ENABLED, sseEnabled);
         config.set(S3_SSE_TYPE, sseType.name());

presto-hive/src/main/java/com/facebook/presto/hive/s3/PrestoS3FileSystem.java

@@ -23,6 +23,7 @@
 import com.amazonaws.auth.BasicAWSCredentials;
 import com.amazonaws.auth.DefaultAWSCredentialsProviderChain;
 import com.amazonaws.auth.InstanceProfileCredentialsProvider;
+import com.amazonaws.auth.STSAssumeRoleSessionCredentialsProvider;
 import com.amazonaws.client.builder.AwsClientBuilder.EndpointConfiguration;
 import com.amazonaws.event.ProgressEvent;
 import com.amazonaws.event.ProgressEventType;
@@ -104,6 +105,7 @@
 import static com.facebook.presto.hive.s3.S3ConfigurationUpdater.S3_CREDENTIALS_PROVIDER;
 import static com.facebook.presto.hive.s3.S3ConfigurationUpdater.S3_ENCRYPTION_MATERIALS_PROVIDER;
 import static com.facebook.presto.hive.s3.S3ConfigurationUpdater.S3_ENDPOINT;
+import static com.facebook.presto.hive.s3.S3ConfigurationUpdater.S3_IAM_ROLE;
 import static com.facebook.presto.hive.s3.S3ConfigurationUpdater.S3_KMS_KEY_ID;
 import static com.facebook.presto.hive.s3.S3ConfigurationUpdater.S3_MAX_BACKOFF_TIME;
 import static com.facebook.presto.hive.s3.S3ConfigurationUpdater.S3_MAX_CLIENT_RETRIES;
@@ -173,6 +175,7 @@ public class PrestoS3FileSystem
     private Duration maxBackoffTime;
     private Duration maxRetryTime;
     private boolean useInstanceCredentials;
+    private String s3IamRole;
     private boolean pinS3ClientToCurrentRegion;
     private boolean sseEnabled;
     private PrestoS3SseType sseType;
@@ -210,6 +213,9 @@ public void initialize(URI uri, Configuration conf)
         this.isPathStyleAccess = conf.getBoolean(S3_PATH_STYLE_ACCESS, defaults.isS3PathStyleAccess());
         this.useInstanceCredentials = conf.getBoolean(S3_USE_INSTANCE_CREDENTIALS, defaults.isS3UseInstanceCredentials());
         this.pinS3ClientToCurrentRegion = conf.getBoolean(S3_PIN_CLIENT_TO_CURRENT_REGION, defaults.isPinS3ClientToCurrentRegion());
+        this.s3IamRole = conf.get(S3_IAM_ROLE, defaults.getS3IamRole());
+        verify(!(useInstanceCredentials && this.s3IamRole != null),
+                "Invalid configuration: either use instance credentials or specify an iam role");
         verify((pinS3ClientToCurrentRegion && conf.get(S3_ENDPOINT) == null) || !pinS3ClientToCurrentRegion,
                 "Invalid configuration: either endpoint can be set or S3 client can be pinned to the current region");
         this.sseEnabled = conf.getBoolean(S3_SSE_ENABLED, defaults.isS3SseEnabled());
@@ -794,6 +800,10 @@ private AWSCredentialsProvider createAwsCredentialsProvider(URI uri, Configurati
             return InstanceProfileCredentialsProvider.getInstance();
         }
 
+        if (s3IamRole != null) {
+            return new STSAssumeRoleSessionCredentialsProvider.Builder(this.s3IamRole, "presto-session").build();
+        }
+
         String providerClass = conf.get(S3_CREDENTIALS_PROVIDER);
         if (!isNullOrEmpty(providerClass)) {
             return getCustomAWSCredentialsProvider(uri, conf, providerClass);

presto-hive/src/main/java/com/facebook/presto/hive/s3/S3ConfigurationUpdater.java

@@ -27,6 +27,7 @@ public interface S3ConfigurationUpdater
     String S3_ENCRYPTION_MATERIALS_PROVIDER = "presto.s3.encryption-materials-provider";
     String S3_PIN_CLIENT_TO_CURRENT_REGION = "presto.s3.pin-client-to-current-region";
     String S3_USE_INSTANCE_CREDENTIALS = "presto.s3.use-instance-credentials";
+    String S3_IAM_ROLE = "presto.hive.s3.iam-role";
     String S3_MULTIPART_MIN_PART_SIZE = "presto.s3.multipart.min-part-size";
     String S3_MULTIPART_MIN_FILE_SIZE = "presto.s3.multipart.min-file-size";
     String S3_STAGING_DIRECTORY = "presto.s3.staging-directory";

presto-hive/src/test/java/com/facebook/presto/hive/s3/TestHiveS3Config.java

@@ -40,6 +40,7 @@ public void testDefaults()
                 .setS3SignerType(null)
                 .setS3PathStyleAccess(false)
                 .setS3UseInstanceCredentials(true)
+                .setS3IamRole(null)
                 .setS3SslEnabled(true)
                 .setS3SseEnabled(false)
                 .setS3SseType(PrestoS3SseType.S3)
@@ -72,6 +73,7 @@ public void testExplicitPropertyMappings()
                 .put("hive.s3.signer-type", "S3SignerType")
                 .put("hive.s3.path-style-access", "true")
                 .put("hive.s3.use-instance-credentials", "false")
+                .put("hive.s3.iam-role", "roleArn")
                 .put("hive.s3.ssl.enabled", "false")
                 .put("hive.s3.sse.enabled", "true")
                 .put("hive.s3.sse.type", "KMS")
@@ -101,6 +103,7 @@ public void testExplicitPropertyMappings()
                 .setS3SignerType(PrestoS3SignerType.S3SignerType)
                 .setS3PathStyleAccess(true)
                 .setS3UseInstanceCredentials(false)
+                .setS3IamRole("roleArn")
                 .setS3SslEnabled(false)
                 .setS3SseEnabled(true)
                 .setS3SseType(PrestoS3SseType.KMS)

presto-hive/src/test/java/com/facebook/presto/hive/s3/TestPrestoS3FileSystem.java

@@ -20,6 +20,7 @@
 import com.amazonaws.auth.AWSStaticCredentialsProvider;
 import com.amazonaws.auth.DefaultAWSCredentialsProviderChain;
 import com.amazonaws.auth.InstanceProfileCredentialsProvider;
+import com.amazonaws.auth.STSAssumeRoleSessionCredentialsProvider;
 import com.amazonaws.services.s3.AmazonS3Client;
 import com.amazonaws.services.s3.AmazonS3EncryptionClient;
 import com.amazonaws.services.s3.S3ClientOptions;
@@ -71,6 +72,7 @@
 import static com.facebook.presto.hive.s3.S3ConfigurationUpdater.S3_CREDENTIALS_PROVIDER;
 import static com.facebook.presto.hive.s3.S3ConfigurationUpdater.S3_ENCRYPTION_MATERIALS_PROVIDER;
 import static com.facebook.presto.hive.s3.S3ConfigurationUpdater.S3_ENDPOINT;
+import static com.facebook.presto.hive.s3.S3ConfigurationUpdater.S3_IAM_ROLE;
 import static com.facebook.presto.hive.s3.S3ConfigurationUpdater.S3_KMS_KEY_ID;
 import static com.facebook.presto.hive.s3.S3ConfigurationUpdater.S3_MAX_BACKOFF_TIME;
 import static com.facebook.presto.hive.s3.S3ConfigurationUpdater.S3_MAX_CLIENT_RETRIES;
@@ -144,6 +146,20 @@ public void testEndpointWithPinToCurrentRegionConfiguration()
         }
     }
 
+    @Test
+    public void testAssumeRoleCredentials()
+            throws Exception
+    {
+        Configuration config = new Configuration();
+        config.set(S3_IAM_ROLE, "role");
+        config.setBoolean(S3_USE_INSTANCE_CREDENTIALS, false);
+
+        try (PrestoS3FileSystem fs = new PrestoS3FileSystem()) {
+            fs.initialize(new URI("s3n://test-bucket/"), config);
+            assertInstanceOf(getAwsCredentialsProvider(fs), STSAssumeRoleSessionCredentialsProvider.class);
+        }
+    }
+
     @Test
     public void testInstanceCredentialsEnabled()
             throws Exception