code review changes

snowflakedb · Apr 11, 2024 · 28ee513 · 28ee513
1 parent efff532
commit 28ee513
Show file tree

Hide file tree

Showing 4 changed files with 39 additions and 44 deletions.
diff --git a/src/main/java/net/snowflake/client/core/SessionUtil.java b/src/main/java/net/snowflake/client/core/SessionUtil.java
@@ -1390,7 +1390,7 @@ private static String getSamlResponseUsingOkta(SFLoginInput loginInput)
         // This error gets thrown if the okta request encountered a retry-able error that
         // requires getting a new one-time token.
         if (ex.getErrorCode() == ErrorCode.AUTHENTICATOR_REQUEST_TIMEOUT.getMessageCode()) {
-          logger.debug("Failed to get Okta SAML response. Retrying.");
+          logger.debug("Failed to get Okta SAML response. Retrying without changing retry count.");
         } else {
           throw ex;
         }

diff --git a/src/test/java/net/snowflake/client/core/SessionUtilLatestIT.java b/src/test/java/net/snowflake/client/core/SessionUtilLatestIT.java
@@ -407,7 +407,7 @@ private SFLoginInput createOktaLoginInput() {
   }
 
   // Testing retry with Okta calls the service to get a new unique token. This is valid after
-  // version 3.15.0.
+  // version 3.15.1.
   @Test
   public void testOktaAuthRetry() throws Throwable {
     SFLoginInput loginInput = createOktaLoginInput();

diff --git a/src/test/java/net/snowflake/client/jdbc/ConnectionIT.java b/src/test/java/net/snowflake/client/jdbc/ConnectionIT.java
@@ -389,40 +389,6 @@ public void testDataSourceOktaSerialization() throws Exception {
     con.close();
   }
 
-  @Test
-  @Ignore
-  public void testDataSourceOktaGenerates429StatusCode() throws Exception {
-    // test with username/password authentication
-    // set up DataSource object and ensure connection works
-    Map<String, String> params = getConnectionParameters();
-    SnowflakeBasicDataSource ds = new SnowflakeBasicDataSource();
-    ds.setServerName(params.get("host"));
-    ds.setSsl("on".equals(params.get("ssl")));
-    ds.setAccount(params.get("account"));
-    ds.setPortNumber(Integer.parseInt(params.get("port")));
-    ds.setUser(params.get("ssoUser"));
-    ds.setPassword(params.get("ssoPassword"));
-    ds.setAuthenticator("<okta address>");
-    Runnable r =
-        () -> {
-          try {
-            ds.getConnection();
-          } catch (SQLException e) {
-            throw new RuntimeException(e);
-          }
-        };
-    List<Thread> threadList = new ArrayList<>();
-    for (int i = 0;
-        i < 30;
-        ++i) { // https://docs.snowflake.com/en/user-guide/admin-security-fed-auth-use#http-429-errors
-      threadList.add(new Thread(r));
-    }
-    threadList.forEach(Thread::start);
-    for (Thread thread : threadList) {
-      thread.join();
-    }
-  }
-
   @Test
   @ConditionalIgnore(condition = RunningOnGithubAction.class)
   public void testConnectUsingKeyPair() throws Exception {

diff --git a/src/test/java/net/snowflake/client/jdbc/ConnectionLatestIT.java b/src/test/java/net/snowflake/client/jdbc/ConnectionLatestIT.java
@@ -37,10 +37,7 @@
 import java.sql.SQLException;
 import java.sql.Statement;
 import java.time.Duration;
-import java.util.Collections;
-import java.util.Enumeration;
-import java.util.Map;
-import java.util.Properties;
+import java.util.*;
 import java.util.concurrent.TimeUnit;
 import net.snowflake.client.ConditionalIgnoreRule;
 import net.snowflake.client.RunningOnGithubAction;
@@ -59,10 +56,7 @@
 import org.apache.http.client.methods.HttpPost;
 import org.apache.http.client.utils.URIBuilder;
 import org.apache.http.entity.StringEntity;
-import org.junit.After;
-import org.junit.Before;
-import org.junit.Rule;
-import org.junit.Test;
+import org.junit.*;
 import org.junit.experimental.categories.Category;
 import org.junit.rules.TemporaryFolder;
 
@@ -1152,4 +1146,39 @@ public void testIsAsyncSession() throws SQLException, InterruptedException {
       assertFalse(snowflakeConnection.getSfSession().isAsyncSession());
     }
   }
+
+  // Test for regenerating okta one-time token for versions > 3.15.1
+  @Test
+  @Ignore
+  public void testDataSourceOktaGenerates429StatusCode() throws Exception {
+    // test with username/password authentication
+    // set up DataSource object and ensure connection works
+    Map<String, String> params = getConnectionParameters();
+    SnowflakeBasicDataSource ds = new SnowflakeBasicDataSource();
+    ds.setServerName(params.get("host"));
+    ds.setSsl("on".equals(params.get("ssl")));
+    ds.setAccount(params.get("account"));
+    ds.setPortNumber(Integer.parseInt(params.get("port")));
+    ds.setUser(params.get("ssoUser"));
+    ds.setPassword(params.get("ssoPassword"));
+    ds.setAuthenticator("<okta address>");
+    Runnable r =
+            () -> {
+              try {
+                ds.getConnection();
+              } catch (SQLException e) {
+                throw new RuntimeException(e);
+              }
+            };
+    List<Thread> threadList = new ArrayList<>();
+    for (int i = 0;
+         i < 30;
+         ++i) { // https://docs.snowflake.com/en/user-guide/admin-security-fed-auth-use#http-429-errors
+      threadList.add(new Thread(r));
+    }
+    threadList.forEach(Thread::start);
+    for (Thread thread : threadList) {
+      thread.join();
+    }
+  }
 }