Code review changes

snowflakedb · Aug 14, 2024 · 663949d · 663949d
1 parent c665821
commit 663949d
Show file tree

Hide file tree

Showing 3 changed files with 14 additions and 16 deletions.
diff --git a/src/main/java/net/snowflake/client/jdbc/SnowflakeBasicDataSource.java b/src/main/java/net/snowflake/client/jdbc/SnowflakeBasicDataSource.java
@@ -388,13 +388,17 @@ public void setGetDateUseNullTimezone(Boolean getDateUseNullTimezone) {
 
   public void setEnableClientRequestMfaToken(boolean enableClientRequestMfaToken) {
     this.setAuthenticator(AUTHENTICATOR_USERNAME_PASSWORD_MFA);
-    this.properties.put("clientRequestMfaToken", enableClientRequestMfaToken);
+    this.properties.put(
+        SFSessionProperty.ENABLE_CLIENT_REQUEST_MFA_TOKEN.getPropertyKey(),
+        enableClientRequestMfaToken);
   }
 
   public void setEnableClientStoreTemporaryCredential(
       boolean enableClientStoreTemporaryCredential) {
     this.setAuthenticator(AUTHENTICATOR_EXTERNAL_BROWSER);
-    this.properties.put("clientStoreTemporaryCredential", enableClientStoreTemporaryCredential);
+    this.properties.put(
+        SFSessionProperty.ENABLE_CLIENT_STORE_TEMPORARY_CREDENTIAL.getPropertyKey(),
+        enableClientStoreTemporaryCredential);
   }
 
   public void setBrowserResponseTimeout(int seconds) {

diff --git a/src/test/java/net/snowflake/client/core/SessionUtilExternalBrowserTest.java b/src/test/java/net/snowflake/client/core/SessionUtilExternalBrowserTest.java
@@ -24,7 +24,6 @@
 import java.nio.charset.StandardCharsets;
 import java.sql.Connection;
 import java.sql.ResultSet;
-import java.sql.Statement;
 import java.util.Map;
 import net.snowflake.client.AbstractDriverIT;
 import net.snowflake.client.jdbc.SnowflakeBasicDataSource;
@@ -259,12 +258,10 @@ public void testEnableClientStoreTemporaryCredential() throws Exception {
     ds.setUser(params.get("user"));
     ds.setEnableClientStoreTemporaryCredential(false);
 
-    try (Connection con = ds.getConnection()) {
-      Statement stmt = con.createStatement();
-      for (int i = 0; i < 3; i++) {
-        try (ResultSet rs = stmt.executeQuery("SELECT CURRENT_USER()")) {
-          assertTrue(rs.next());
-        }
+    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/core/SnowflakeMFACacheTest.java b/src/test/java/net/snowflake/client/core/SnowflakeMFACacheTest.java
@@ -22,7 +22,6 @@
 import java.sql.DriverManager;
 import java.sql.ResultSet;
 import java.sql.SQLException;
-import java.sql.Statement;
 import java.util.Map;
 import java.util.Properties;
 import net.snowflake.client.AbstractDriverIT;
@@ -348,12 +347,10 @@ public void testEnableClientRequestMfaToken() throws SQLException {
     ds.setPassword(params.get("password"));
     ds.setEnableClientRequestMfaToken(false);
 
-    try (Connection con = ds.getConnection()) {
-      for (int i = 0; i < 3; i++) {
-        try (Statement stmt = con.createStatement();
-            ResultSet rs = stmt.executeQuery("SELECT CURRENT_USER()")) {
-          assertTrue(rs.next());
-        }
+    for (int i = 0; i < 3; i++) {
+      try (Connection con = ds.getConnection();
+          ResultSet rs = con.createStatement().executeQuery("SELECT 1")) {
+        assertTrue(rs.next());
       }
     }
   }