Revert changes of Proxy

bloomberg · Dec 8, 2023 · b637985 · b637985
1 parent 043350b
commit b637985
Show file tree

Hide file tree

Showing 4 changed files with 46 additions and 0 deletions.
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
@@ -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
@@ -24,6 +24,7 @@ public class SnowflakeConfig
     private String role;
     private String warehouse;
     private Boolean timestampNoTimezoneAsUTC;
+    private String httpProxy;
 
     public Optional<String> getAccount()
     {
@@ -84,4 +85,16 @@ public SnowflakeConfig setTimestampNoTimezoneAsUTC(Boolean timestampNoTimezoneAs
         this.timestampNoTimezoneAsUTC = timestampNoTimezoneAsUTC;
         return this;
     }
+
+    public Optional<String> 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
@@ -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
@@ -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);