SNOW-1094021: Add Properties setters in SnowflakeBasicDataSource (#1800)

SNOW-1094021: Add Properties setters in SnowflakeBasicaDataSource
snowflakedb · Jul 15, 2024 · 579cb8f · 579cb8f
1 parent c3b9f8b
commit 579cb8f
Show file tree

Hide file tree

Showing 2 changed files with 242 additions and 15 deletions.
diff --git a/src/main/java/net/snowflake/client/jdbc/SnowflakeBasicDataSource.java b/src/main/java/net/snowflake/client/jdbc/SnowflakeBasicDataSource.java
@@ -13,6 +13,7 @@
 import java.util.Properties;
 import java.util.logging.Logger;
 import javax.sql.DataSource;
+import net.snowflake.client.core.SFSessionProperty;
 import net.snowflake.client.log.ArgSupplier;
 import net.snowflake.client.log.SFLogger;
 import net.snowflake.client.log.SFLoggerFactory;
@@ -22,6 +23,7 @@ public class SnowflakeBasicDataSource implements DataSource, Serializable {
   private static final long serialversionUID = 1L;
   private static final String AUTHENTICATOR_SNOWFLAKE_JWT = "SNOWFLAKE_JWT";
   private static final String AUTHENTICATOR_OAUTH = "OAUTH";
+  private static final String AUTHENTICATOR_USERNAME_PASSWORD_MFA = "USERNAME_PASSWORD_MFA";
   private String url;
 
   private String serverName;
@@ -88,12 +90,12 @@ public Connection getConnection() throws SQLException {
   public Connection getConnection(String username, String password) throws SQLException {
     if (!AUTHENTICATOR_OAUTH.equalsIgnoreCase(
         authenticator)) { // For OAuth, no username is required
-      properties.put("user", username);
+      properties.put(SFSessionProperty.USER.getPropertyKey(), username);
     }
 
     // The driver needs password for OAUTH as part of SNOW-533673 feature request.
     if (!AUTHENTICATOR_SNOWFLAKE_JWT.equalsIgnoreCase(authenticator)) {
-      properties.put("password", password);
+      properties.put(SFSessionProperty.PASSWORD.getPropertyKey(), password);
     }
 
     try {
@@ -119,15 +121,16 @@ public void setLogWriter(PrintWriter out) throws SQLException {
   @Override
   public int getLoginTimeout() throws SQLException {
     try {
-      return Integer.parseInt(properties.getProperty("loginTimeout"));
+      return Integer.parseInt(
+          properties.getProperty(SFSessionProperty.LOGIN_TIMEOUT.getPropertyKey()));
     } catch (NumberFormatException e) {
       return 0;
     }
   }
 
   @Override
   public void setLoginTimeout(int seconds) throws SQLException {
-    properties.put("loginTimeout", Integer.toString(seconds));
+    properties.put(SFSessionProperty.LOGIN_TIMEOUT.getPropertyKey(), Integer.toString(seconds));
   }
 
   @Override
@@ -150,19 +153,19 @@ public void setUrl(String url) {
   }
 
   public void setDatabaseName(String databaseName) {
-    properties.put("db", databaseName);
+    properties.put(SFSessionProperty.DATABASE.getPropertyKey(), databaseName);
   }
 
   public void setSchema(String schema) {
-    properties.put("schema", schema);
+    properties.put(SFSessionProperty.SCHEMA.getPropertyKey(), schema);
   }
 
   public void setWarehouse(String warehouse) {
-    properties.put("warehouse", warehouse);
+    properties.put(SFSessionProperty.WAREHOUSE.getPropertyKey(), warehouse);
   }
 
   public void setRole(String role) {
-    properties.put("role", role);
+    properties.put(SFSessionProperty.ROLE.getPropertyKey(), role);
   }
 
   public void setUser(String user) {
@@ -182,7 +185,7 @@ public void setPortNumber(int portNumber) {
   }
 
   public void setAccount(String account) {
-    this.properties.put("account", account);
+    this.properties.put(SFSessionProperty.ACCOUNT.getPropertyKey(), account);
   }
 
   public void setSsl(boolean ssl) {
@@ -191,12 +194,12 @@ public void setSsl(boolean ssl) {
 
   public void setAuthenticator(String authenticator) {
     this.authenticator = authenticator;
-    this.properties.put("authenticator", authenticator);
+    this.properties.put(SFSessionProperty.AUTHENTICATOR.getPropertyKey(), authenticator);
   }
 
   public void setOauthToken(String oauthToken) {
     this.setAuthenticator(AUTHENTICATOR_OAUTH);
-    this.properties.put("token", oauthToken);
+    this.properties.put(SFSessionProperty.TOKEN.getPropertyKey(), oauthToken);
   }
 
   public String getUrl() {
@@ -217,18 +220,155 @@ public String getUrl() {
 
   public void setPrivateKey(PrivateKey privateKey) {
     this.setAuthenticator(AUTHENTICATOR_SNOWFLAKE_JWT);
-    this.properties.put("privateKey", privateKey);
+    this.properties.put(SFSessionProperty.PRIVATE_KEY.getPropertyKey(), privateKey);
   }
 
   public void setPrivateKeyFile(String location, String password) {
     this.setAuthenticator(AUTHENTICATOR_SNOWFLAKE_JWT);
-    this.properties.put("private_key_file", location);
+    this.properties.put(SFSessionProperty.PRIVATE_KEY_FILE.getPropertyKey(), location);
     if (!Strings.isNullOrEmpty(password)) {
-      this.properties.put("private_key_file_pwd", password);
+      this.properties.put(SFSessionProperty.PRIVATE_KEY_FILE_PWD.getPropertyKey(), password);
     }
   }
 
   public void setTracing(String tracing) {
-    this.properties.put("tracing", tracing);
+    this.properties.put(SFSessionProperty.TRACING.getPropertyKey(), tracing);
+  }
+
+  protected Properties getProperties() {
+    return this.properties;
+  }
+
+  public void setAllowUnderscoresInHost(boolean allowUnderscoresInHost) {
+    this.properties.put(
+        SFSessionProperty.ALLOW_UNDERSCORES_IN_HOST.getPropertyKey(),
+        String.valueOf(allowUnderscoresInHost));
+  }
+
+  public void setDisableGcsDefaultCredentials(boolean isGcsDefaultCredentialsDisabled) {
+    this.properties.put(
+        SFSessionProperty.DISABLE_GCS_DEFAULT_CREDENTIALS.getPropertyKey(),
+        String.valueOf(isGcsDefaultCredentialsDisabled));
+  }
+
+  public void setDisableSamlURLCheck(boolean disableSamlURLCheck) {
+    this.properties.put(
+        SFSessionProperty.DISABLE_SAML_URL_CHECK.getPropertyKey(),
+        String.valueOf(disableSamlURLCheck));
+  }
+
+  public void setPasscode(String passcode) {
+    this.setAuthenticator(AUTHENTICATOR_USERNAME_PASSWORD_MFA);
+    this.properties.put(SFSessionProperty.PASSCODE.getPropertyKey(), passcode);
+  }
+
+  public void setPasscodeInPassword(boolean isPasscodeInPassword) {
+    this.properties.put(
+        SFSessionProperty.PASSCODE_IN_PASSWORD.getPropertyKey(),
+        String.valueOf(isPasscodeInPassword));
+    if (isPasscodeInPassword) {
+      this.setAuthenticator(AUTHENTICATOR_USERNAME_PASSWORD_MFA);
+    }
+  }
+
+  public void setDisableSocksProxy(boolean ignoreJvmSocksProxy) {
+    this.properties.put(
+        SFSessionProperty.DISABLE_SOCKS_PROXY.getPropertyKey(),
+        String.valueOf(ignoreJvmSocksProxy));
+  }
+
+  public void setNonProxyHosts(String nonProxyHosts) {
+    this.properties.put(SFSessionProperty.NON_PROXY_HOSTS.getPropertyKey(), nonProxyHosts);
+  }
+
+  public void setProxyHost(String proxyHost) {
+    this.properties.put(SFSessionProperty.PROXY_HOST.getPropertyKey(), proxyHost);
+  }
+
+  public void setProxyPassword(String proxyPassword) {
+    this.properties.put(SFSessionProperty.PROXY_PASSWORD.getPropertyKey(), proxyPassword);
+  }
+
+  public void setProxyPort(int proxyPort) {
+    this.properties.put(SFSessionProperty.PROXY_PORT.getPropertyKey(), Integer.toString(proxyPort));
+  }
+
+  public void setProxyProtocol(String proxyProtocol) {
+    this.properties.put(SFSessionProperty.PROXY_PROTOCOL.getPropertyKey(), proxyProtocol);
+  }
+
+  public void setProxyUser(String proxyUser) {
+    this.properties.put(SFSessionProperty.PROXY_USER.getPropertyKey(), proxyUser);
+  }
+
+  public void setUseProxy(boolean useProxy) {
+    this.properties.put(SFSessionProperty.USE_PROXY.getPropertyKey(), String.valueOf(useProxy));
+  }
+
+  public void setNetworkTimeout(int networkTimeoutSeconds) {
+    this.properties.put(
+        SFSessionProperty.NETWORK_TIMEOUT.getPropertyKey(),
+        Integer.toString(networkTimeoutSeconds));
+  }
+
+  public void setQueryTimeout(int queryTimeoutSeconds) {
+    this.properties.put(
+        SFSessionProperty.QUERY_TIMEOUT.getPropertyKey(), Integer.toString(queryTimeoutSeconds));
+  }
+
+  public void setApplication(String application) {
+    this.properties.put(SFSessionProperty.APPLICATION.getPropertyKey(), application);
+  }
+
+  public void setClientConfigFile(String clientConfigFile) {
+    this.properties.put(SFSessionProperty.CLIENT_CONFIG_FILE.getPropertyKey(), clientConfigFile);
+  }
+
+  public void setEnablePatternSearch(boolean enablePatternSearch) {
+    this.properties.put(
+        SFSessionProperty.ENABLE_PATTERN_SEARCH.getPropertyKey(),
+        String.valueOf(enablePatternSearch));
+  }
+
+  public void setEnablePutGet(boolean enablePutGet) {
+    this.properties.put(
+        SFSessionProperty.ENABLE_PUT_GET.getPropertyKey(), String.valueOf(enablePutGet));
+  }
+
+  public void setArrowTreatDecimalAsInt(boolean treatDecimalAsInt) {
+    this.properties.put(
+        SFSessionProperty.JDBC_ARROW_TREAT_DECIMAL_AS_INT.getPropertyKey(),
+        String.valueOf(treatDecimalAsInt));
+  }
+
+  public void setMaxHttpRetries(int maxHttpRetries) {
+    this.properties.put(
+        SFSessionProperty.MAX_HTTP_RETRIES.getPropertyKey(), Integer.toString(maxHttpRetries));
+  }
+
+  public void setOcspFailOpen(boolean ocspFailOpen) {
+    this.properties.put(
+        SFSessionProperty.OCSP_FAIL_OPEN.getPropertyKey(), String.valueOf(ocspFailOpen));
+  }
+
+  public void setPutGetMaxRetries(int putGetMaxRetries) {
+    this.properties.put(
+        SFSessionProperty.PUT_GET_MAX_RETRIES.getPropertyKey(), Integer.toString(putGetMaxRetries));
+  }
+
+  public void setStringsQuotedForColumnDef(boolean stringsQuotedForColumnDef) {
+    this.properties.put(
+        SFSessionProperty.STRINGS_QUOTED.getPropertyKey(),
+        String.valueOf(stringsQuotedForColumnDef));
+  }
+
+  public void setEnableDiagnostics(boolean enableDiagnostics) {
+    this.properties.put(
+        SFSessionProperty.ENABLE_DIAGNOSTICS.getPropertyKey(), String.valueOf(enableDiagnostics));
+  }
+
+  public void setDiagnosticsAllowlistFile(String diagnosticsAllowlistFile) {
+    this.properties.put(
+        SFSessionProperty.DIAGNOSTICS_ALLOWLIST_FILE.getPropertyKey(), diagnosticsAllowlistFile);
   }
 }
diff --git a/src/test/java/net/snowflake/client/jdbc/ConnectionLatestIT.java b/src/test/java/net/snowflake/client/jdbc/ConnectionLatestIT.java
@@ -57,6 +57,7 @@
 import net.snowflake.client.core.ObjectMapperFactory;
 import net.snowflake.client.core.QueryStatus;
 import net.snowflake.client.core.SFSession;
+import net.snowflake.client.core.SFSessionProperty;
 import net.snowflake.client.core.SecurityUtil;
 import net.snowflake.client.core.SessionUtil;
 import net.snowflake.client.jdbc.telemetryOOB.TelemetryService;
@@ -1379,6 +1380,92 @@ public void testDataSourceOktaGenerates429StatusCode() throws Exception {
     }
   }
 
+  /** Test added in JDBC driver version > 3.16.1 */
+  @Test
+  public void testDataSourceSetters() {
+    Map<String, String> params = getConnectionParameters();
+    SnowflakeBasicDataSource ds = new SnowflakeBasicDataSource();
+
+    ds.setTracing("all");
+    ds.setApplication("application_name");
+    ds.setAccount(params.get("account"));
+    ds.setAuthenticator("snowflake");
+    ds.setArrowTreatDecimalAsInt(true);
+    ds.setAllowUnderscoresInHost(true);
+    ds.setClientConfigFile("/some/path/file.json");
+    ds.setDisableGcsDefaultCredentials(false);
+    ds.setDisableSamlURLCheck(false);
+    ds.setDisableSocksProxy(false);
+    ds.setEnablePatternSearch(true);
+    ds.setDatabaseName("DB_NAME");
+    ds.setEnablePutGet(false);
+    ds.setMaxHttpRetries(5);
+    ds.setNetworkTimeout(10);
+    ds.setOcspFailOpen(false);
+    ds.setProxyHost("proxyHost.com");
+    ds.setProxyPort(8080);
+    ds.setProxyProtocol("http");
+    ds.setProxyUser("proxyUser");
+    ds.setProxyPassword("proxyPassword");
+    ds.setPutGetMaxRetries(3);
+    ds.setStringsQuotedForColumnDef(true);
+    ds.setEnableDiagnostics(true);
+    ds.setDiagnosticsAllowlistFile("/some/path/allowlist.json");
+
+    Properties props = ds.getProperties();
+    assertEquals(params.get("account"), props.get("account"));
+    assertEquals("snowflake", props.get("authenticator"));
+    assertEquals("all", props.get("tracing"));
+    assertEquals("application_name", props.get(SFSessionProperty.APPLICATION.getPropertyKey()));
+    assertEquals("snowflake", props.get(SFSessionProperty.AUTHENTICATOR.getPropertyKey()));
+    assertEquals(
+        "true", props.get(SFSessionProperty.JDBC_ARROW_TREAT_DECIMAL_AS_INT.getPropertyKey()));
+    assertEquals("true", props.get(SFSessionProperty.ALLOW_UNDERSCORES_IN_HOST.getPropertyKey()));
+    assertEquals(
+        "/some/path/file.json", props.get(SFSessionProperty.CLIENT_CONFIG_FILE.getPropertyKey()));
+    assertEquals(
+        "false", props.get(SFSessionProperty.DISABLE_GCS_DEFAULT_CREDENTIALS.getPropertyKey()));
+    assertEquals("false", props.get(SFSessionProperty.DISABLE_SAML_URL_CHECK.getPropertyKey()));
+    assertEquals("false", props.get(SFSessionProperty.DISABLE_SOCKS_PROXY.getPropertyKey()));
+    assertEquals("true", props.get(SFSessionProperty.ENABLE_PATTERN_SEARCH.getPropertyKey()));
+    assertEquals("DB_NAME", props.get(SFSessionProperty.DATABASE.getPropertyKey()));
+    assertEquals("false", props.get(SFSessionProperty.ENABLE_PUT_GET.getPropertyKey()));
+    assertEquals("5", props.get(SFSessionProperty.MAX_HTTP_RETRIES.getPropertyKey()));
+    assertEquals("10", props.get(SFSessionProperty.NETWORK_TIMEOUT.getPropertyKey()));
+    assertEquals("false", props.get(SFSessionProperty.OCSP_FAIL_OPEN.getPropertyKey()));
+    assertEquals("proxyHost.com", props.get(SFSessionProperty.PROXY_HOST.getPropertyKey()));
+    assertEquals("8080", props.get(SFSessionProperty.PROXY_PORT.getPropertyKey()));
+    assertEquals("http", props.get(SFSessionProperty.PROXY_PROTOCOL.getPropertyKey()));
+    assertEquals("proxyUser", props.get(SFSessionProperty.PROXY_USER.getPropertyKey()));
+    assertEquals("proxyPassword", props.get(SFSessionProperty.PROXY_PASSWORD.getPropertyKey()));
+    assertEquals("3", props.get(SFSessionProperty.PUT_GET_MAX_RETRIES.getPropertyKey()));
+    assertEquals("true", props.get(SFSessionProperty.STRINGS_QUOTED.getPropertyKey()));
+    assertEquals("true", props.get(SFSessionProperty.ENABLE_DIAGNOSTICS.getPropertyKey()));
+    assertEquals(
+        "/some/path/allowlist.json",
+        props.get(SFSessionProperty.DIAGNOSTICS_ALLOWLIST_FILE.getPropertyKey()));
+
+    ds.setOauthToken("a_token");
+    assertEquals("OAUTH", props.get(SFSessionProperty.AUTHENTICATOR.getPropertyKey()));
+    assertEquals("a_token", props.get(SFSessionProperty.TOKEN.getPropertyKey()));
+
+    ds.setPasscodeInPassword(true);
+    assertEquals("true", props.get(SFSessionProperty.PASSCODE_IN_PASSWORD.getPropertyKey()));
+    assertEquals(
+        "USERNAME_PASSWORD_MFA", props.get(SFSessionProperty.AUTHENTICATOR.getPropertyKey()));
+
+    ds.setPrivateKeyFile("key.p8", "pwd");
+    assertEquals("key.p8", props.get(SFSessionProperty.PRIVATE_KEY_FILE.getPropertyKey()));
+    assertEquals("pwd", props.get(SFSessionProperty.PRIVATE_KEY_FILE_PWD.getPropertyKey()));
+    assertEquals("SNOWFLAKE_JWT", props.get(SFSessionProperty.AUTHENTICATOR.getPropertyKey()));
+
+    ds.setPasscodeInPassword(false);
+    ds.setPasscode("a_passcode");
+    assertEquals("false", props.get(SFSessionProperty.PASSCODE_IN_PASSWORD.getPropertyKey()));
+    assertEquals(
+        "USERNAME_PASSWORD_MFA", props.get(SFSessionProperty.AUTHENTICATOR.getPropertyKey()));
+    assertEquals("a_passcode", props.get(SFSessionProperty.PASSCODE.getPropertyKey()));
+  }
   /**
    * SNOW-1465374: For TIMESTAMP_LTZ we were returning timestamps without timezone when scale was
    * set e.g. to 6 in Arrow format The problem wasn't visible when calling getString, but was