SNOW-732360: Allow connection caching to be disabled by user (#1845)

snowflakedb · Aug 29, 2024 · 3ede535 · 3ede535
1 parent 41582e8
commit 3ede535
Show file tree

Hide file tree

Showing 8 changed files with 112 additions and 2 deletions.
diff --git a/src/main/java/net/snowflake/client/core/SFLoginInput.java b/src/main/java/net/snowflake/client/core/SFLoginInput.java
@@ -51,6 +51,8 @@ public class SFLoginInput {
 
   private boolean disableConsoleLogin = true;
   private boolean disableSamlURLCheck = false;
+  private boolean enableClientStoreTemporaryCredential;
+  private boolean enableClientRequestMfaToken;
 
   private Duration browserResponseTimeout;
 
@@ -464,4 +466,23 @@ String getHostFromServerUrl() throws SFException {
     }
     return url.getHost();
   }
+
+  boolean isEnableClientStoreTemporaryCredential() {
+    return enableClientStoreTemporaryCredential;
+  }
+
+  SFLoginInput setEnableClientStoreTemporaryCredential(
+      boolean enableClientStoreTemporaryCredential) {
+    this.enableClientStoreTemporaryCredential = enableClientStoreTemporaryCredential;
+    return this;
+  }
+
+  boolean isEnableClientRequestMfaToken() {
+    return enableClientRequestMfaToken;
+  }
+
+  SFLoginInput setEnableClientRequestMfaToken(boolean enableClientRequestMfaToken) {
+    this.enableClientRequestMfaToken = enableClientRequestMfaToken;
+    return this;
+  }
 }
diff --git a/src/main/java/net/snowflake/client/core/SFSession.java b/src/main/java/net/snowflake/client/core/SFSession.java
@@ -142,6 +142,9 @@ public class SFSession extends SFBaseSession {
    */
   private int retryTimeout = 300;
 
+  private boolean enableClientStoreTemporaryCredential = true;
+  private boolean enableClientRequestMfaToken = true;
+
   /**
    * Max timeout for external browser authentication in seconds
    *
@@ -522,6 +525,18 @@ public void addSFSessionProperty(String propertyName, Object propertyValue) thro
           }
           break;
 
+        case ENABLE_CLIENT_STORE_TEMPORARY_CREDENTIAL:
+          if (propertyValue != null) {
+            enableClientStoreTemporaryCredential = getBooleanValue(propertyValue);
+          }
+          break;
+
+        case ENABLE_CLIENT_REQUEST_MFA_TOKEN:
+          if (propertyValue != null) {
+            enableClientRequestMfaToken = getBooleanValue(propertyValue);
+          }
+          break;
+
         default:
           break;
       }
@@ -677,6 +692,8 @@ public synchronized void open() throws SFException, SnowflakeSQLException {
                 ? getBooleanValue(
                     connectionPropertiesMap.get(SFSessionProperty.DISABLE_SAML_URL_CHECK))
                 : false)
+        .setEnableClientStoreTemporaryCredential(enableClientStoreTemporaryCredential)
+        .setEnableClientRequestMfaToken(enableClientRequestMfaToken)
         .setBrowserResponseTimeout(browserResponseTimeout);
 
     logger.info(

diff --git a/src/main/java/net/snowflake/client/core/SFSessionProperty.java b/src/main/java/net/snowflake/client/core/SFSessionProperty.java
@@ -106,6 +106,10 @@ public enum SFSessionProperty {
 
   BROWSER_RESPONSE_TIMEOUT("BROWSER_RESPONSE_TIMEOUT", false, Integer.class),
 
+  ENABLE_CLIENT_STORE_TEMPORARY_CREDENTIAL("clientStoreTemporaryCredential", false, Boolean.class),
+
+  ENABLE_CLIENT_REQUEST_MFA_TOKEN("clientRequestMfaToken", false, Boolean.class),
+
   HTTP_CLIENT_CONNECTION_TIMEOUT("HTTP_CLIENT_CONNECTION_TIMEOUT", false, Integer.class),
 
   HTTP_CLIENT_SOCKET_TIMEOUT("HTTP_CLIENT_SOCKET_TIMEOUT", false, Integer.class);

diff --git a/src/main/java/net/snowflake/client/core/SessionUtil.java b/src/main/java/net/snowflake/client/core/SessionUtil.java
@@ -278,7 +278,8 @@ static SFLoginOutput openSession(
           "missing token or password for opening session");
     }
     if (authenticator.equals(ClientAuthnDTO.AuthenticatorType.EXTERNALBROWSER)) {
-      if (Constants.getOS() == Constants.OS.MAC || Constants.getOS() == Constants.OS.WINDOWS) {
+      if ((Constants.getOS() == Constants.OS.MAC || Constants.getOS() == Constants.OS.WINDOWS)
+          && loginInput.isEnableClientStoreTemporaryCredential()) {
         // force to set the flag for Mac/Windows users
         loginInput.getSessionParameters().put(CLIENT_STORE_TEMPORARY_CREDENTIAL, true);
       } else {
@@ -299,7 +300,8 @@ static SFLoginOutput openSession(
     }
 
     if (authenticator.equals(ClientAuthnDTO.AuthenticatorType.USERNAME_PASSWORD_MFA)) {
-      if (Constants.getOS() == Constants.OS.MAC || Constants.getOS() == Constants.OS.WINDOWS) {
+      if ((Constants.getOS() == Constants.OS.MAC || Constants.getOS() == Constants.OS.WINDOWS)
+          && loginInput.isEnableClientRequestMfaToken()) {
         loginInput.getSessionParameters().put(CLIENT_REQUEST_MFA_TOKEN, true);
       }
     }

diff --git a/src/main/java/net/snowflake/client/jdbc/SnowflakeBasicDataSource.java b/src/main/java/net/snowflake/client/jdbc/SnowflakeBasicDataSource.java
@@ -394,6 +394,21 @@ public void setGetDateUseNullTimezone(Boolean getDateUseNullTimezone) {
     this.properties.put("JDBC_GET_DATE_USE_NULL_TIMEZONE", getDateUseNullTimezone);
   }
 
+  public void setEnableClientRequestMfaToken(boolean enableClientRequestMfaToken) {
+    this.setAuthenticator(AUTHENTICATOR_USERNAME_PASSWORD_MFA);
+    this.properties.put(
+        SFSessionProperty.ENABLE_CLIENT_REQUEST_MFA_TOKEN.getPropertyKey(),
+        enableClientRequestMfaToken);
+  }
+
+  public void setEnableClientStoreTemporaryCredential(
+      boolean enableClientStoreTemporaryCredential) {
+    this.setAuthenticator(AUTHENTICATOR_EXTERNAL_BROWSER);
+    this.properties.put(
+        SFSessionProperty.ENABLE_CLIENT_STORE_TEMPORARY_CREDENTIAL.getPropertyKey(),
+        enableClientStoreTemporaryCredential);
+  }
+
   public void setBrowserResponseTimeout(int seconds) {
     this.setAuthenticator(AUTHENTICATOR_EXTERNAL_BROWSER);
     this.properties.put("BROWSER_RESPONSE_TIMEOUT", Integer.toString(seconds));

diff --git a/src/test/java/net/snowflake/client/core/SessionUtilExternalBrowserTest.java b/src/test/java/net/snowflake/client/core/SessionUtilExternalBrowserTest.java
@@ -22,6 +22,8 @@
 import java.net.URI;
 import java.net.URISyntaxException;
 import java.nio.charset.StandardCharsets;
+import java.sql.Connection;
+import java.sql.ResultSet;
 import java.util.Map;
 import net.snowflake.client.AbstractDriverIT;
 import net.snowflake.client.jdbc.SnowflakeBasicDataSource;
@@ -243,6 +245,27 @@ private SFLoginInput initMockLoginInput() {
     return loginInput;
   }
 
+  // Run this test manually to test disabling storing temporary credetials with external browser
+  // auth. This is valid for versions after 3.18.0.
+  @Test
+  @Ignore
+  public void testEnableClientStoreTemporaryCredential() throws Exception {
+    Map<String, String> params = AbstractDriverIT.getConnectionParameters();
+    SnowflakeBasicDataSource ds = new SnowflakeBasicDataSource();
+    ds.setServerName(params.get("host"));
+    ds.setAccount(params.get("account"));
+    ds.setPortNumber(Integer.parseInt(params.get("port")));
+    ds.setUser(params.get("user"));
+    ds.setEnableClientStoreTemporaryCredential(false);
+
+    for (int i = 0; i < 3; i++) {
+      try (Connection con = ds.getConnection();
+          ResultSet rs = con.createStatement().executeQuery("SELECT 1")) {
+        assertTrue(rs.next());
+      }
+    }
+  }
+
   // Run this test manually to confirm external browser timeout is working. When test runs it will
   // open a browser window for authentication, close the window, and you should get the expected
   // error message within the set timeout. Valid for driver versions after 3.18.0.

diff --git a/src/test/java/net/snowflake/client/core/SnowflakeMFACacheTest.java b/src/test/java/net/snowflake/client/core/SnowflakeMFACacheTest.java
@@ -20,12 +20,17 @@
 import java.io.StringWriter;
 import java.sql.Connection;
 import java.sql.DriverManager;
+import java.sql.ResultSet;
 import java.sql.SQLException;
+import java.util.Map;
 import java.util.Properties;
+import net.snowflake.client.AbstractDriverIT;
+import net.snowflake.client.jdbc.SnowflakeBasicDataSource;
 import net.snowflake.client.jdbc.SnowflakeSQLException;
 import org.apache.commons.io.IOUtils;
 import org.apache.http.client.methods.HttpPost;
 import org.junit.Assert;
+import org.junit.Ignore;
 import org.junit.Test;
 import org.mockito.MockedStatic;
 import org.mockito.Mockito;
@@ -327,4 +332,26 @@ public void testUnavailableLocalSecureStorage() throws SQLException {
       CredentialManager.getInstance().resetSecureStorageManager();
     }
   }
+
+  // Run this test manually to test disabling the client request MFA token. Use an MFA
+  // authentication enabled user. This is valid for versions after 3.18.0.
+  @Test
+  @Ignore
+  public void testEnableClientRequestMfaToken() throws SQLException {
+    Map<String, String> params = AbstractDriverIT.getConnectionParameters();
+    SnowflakeBasicDataSource ds = new SnowflakeBasicDataSource();
+    ds.setServerName(params.get("host"));
+    ds.setAccount(params.get("account"));
+    ds.setPortNumber(Integer.parseInt(params.get("port")));
+    ds.setUser(params.get("user"));
+    ds.setPassword(params.get("password"));
+    ds.setEnableClientRequestMfaToken(false);
+
+    for (int i = 0; i < 3; i++) {
+      try (Connection con = ds.getConnection();
+          ResultSet rs = con.createStatement().executeQuery("SELECT 1")) {
+        assertTrue(rs.next());
+      }
+    }
+  }
 }
diff --git a/src/test/java/net/snowflake/client/jdbc/ConnectionIT.java b/src/test/java/net/snowflake/client/jdbc/ConnectionIT.java
@@ -1022,6 +1022,7 @@ private Properties setCommonConnectionParameters(boolean validateDefaultParamete
     return kvMap2Properties(params, validateDefaultParameters);
   }
 
+  // TODO: This is a temporary Ignore to unblock merging PRs until ORG account is unlocked.
   @Ignore
   @Test
   public void testFailOverOrgAccount() throws SQLException {