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 index ce9857e13fd87..ad811dbd17b96 100644 --- 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 getAdlProxyHost() + { + return Optional.ofNullable(adlProxyHost); + } + @ConfigSecuritySensitive @Config("hive.azure.abfs.oauth.endpoint") public HiveAzureConfig setAbfsOAuthClientEndpoint(String endpoint) diff --git a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/azure/TrinoAzureConfigurationInitializer.java b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/azure/TrinoAzureConfigurationInitializer.java index 7810d20bb0f9e..e792e6670e15b 100644 --- a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/azure/TrinoAzureConfigurationInitializer.java +++ b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/azure/TrinoAzureConfigurationInitializer.java @@ -13,6 +13,7 @@ */ 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; @@ -20,6 +21,9 @@ 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 dropEmpty(Optional 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 index 240771c401f4a..e43fc2b5d3a3d 100644 --- 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");