Allow setting the Glue STS endpoint and region

trinodb · Sep 30, 2022 · 817e4a8 · 817e4a8
1 parent 072fef2
commit 817e4a8
Show file tree

Hide file tree

Showing 4 changed files with 59 additions and 1 deletion.
diff --git a/docs/src/main/sphinx/connector/hive.rst b/docs/src/main/sphinx/connector/hive.rst
@@ -606,6 +606,13 @@ Property Name                                        Description
 ``hive.metastore.glue.endpoint-url``                 Glue API endpoint URL (optional).
                                                      Example: ``https://glue.us-east-1.amazonaws.com``
 
+``hive.metastore.glue.sts.region``                   AWS region of the STS service to authenticate with. This is
+                                                     required when running in a GovCloud region.
+                                                     Example: ``us-gov-east-1``
+
+``hive.metastore.glue.sts.endpoint``                 STS endpoint URL to use when authenticating to Glue (optional).
+                                                     Example: ``https://sts.us-gov-east-1.amazonaws.com``
+
 ``hive.metastore.glue.pin-client-to-current-region`` Pin Glue requests to the same region as the EC2 instance
                                                      where Trino is running, defaults to ``false``.
 

diff --git a/...trino-hive/src/main/java/io/trino/plugin/hive/metastore/glue/GlueCredentialsProvider.java b/...trino-hive/src/main/java/io/trino/plugin/hive/metastore/glue/GlueCredentialsProvider.java
@@ -18,10 +18,13 @@
 import com.amazonaws.auth.BasicAWSCredentials;
 import com.amazonaws.auth.DefaultAWSCredentialsProviderChain;
 import com.amazonaws.auth.STSAssumeRoleSessionCredentialsProvider;
+import com.amazonaws.client.builder.AwsClientBuilder;
+import com.amazonaws.services.securitytoken.AWSSecurityTokenServiceClientBuilder;
 
 import javax.inject.Inject;
 import javax.inject.Provider;
 
+import static io.trino.plugin.hive.aws.AwsCurrentRegionHolder.getCurrentRegionFromEC2Metadata;
 import static java.lang.String.format;
 
 public class GlueCredentialsProvider
@@ -45,10 +48,24 @@ public GlueCredentialsProvider(GlueHiveMetastoreConfig config)
                 provider = DefaultAWSCredentialsProviderChain.getInstance();
             }
             if (config.getIamRole().isPresent()) {
+                AWSSecurityTokenServiceClientBuilder stsClientBuilder = AWSSecurityTokenServiceClientBuilder
+                        .standard()
+                        .withCredentials(provider);
+
+                if (config.getGlueStsEndpointUrl().isPresent() && config.getGlueStsRegion().isPresent()) {
+                    stsClientBuilder.setEndpointConfiguration(new AwsClientBuilder.EndpointConfiguration(config.getGlueStsEndpointUrl().get(), config.getGlueStsRegion().get()));
+                }
+                else if (config.getGlueStsRegion().isPresent()) {
+                    stsClientBuilder.setRegion(config.getGlueStsRegion().get());
+                }
+                else if (config.getPinGlueClientToCurrentRegion()) {
+                    stsClientBuilder.setRegion(getCurrentRegionFromEC2Metadata().getName());
+                }
+
                 provider = new STSAssumeRoleSessionCredentialsProvider
                         .Builder(config.getIamRole().get(), "trino-session")
                         .withExternalId(config.getExternalId().orElse(null))
-                        .withLongLivedCredentialsProvider(provider)
+                        .withStsClient(stsClientBuilder.build())
                         .build();
             }
             this.credentialsProvider = provider;

diff --git a/...trino-hive/src/main/java/io/trino/plugin/hive/metastore/glue/GlueHiveMetastoreConfig.java b/...trino-hive/src/main/java/io/trino/plugin/hive/metastore/glue/GlueHiveMetastoreConfig.java
@@ -28,6 +28,8 @@ public class GlueHiveMetastoreConfig
 {
     private Optional<String> glueRegion = Optional.empty();
     private Optional<String> glueEndpointUrl = Optional.empty();
+    private Optional<String> glueStsRegion = Optional.empty();
+    private Optional<String> glueStsEndpointUrl = Optional.empty();
     private boolean pinGlueClientToCurrentRegion;
     private int maxGlueErrorRetries = 10;
     private int maxGlueConnections = 30;
@@ -70,6 +72,32 @@ public GlueHiveMetastoreConfig setGlueEndpointUrl(String glueEndpointUrl)
         return this;
     }
 
+    public Optional<String> getGlueStsRegion()
+    {
+        return glueStsRegion;
+    }
+
+    @Config("hive.metastore.glue.sts.region")
+    @ConfigDescription("AWS STS signing region for Glue authentication")
+    public GlueHiveMetastoreConfig setGlueStsRegion(String glueStsRegion)
+    {
+        this.glueStsRegion = Optional.ofNullable(glueStsRegion);
+        return this;
+    }
+
+    public Optional<String> getGlueStsEndpointUrl()
+    {
+        return glueStsEndpointUrl;
+    }
+
+    @Config("hive.metastore.glue.sts.endpoint")
+    @ConfigDescription("AWS STS endpoint for Glue authentication")
+    public GlueHiveMetastoreConfig setGlueStsEndpointUrl(String glueStsEndpointUrl)
+    {
+        this.glueStsEndpointUrl = Optional.ofNullable(glueStsEndpointUrl);
+        return this;
+    }
+
     public boolean getPinGlueClientToCurrentRegion()
     {
         return pinGlueClientToCurrentRegion;

diff --git a/...o-hive/src/test/java/io/trino/plugin/hive/metastore/glue/TestGlueHiveMetastoreConfig.java b/...o-hive/src/test/java/io/trino/plugin/hive/metastore/glue/TestGlueHiveMetastoreConfig.java
@@ -30,6 +30,8 @@ public void testDefaults()
         assertRecordedDefaults(recordDefaults(GlueHiveMetastoreConfig.class)
                 .setGlueRegion(null)
                 .setGlueEndpointUrl(null)
+                .setGlueStsRegion(null)
+                .setGlueStsEndpointUrl(null)
                 .setPinGlueClientToCurrentRegion(false)
                 .setMaxGlueConnections(30)
                 .setMaxGlueErrorRetries(10)
@@ -53,6 +55,8 @@ public void testExplicitPropertyMapping()
         Map<String, String> properties = ImmutableMap.<String, String>builder()
                 .put("hive.metastore.glue.region", "us-east-1")
                 .put("hive.metastore.glue.endpoint-url", "http://foo.bar")
+                .put("hive.metastore.glue.sts.region", "us-east-3")
+                .put("hive.metastore.glue.sts.endpoint", "http://sts.foo.bar")
                 .put("hive.metastore.glue.pin-client-to-current-region", "true")
                 .put("hive.metastore.glue.max-connections", "10")
                 .put("hive.metastore.glue.max-error-retries", "20")
@@ -73,6 +77,8 @@ public void testExplicitPropertyMapping()
         GlueHiveMetastoreConfig expected = new GlueHiveMetastoreConfig()
                 .setGlueRegion("us-east-1")
                 .setGlueEndpointUrl("http://foo.bar")
+                .setGlueStsRegion("us-east-3")
+                .setGlueStsEndpointUrl("http://sts.foo.bar")
                 .setPinGlueClientToCurrentRegion(true)
                 .setMaxGlueConnections(10)
                 .setMaxGlueErrorRetries(20)