diff --git a/plugin/trino-snowflake/src/main/java/io/trino/plugin/snowflake/SnowflakeClientModule.java b/plugin/trino-snowflake/src/main/java/io/trino/plugin/snowflake/SnowflakeClientModule.java index 0b3432bb8b332..19fc358471915 100644 --- a/plugin/trino-snowflake/src/main/java/io/trino/plugin/snowflake/SnowflakeClientModule.java +++ b/plugin/trino-snowflake/src/main/java/io/trino/plugin/snowflake/SnowflakeClientModule.java @@ -25,12 +25,15 @@ import io.trino.plugin.jdbc.JdbcClient; import io.trino.plugin.jdbc.TypeHandlingJdbcConfig; import io.trino.plugin.jdbc.credential.CredentialProvider; +import io.trino.spi.TrinoException; import net.snowflake.client.jdbc.SnowflakeDriver; import java.net.MalformedURLException; +import java.net.URL; import java.util.Properties; import static io.airlift.configuration.ConfigBinder.configBinder; +import static io.trino.spi.StandardErrorCode.NOT_SUPPORTED; public class SnowflakeClientModule implements Module @@ -63,6 +66,30 @@ public ConnectionFactory getConnectionFactory(BaseJdbcConfig baseJdbcConfig, Sno properties.setProperty("TIME_OUTPUT_FORMAT", "HH24:MI:SS.FF9"); snowflakeConfig.getTimestampNoTimezoneAsUTC().ifPresent(as_utc -> properties.setProperty("JDBC_TREAT_TIMESTAMP_NTZ_AS_UTC", as_utc ? "true" : "false")); + // Support for Corporate proxies + if (snowflakeConfig.getHTTPProxy().isPresent()) { + String proxy = snowflakeConfig.getHTTPProxy().get(); + + URL url = new URL(proxy); + + properties.setProperty("useProxy", "true"); + properties.setProperty("proxyHost", url.getHost()); + properties.setProperty("proxyPort", Integer.toString(url.getPort())); + properties.setProperty("proxyProtocol", url.getProtocol()); + + String userInfo = url.getUserInfo(); + if (userInfo != null) { + String[] usernamePassword = userInfo.split(":", 2); + + if (usernamePassword.length != 2) { + throw new TrinoException(NOT_SUPPORTED, "Improper snowflake.http_proxy. username:password@ is optional but what was entered was not correct"); + } + + properties.setProperty("proxyUser", usernamePassword[0]); + properties.setProperty("proxyPassword", usernamePassword[1]); + } + } + return new DriverConnectionFactory(new SnowflakeDriver(), baseJdbcConfig.getConnectionUrl(), properties, credentialProvider); } } diff --git a/plugin/trino-snowflake/src/main/java/io/trino/plugin/snowflake/SnowflakeConfig.java b/plugin/trino-snowflake/src/main/java/io/trino/plugin/snowflake/SnowflakeConfig.java index 86f97efdbd24b..6dbf125201774 100644 --- a/plugin/trino-snowflake/src/main/java/io/trino/plugin/snowflake/SnowflakeConfig.java +++ b/plugin/trino-snowflake/src/main/java/io/trino/plugin/snowflake/SnowflakeConfig.java @@ -24,6 +24,7 @@ public class SnowflakeConfig private String role; private String warehouse; private Boolean timestampNoTimezoneAsUTC; + private String httpProxy; public Optional getAccount() { @@ -84,4 +85,16 @@ public SnowflakeConfig setTimestampNoTimezoneAsUTC(Boolean timestampNoTimezoneAs this.timestampNoTimezoneAsUTC = timestampNoTimezoneAsUTC; return this; } + + public Optional getHTTPProxy() + { + return Optional.ofNullable(httpProxy); + } + + @Config("snowflake.http-proxy") + public SnowflakeConfig setHTTPProxy(String httpProxy) + { + this.httpProxy = httpProxy; + return this; + } } diff --git a/plugin/trino-snowflake/src/test/java/io/trino/plugin/snowflake/SnowflakeQueryRunner.java b/plugin/trino-snowflake/src/test/java/io/trino/plugin/snowflake/SnowflakeQueryRunner.java index 231c46b5d6f33..a50debaf003b6 100644 --- a/plugin/trino-snowflake/src/test/java/io/trino/plugin/snowflake/SnowflakeQueryRunner.java +++ b/plugin/trino-snowflake/src/test/java/io/trino/plugin/snowflake/SnowflakeQueryRunner.java @@ -56,6 +56,9 @@ public static DistributedQueryRunner createSnowflakeQueryRunner( connectorProperties.putIfAbsent("snowflake.database", TestingSnowflakeServer.TEST_DATABASE); connectorProperties.putIfAbsent("snowflake.role", TestingSnowflakeServer.TEST_ROLE); connectorProperties.putIfAbsent("snowflake.warehouse", TestingSnowflakeServer.TEST_WAREHOUSE); + if (TestingSnowflakeServer.TEST_PROXY != null) { + connectorProperties.putIfAbsent("snowflake.httpProxy", TestingSnowflakeServer.TEST_PROXY); + } queryRunner.installPlugin(new SnowflakePlugin()); queryRunner.createCatalog("snowflake", "snowflake", connectorProperties); diff --git a/plugin/trino-snowflake/src/test/java/io/trino/plugin/snowflake/TestSnowflakeConfig.java b/plugin/trino-snowflake/src/test/java/io/trino/plugin/snowflake/TestSnowflakeConfig.java index 5d019ab1e7de2..eb5c32a3d063c 100644 --- a/plugin/trino-snowflake/src/test/java/io/trino/plugin/snowflake/TestSnowflakeConfig.java +++ b/plugin/trino-snowflake/src/test/java/io/trino/plugin/snowflake/TestSnowflakeConfig.java @@ -32,6 +32,7 @@ public void testDefaults() .setDatabase(null) .setRole(null) .setWarehouse(null) + .setHTTPProxy(null) .setTimestampNoTimezoneAsUTC(null)); } @@ -43,6 +44,7 @@ public void testExplicitPropertyMappings() .put("snowflake.database", "MYDATABASE") .put("snowflake.role", "MYROLE") .put("snowflake.warehouse", "MYWAREHOUSE") + .put("snowflake.http-proxy", "MYPROXY") .put("snowflake.timestamp-no-timezone-as-utc", "true") .buildOrThrow(); @@ -51,6 +53,7 @@ public void testExplicitPropertyMappings() .setDatabase("MYDATABASE") .setRole("MYROLE") .setWarehouse("MYWAREHOUSE") + .setHTTPProxy("MYPROXY") .setTimestampNoTimezoneAsUTC(true); assertFullMapping(properties, expected);