Added retry logic for authentication with Okta

CHANGELOG.rst

@@ -1,59 +1,59 @@
 **JDBC Driver 3.15.0**
 
-- \||Please Refer to Release Notes at https://docs.snowflake.com/en/release-notes/clients-drivers/jdbc
+- \||Please Refer to Release Notes at https://community.snowflake.com/s/article/JDBC-Driver-Release-Notes
 
 **JDBC Driver 3.14.5**
 
-- \||Please Refer to Release Notes at https://docs.snowflake.com/en/release-notes/clients-drivers/jdbc
+- \||Please Refer to Release Notes at https://community.snowflake.com/s/article/JDBC-Driver-Release-Notes
 
 **JDBC Driver 3.14.4**
 
-- \||Please Refer to Release Notes at https://docs.snowflake.com/en/release-notes/clients-drivers/jdbc
+- \||Please Refer to Release Notes at https://community.snowflake.com/s/article/JDBC-Driver-Release-Notes
 
 **JDBC Driver 3.14.3**
 
-- \||Please Refer to Release Notes at https://docs.snowflake.com/en/release-notes/clients-drivers/jdbc
+- \||Please Refer to Release Notes at https://community.snowflake.com/s/article/JDBC-Driver-Release-Notes
 
 **JDBC Driver 3.14.2**
 
-- \||Please Refer to Release Notes at https://docs.snowflake.com/en/release-notes/clients-drivers/jdbc
+- \||Please Refer to Release Notes at https://community.snowflake.com/s/article/JDBC-Driver-Release-Notes
 
 **JDBC Driver 3.14.1**
 
-- \||Please Refer to Release Notes at https://docs.snowflake.com/en/release-notes/clients-drivers/jdbc
+- \||Please Refer to Release Notes at https://community.snowflake.com/s/article/JDBC-Driver-Release-Notes
 
 **JDBC Driver 3.14.0**
 
-- \||Please Refer to Release Notes at https://docs.snowflake.com/en/release-notes/clients-drivers/jdbc
+- \||Please Refer to Release Notes at https://community.snowflake.com/s/article/JDBC-Driver-Release-Notes
 
 **JDBC Driver 3.13.33**
 
-- \|| Please Refer to Release Notes at https://docs.snowflake.com/en/release-notes/clients-drivers/jdbc
+- \|| Please Refer to Release Notes at https://community.snowflake.com/s/article/JDBC-Driver-Release-Notes
 
 **JDBC Driver 3.13.32**
 
 
-- \|| Please Refer to Release Notes at https://docs.snowflake.com/en/release-notes/clients-drivers/jdbc
+- \|| Please Refer to Release Notes at https://community.snowflake.com/s/article/JDBC-Driver-Release-Notes
 
 **JDBC Driver 3.13.31**
 
-- \|| Please Refer to Release Notes at https://docs.snowflake.com/en/release-notes/clients-drivers/jdbc
+- \|| Please Refer to Release Notes at https://community.snowflake.com/s/article/JDBC-Driver-Release-Notes
 
 **JDBC Driver 3.13.30**
 
-- \|| Please Refer to Release Notes at https://docs.snowflake.com/en/release-notes/clients-drivers/jdbc
+- \|| Please Refer to Release Notes at https://community.snowflake.com/s/article/JDBC-Driver-Release-Notes
 
 **JDBC Driver 3.13.29**
 
-- \|| Please Refer to Release Notes at https://docs.snowflake.com/en/release-notes/clients-drivers/jdbc
+- \|| Please Refer to Release Notes at https://community.snowflake.com/s/article/JDBC-Driver-Release-Notes
 
 **JDBC Driver 3.13.28**
 
-- \|| Please Refer to Release Notes at https://docs.snowflake.com/en/release-notes/clients-drivers/jdbc
+- \|| Please Refer to Release Notes at https://community.snowflake.com/s/article/JDBC-Driver-Release-Notes
 
 **JDBC Driver 3.13.27**
 
-- \|| Please Refer to Release Notes at https://docs.snowflake.com/en/release-notes/clients-drivers/jdbc
+- \|| Please Refer to Release Notes at https://community.snowflake.com/s/article/JDBC-Driver-Release-Notes
 
 **JDBC Driver 3.13.26**
 

parent-pom.xml

@@ -38,7 +38,6 @@
     <commons.pool.version>1.5.4</commons.pool.version>
     <c3p0.version>0.9.5.4</c3p0.version>
     <google.api.grpc.version>2.22.0</google.api.grpc.version>
-    <google.auth.library.oauth2.http.version>1.19.0</google.auth.library.oauth2.http.version>
     <google.cloud.core.version>2.21.0</google.cloud.core.version>
     <google.cloud.storage.version>2.22.6</google.cloud.storage.version>
     <google.code.gson.version>2.10.1</google.code.gson.version>
@@ -162,11 +161,6 @@
         <artifactId>proto-google-common-protos</artifactId>
         <version>${google.api.grpc.version}</version>
       </dependency>
-      <dependency>
-        <groupId>com.google.auth</groupId>
-        <artifactId>google-auth-library-oauth2-http</artifactId>
-        <version>${google.auth.library.oauth2.http.version}</version>
-      </dependency>
       <dependency>
         <groupId>com.google.cloud</groupId>
         <artifactId>google-cloud-core</artifactId>
@@ -520,10 +514,6 @@
       <groupId>com.google.api</groupId>
       <artifactId>gax</artifactId>
     </dependency>
-    <dependency>
-      <groupId>com.google.auth</groupId>
-      <artifactId>google-auth-library-oauth2-http</artifactId>
-    </dependency>
     <dependency>
       <groupId>com.google.cloud</groupId>
       <artifactId>google-cloud-core</artifactId>

src/main/java/net/snowflake/client/core/JsonSqlInput.java

@@ -190,9 +190,11 @@ public Timestamp readTimestamp(TimeZone tz) throws SQLException {
           int columnType = ColumnTypeHelper.getColumnType(fieldMetadata.getType(), session);
           int columnSubType = fieldMetadata.getType();
           int scale = fieldMetadata.getScale();
+
           Timestamp result =
               SqlInputTimestampUtil.getTimestampFromType(
                   columnSubType, (String) value, session, sessionTimeZone, tz);
+
           if (result != null) {
             return result;
           }

src/main/java/net/snowflake/client/core/SFArrowResultSet.java

@@ -13,7 +13,6 @@
 import java.io.IOException;
 import java.math.BigDecimal;
 import java.math.RoundingMode;
-import java.sql.Array;
 import java.sql.Date;
 import java.sql.SQLException;
 import java.sql.SQLInput;

src/main/java/net/snowflake/client/core/SFBaseResultSet.java

@@ -11,7 +11,6 @@
 import com.fasterxml.jackson.databind.ObjectMapper;
 import com.fasterxml.jackson.databind.node.ArrayNode;
 import java.math.BigDecimal;
-import java.sql.Array;
 import java.sql.Date;
 import java.sql.SQLException;
 import java.sql.SQLInput;

src/main/java/net/snowflake/client/core/SFBaseSession.java

@@ -137,9 +137,6 @@ public abstract class SFBaseSession {
   // we need to allow for it to maintain backwards compatibility.
   private boolean enablePatternSearch = true;
 
-  /** Disable lookup for default credentials by GCS library */
-  private boolean disableGcsDefaultCredentials = false;
-
   private Map<String, Object> commonParameters;
 
   private boolean isJdbcArrowTreatDecimalAsInt = true;
@@ -748,14 +745,6 @@ public void setEnablePatternSearch(boolean enablePatternSearch) {
     this.enablePatternSearch = enablePatternSearch;
   }
 
-  public boolean getDisableGcsDefaultCredentials() {
-    return disableGcsDefaultCredentials;
-  }
-
-  public void setDisableGcsDefaultCredentials(boolean disableGcsDefaultCredentials) {
-    this.disableGcsDefaultCredentials = disableGcsDefaultCredentials;
-  }
-
   public int getClientResultChunkSize() {
     return clientResultChunkSize;
   }

src/main/java/net/snowflake/client/core/SFJsonResultSet.java

@@ -7,7 +7,6 @@
 import com.fasterxml.jackson.core.JsonProcessingException;
 import com.fasterxml.jackson.databind.JsonNode;
 import java.math.BigDecimal;
-import java.sql.Array;
 import java.sql.Date;
 import java.sql.SQLInput;
 import java.sql.Time;
@@ -88,6 +87,7 @@ public Object getObject(int columnIndex) throws SFException {
         } else {
           throw new SFException(ErrorCode.FEATURE_UNSUPPORTED, "data type: " + type);
         }
+
       case Types.ARRAY:
         if (StructureTypeHelper.isStructureTypeEnabled()) {
           return getArray(columnIndex);

src/main/java/net/snowflake/client/core/SFResultSet.java

@@ -155,7 +155,22 @@ public SFResultSet(
       Telemetry telemetryClient,
       boolean sortResult)
       throws SQLException {
-    super(resultSetSerializable.getTimeZone(), new Converters(session, resultSetSerializable));
+    super(
+        resultSetSerializable.getTimeZone(),
+        new Converters(
+            resultSetSerializable.getTimeZone(),
+            session,
+            resultSetSerializable.getResultVersion(),
+            resultSetSerializable.isHonorClientTZForTimestampNTZ(),
+            resultSetSerializable.getTreatNTZAsUTC(),
+            resultSetSerializable.getUseSessionTimezone(),
+            resultSetSerializable.getFormatDateWithTimeZone(),
+            resultSetSerializable.getBinaryFormatter(),
+            resultSetSerializable.getDateFormatter(),
+            resultSetSerializable.getTimeFormatter(),
+            resultSetSerializable.getTimestampNTZFormatter(),
+            resultSetSerializable.getTimestampLTZFormatter(),
+            resultSetSerializable.getTimestampTZFormatter()));
     this.resultSetSerializable = resultSetSerializable;
     this.columnCount = 0;
     this.sortResult = sortResult;

src/main/java/net/snowflake/client/core/SFResultSetMetaData.java

@@ -473,8 +473,7 @@ public List<Boolean> getIsAutoIncrementList() {
     return isAutoIncrementList;
   }
 
-  @SnowflakeJdbcInternalApi
-  public List<SnowflakeColumnMetadata> getColumnMetadata() {
+  List<SnowflakeColumnMetadata> getColumnMetadata() {
     return columnMetadata;
   }
 }

src/main/java/net/snowflake/client/core/SFSession.java

@@ -469,11 +469,6 @@ public void addSFSessionProperty(String propertyName, Object propertyValue) thro
             setEnablePatternSearch(getBooleanValue(propertyValue));
           }
           break;
-        case DISABLE_GCS_DEFAULT_CREDENTIALS:
-          if (propertyValue != null) {
-            setDisableGcsDefaultCredentials(getBooleanValue(propertyValue));
-          }
-          break;
 
         case JDBC_ARROW_TREAT_DECIMAL_AS_INT:
           if (propertyValue != null) {

src/main/java/net/snowflake/client/core/SessionUtil.java

@@ -401,7 +401,21 @@ private static SFLoginOutput newSession(
         }
       } else if (authenticatorType == ClientAuthnDTO.AuthenticatorType.OKTA) {
         // okta authenticator v1
-        tokenOrSamlResponse = getSamlResponseUsingOkta(loginInput);
+        while (true) {
+          try {
+            tokenOrSamlResponse = getSamlResponseUsingOkta(loginInput);
+          } catch (SnowflakeSQLException ex) {
+            // 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.");
+              continue;
+            } else {
+              throw ex;
+            }
+          }
+          break;
+        }
       } else if (authenticatorType == ClientAuthnDTO.AuthenticatorType.SNOWFLAKE_JWT) {
         SessionUtilKeyPair s =
             new SessionUtilKeyPair(
@@ -651,16 +665,31 @@ private static SFLoginOutput newSession(
                   loginInput.getHttpClientSettingsKey());
         } catch (SnowflakeSQLException ex) {
           if (ex.getErrorCode() == ErrorCode.AUTHENTICATOR_REQUEST_TIMEOUT.getMessageCode()) {
-            if (authenticatorType == ClientAuthnDTO.AuthenticatorType.SNOWFLAKE_JWT) {
-              SessionUtilKeyPair s =
-                  new SessionUtilKeyPair(
-                      loginInput.getPrivateKey(),
-                      loginInput.getPrivateKeyFile(),
-                      loginInput.getPrivateKeyFilePwd(),
-                      loginInput.getAccountName(),
-                      loginInput.getUserName());
-
-              data.put(ClientAuthnParameter.TOKEN.name(), s.issueJwtToken());
+            if (authenticatorType == ClientAuthnDTO.AuthenticatorType.SNOWFLAKE_JWT
+                || authenticatorType == ClientAuthnDTO.AuthenticatorType.OKTA) {
+
+              if (authenticatorType == ClientAuthnDTO.AuthenticatorType.SNOWFLAKE_JWT) {
+                SessionUtilKeyPair s =
+                    new SessionUtilKeyPair(
+                        loginInput.getPrivateKey(),
+                        loginInput.getPrivateKeyFile(),
+                        loginInput.getPrivateKeyFilePwd(),
+                        loginInput.getAccountName(),
+                        loginInput.getUserName());
+
+                data.put(ClientAuthnParameter.TOKEN.name(), s.issueJwtToken());
+              } else if (authenticatorType == ClientAuthnDTO.AuthenticatorType.OKTA) {
+                logger.debug("Retrieve new token for Okta authentication.");
+                // If we need to retry, we need to get a new Okta token
+                tokenOrSamlResponse = getSamlResponseUsingOkta(loginInput);
+                data.put(ClientAuthnParameter.RAW_SAML_RESPONSE.name(), tokenOrSamlResponse);
+                authnData.setData(data);
+                String updatedJson = mapper.writeValueAsString(authnData);
+
+                StringEntity updatedInput = new StringEntity(updatedJson, StandardCharsets.UTF_8);
+                updatedInput.setContentType("application/json");
+                postRequest.setEntity(updatedInput);
+              }
 
               long elapsedSeconds = ex.getElapsedSeconds();
 
@@ -690,7 +719,8 @@ private static SFLoginOutput newSession(
                 }
               }
 
-              // JWT renew should not count as a retry, so we pass back the current retry count.
+              // JWT or Okta renew should not count as a retry, so we pass back the current retry
+              // count.
               retryCount = ex.getRetryCount();
 
               continue;