Add support for Glue endpoint URL

trinodb · Apr 1, 2020 · 728be1a · 728be1a
1 parent f922c85
commit 728be1a
Show file tree

Hide file tree

Showing 3 changed files with 26 additions and 2 deletions.
diff --git a/presto-hive/src/main/java/io/prestosql/plugin/hive/metastore/glue/GlueHiveMetastore.java b/presto-hive/src/main/java/io/prestosql/plugin/hive/metastore/glue/GlueHiveMetastore.java
@@ -21,6 +21,7 @@
 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.services.glue.AWSGlueAsync;
 import com.amazonaws.services.glue.AWSGlueAsyncClientBuilder;
 import com.amazonaws.services.glue.model.AlreadyExistsException;
@@ -112,6 +113,7 @@
 import java.util.concurrent.Future;
 import java.util.function.Function;
 
+import static com.google.common.base.Preconditions.checkArgument;
 import static com.google.common.base.Strings.isNullOrEmpty;
 import static com.google.common.collect.Comparators.lexicographical;
 import static com.google.common.collect.ImmutableMap.toImmutableMap;
@@ -182,8 +184,13 @@ private static AWSGlueAsync createAsyncGlueClient(GlueHiveMetastoreConfig config
         ClientConfiguration clientConfig = new ClientConfiguration().withMaxConnections(config.getMaxGlueConnections());
         AWSGlueAsyncClientBuilder asyncGlueClientBuilder = AWSGlueAsyncClientBuilder.standard()
                 .withClientConfiguration(clientConfig);
-
-        if (config.getGlueRegion().isPresent()) {
+        if (config.getGlueEndpointUrl().isPresent()) {
+            checkArgument(config.getGlueRegion().isPresent(), "Glue region must be set when Glue endpoint URL is set");
+            asyncGlueClientBuilder.setEndpointConfiguration(new EndpointConfiguration(
+                    config.getGlueEndpointUrl().get(),
+                    config.getGlueRegion().get()));
+        }
+        else if (config.getGlueRegion().isPresent()) {
             asyncGlueClientBuilder.setRegion(config.getGlueRegion().get());
         }
         else if (config.getPinGlueClientToCurrentRegion()) {

diff --git a/...o-hive/src/main/java/io/prestosql/plugin/hive/metastore/glue/GlueHiveMetastoreConfig.java b/...o-hive/src/main/java/io/prestosql/plugin/hive/metastore/glue/GlueHiveMetastoreConfig.java
@@ -25,6 +25,7 @@
 public class GlueHiveMetastoreConfig
 {
     private Optional<String> glueRegion = Optional.empty();
+    private Optional<String> glueEndpointUrl = Optional.empty();
     private boolean pinGlueClientToCurrentRegion;
     private int maxGlueConnections = 5;
     private Optional<String> defaultWarehouseDir = Optional.empty();
@@ -50,6 +51,19 @@ public GlueHiveMetastoreConfig setGlueRegion(String region)
         return this;
     }
 
+    public Optional<String> getGlueEndpointUrl()
+    {
+        return glueEndpointUrl;
+    }
+
+    @Config("hive.metastore.glue.endpoint-url")
+    @ConfigDescription("Glue API endpoint URL")
+    public GlueHiveMetastoreConfig setGlueEndpointUrl(String glueEndpointUrl)
+    {
+        this.glueEndpointUrl = Optional.ofNullable(glueEndpointUrl);
+        return this;
+    }
+
     public boolean getPinGlueClientToCurrentRegion()
     {
         return pinGlueClientToCurrentRegion;

diff --git a/...ve/src/test/java/io/prestosql/plugin/hive/metastore/glue/TestGlueHiveMetastoreConfig.java b/...ve/src/test/java/io/prestosql/plugin/hive/metastore/glue/TestGlueHiveMetastoreConfig.java
@@ -29,6 +29,7 @@ public void testDefaults()
     {
         assertRecordedDefaults(recordDefaults(GlueHiveMetastoreConfig.class)
                 .setGlueRegion(null)
+                .setGlueEndpointUrl(null)
                 .setPinGlueClientToCurrentRegion(false)
                 .setMaxGlueConnections(5)
                 .setDefaultWarehouseDir(null)
@@ -47,6 +48,7 @@ public void testExplicitPropertyMapping()
     {
         Map<String, String> properties = new ImmutableMap.Builder<String, String>()
                 .put("hive.metastore.glue.region", "us-east-1")
+                .put("hive.metastore.glue.endpoint-url", "http://foo.bar")
                 .put("hive.metastore.glue.pin-client-to-current-region", "true")
                 .put("hive.metastore.glue.max-connections", "10")
                 .put("hive.metastore.glue.default-warehouse-dir", "/location")
@@ -62,6 +64,7 @@ public void testExplicitPropertyMapping()
 
         GlueHiveMetastoreConfig expected = new GlueHiveMetastoreConfig()
                 .setGlueRegion("us-east-1")
+                .setGlueEndpointUrl("http://foo.bar")
                 .setPinGlueClientToCurrentRegion(true)
                 .setMaxGlueConnections(10)
                 .setDefaultWarehouseDir("/location")