Allow setting connection proxy for ADL connection

trinodb · Apr 20, 2021 · 5ed22be · 5ed22be
1 parent 8942fdd
commit 5ed22be
Show file tree

Hide file tree

Showing 3 changed files with 30 additions and 0 deletions.
diff --git a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/azure/HiveAzureConfig.java b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/azure/HiveAzureConfig.java
@@ -13,6 +13,7 @@
  */
 package io.trino.plugin.hive.azure;
 
+import com.google.common.net.HostAndPort;
 import io.airlift.configuration.Config;
 import io.airlift.configuration.ConfigSecuritySensitive;
 
@@ -27,6 +28,7 @@ public class HiveAzureConfig
     private String adlClientId;
     private String adlCredential;
     private String adlRefreshUrl;
+    private HostAndPort adlProxyHost;
     private String abfsOAuthClientEndpoint;
     private String abfsOAuthClientId;
     private String abfsOAuthClientSecret;
@@ -122,6 +124,18 @@ public HiveAzureConfig setAdlRefreshUrl(String adlRefreshUrl)
         return this;
     }
 
+    @Config("hive.azure.adl-proxy-host")
+    public HiveAzureConfig setAdlProxyHost(HostAndPort adlProxyHost)
+    {
+        this.adlProxyHost = adlProxyHost;
+        return this;
+    }
+
+    public Optional<HostAndPort> getAdlProxyHost()
+    {
+        return Optional.ofNullable(adlProxyHost);
+    }
+
     @ConfigSecuritySensitive
     @Config("hive.azure.abfs.oauth.endpoint")
     public HiveAzureConfig setAbfsOAuthClientEndpoint(String endpoint)

diff --git a/...ino-hive/src/main/java/io/trino/plugin/hive/azure/TrinoAzureConfigurationInitializer.java b/...ino-hive/src/main/java/io/trino/plugin/hive/azure/TrinoAzureConfigurationInitializer.java
@@ -13,13 +13,17 @@
  */
 package io.trino.plugin.hive.azure;
 
+import com.google.common.net.HostAndPort;
 import io.trino.plugin.hive.ConfigurationInitializer;
 import org.apache.hadoop.conf.Configuration;
 import org.apache.hadoop.fs.adl.AdlFileSystem;
 import org.apache.hadoop.fs.azurebfs.AzureBlobFileSystem;
 
 import javax.inject.Inject;
 
+import java.net.InetSocketAddress;
+import java.net.Proxy;
+import java.net.Proxy.Type;
 import java.util.Optional;
 
 import static com.google.common.base.Preconditions.checkArgument;
@@ -79,6 +83,9 @@ public TrinoAzureConfigurationInitializer(HiveAzureConfig config)
         checkArgument(
                 !(abfsAccessKey.isPresent() && abfsOAuthClientSecret.isPresent()),
                 "Multiple ABFS authentication methods configured: access key and OAuth2");
+
+        config.getAdlProxyHost().ifPresent(proxyHost ->
+                io.trino.hadoop.$internal.com.microsoft.azure.datalake.store.HttpTransport.setConnectionProxy(proxyForHost(proxyHost)));
     }
 
     @Override
@@ -116,4 +123,9 @@ private static Optional<String> dropEmpty(Optional<String> optional)
     {
         return optional.filter(value -> !value.isEmpty());
     }
+
+    private static Proxy proxyForHost(HostAndPort address)
+    {
+        return new Proxy(Type.HTTP, new InetSocketAddress(address.getHost(), address.getPort()));
+    }
 }
diff --git a/plugin/trino-hive/src/test/java/io/trino/plugin/hive/azure/TestHiveAzureConfig.java b/plugin/trino-hive/src/test/java/io/trino/plugin/hive/azure/TestHiveAzureConfig.java
@@ -14,6 +14,7 @@
 package io.trino.plugin.hive.azure;
 
 import com.google.common.collect.ImmutableMap;
+import com.google.common.net.HostAndPort;
 import org.testng.annotations.Test;
 
 import java.util.Map;
@@ -34,6 +35,7 @@ public void testDefaults()
                 .setAbfsAccessKey(null)
                 .setAdlClientId(null)
                 .setAdlCredential(null)
+                .setAdlProxyHost(null)
                 .setAdlRefreshUrl(null)
                 .setAbfsOAuthClientEndpoint(null)
                 .setAbfsOAuthClientId(null)
@@ -51,6 +53,7 @@ public void testExplicitPropertyMappings()
                 .put("hive.azure.adl-client-id", "adlclientid")
                 .put("hive.azure.adl-credential", "adlcredential")
                 .put("hive.azure.adl-refresh-url", "adlrefreshurl")
+                .put("hive.azure.adl-proxy-host", "proxy-host:9800")
                 .put("hive.azure.abfs.oauth.endpoint", "abfsoauthendpoint")
                 .put("hive.azure.abfs.oauth.client-id", "abfsoauthclientid")
                 .put("hive.azure.abfs.oauth.secret", "abfsoauthsecret")
@@ -64,6 +67,7 @@ public void testExplicitPropertyMappings()
                 .setAdlClientId("adlclientid")
                 .setAdlCredential("adlcredential")
                 .setAdlRefreshUrl("adlrefreshurl")
+                .setAdlProxyHost(HostAndPort.fromParts("proxy-host", 9800))
                 .setAbfsOAuthClientEndpoint("abfsoauthendpoint")
                 .setAbfsOAuthClientId("abfsoauthclientid")
                 .setAbfsOAuthClientSecret("abfsoauthsecret");