SNOW-1853185: Pom test

TestOnly/pom.xml

@@ -22,6 +22,7 @@
     <surefire.version>3.5.1</surefire.version>
     <mockito.version>3.5.6</mockito.version>
     <netty.version>4.1.118.Final</netty.version>
+    <apache.httpclient.version>4.5.14</apache.httpclient.version>
     <shadeBase>net.snowflake.client.jdbc.internal</shadeBase>
   </properties>
 
@@ -31,6 +32,11 @@
       <artifactId>netty-common</artifactId>
       <version>${netty.version}</version>
     </dependency>
+    <dependency>
+      <groupId>org.apache.httpcomponents</groupId>
+      <artifactId>httpclient</artifactId>
+      <version>${apache.httpclient.version}</version>
+    </dependency>
     <dependency>
       <groupId>io.netty</groupId>
       <artifactId>netty-buffer</artifactId>

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

@@ -32,6 +32,7 @@
 import javax.net.ssl.TrustManager;
 import net.snowflake.client.jdbc.ErrorCode;
 import net.snowflake.client.jdbc.RestRequest;
+import net.snowflake.client.jdbc.RetryContext;
 import net.snowflake.client.jdbc.SnowflakeDriver;
 import net.snowflake.client.jdbc.SnowflakeSQLException;
 import net.snowflake.client.jdbc.SnowflakeUtil;
@@ -618,7 +619,8 @@ static String executeRequestWithoutCookies(
         true, // guid? (do we need this?)
         false, // no retry on HTTP 403
         getHttpClient(ocspAndProxyKey),
-        new ExecTimeTelemetryData());
+        new ExecTimeTelemetryData(),
+        null);
   }
 
   /**
@@ -642,6 +644,39 @@ public static String executeGeneralRequest(
       int retryCount,
       HttpClientSettingsKey ocspAndProxyAndGzipKey)
       throws SnowflakeSQLException, IOException {
+    return executeGeneralRequest(
+        httpRequest,
+        retryTimeout,
+        authTimeout,
+        socketTimeout,
+        retryCount,
+        ocspAndProxyAndGzipKey,
+        null);
+  }
+
+  /**
+   * Executes an HTTP request for Snowflake.
+   *
+   * @param httpRequest HttpRequestBase
+   * @param retryTimeout retry timeout
+   * @param authTimeout authenticator specific timeout
+   * @param socketTimeout socket timeout (in ms)
+   * @param retryCount retry count for the request
+   * @param ocspAndProxyAndGzipKey OCSP mode and proxy settings for httpclient
+   * @param retryContext RetryContext used to customize retry handling functionality
+   * @return response
+   * @throws SnowflakeSQLException if Snowflake error occurs
+   * @throws IOException raises if a general IO error occurs
+   */
+  public static String executeGeneralRequest(
+      HttpRequestBase httpRequest,
+      int retryTimeout,
+      int authTimeout,
+      int socketTimeout,
+      int retryCount,
+      HttpClientSettingsKey ocspAndProxyAndGzipKey,
+      RetryContext retryContext)
+      throws SnowflakeSQLException, IOException {
     logger.debug("Executing general request");
     return executeRequest(
         httpRequest,
@@ -654,7 +689,8 @@ public static String executeGeneralRequest(
         false, // no retry parameter
         false, // no retry on HTTP 403
         ocspAndProxyAndGzipKey,
-        new ExecTimeTelemetryData());
+        new ExecTimeTelemetryData(),
+        retryContext);
   }
 
   /**
@@ -692,7 +728,8 @@ public static String executeGeneralRequest(
         true, // include request GUID
         false, // no retry on HTTP 403
         httpClient,
-        new ExecTimeTelemetryData());
+        new ExecTimeTelemetryData(),
+        null);
   }
 
   /**
@@ -726,6 +763,54 @@ public static String executeRequest(
       HttpClientSettingsKey ocspAndProxyKey,
       ExecTimeTelemetryData execTimeData)
       throws SnowflakeSQLException, IOException {
+    return executeRequest(
+        httpRequest,
+        retryTimeout,
+        authTimeout,
+        socketTimeout,
+        maxRetries,
+        injectSocketTimeout,
+        canceling,
+        includeRetryParameters,
+        retryOnHTTP403,
+        ocspAndProxyKey,
+        execTimeData,
+        null);
+  }
+
+  /**
+   * Executes an HTTP request for Snowflake.
+   *
+   * @param httpRequest HttpRequestBase
+   * @param retryTimeout retry timeout
+   * @param authTimeout authenticator timeout
+   * @param socketTimeout socket timeout (in ms)
+   * @param maxRetries retry count for the request
+   * @param injectSocketTimeout injecting socket timeout
+   * @param canceling canceling?
+   * @param includeRetryParameters whether to include retry parameters in retried requests
+   * @param retryOnHTTP403 whether to retry on HTTP 403 or not
+   * @param ocspAndProxyKey OCSP mode and proxy settings for httpclient
+   * @param execTimeData query execution time telemetry data object
+   * @param retryContext RetryContext used to customize retry handling functionality
+   * @return response
+   * @throws SnowflakeSQLException if Snowflake error occurs
+   * @throws IOException raises if a general IO error occurs
+   */
+  public static String executeRequest(
+      HttpRequestBase httpRequest,
+      int retryTimeout,
+      int authTimeout,
+      int socketTimeout,
+      int maxRetries,
+      int injectSocketTimeout,
+      AtomicBoolean canceling,
+      boolean includeRetryParameters,
+      boolean retryOnHTTP403,
+      HttpClientSettingsKey ocspAndProxyKey,
+      ExecTimeTelemetryData execTimeData,
+      RetryContext retryContext)
+      throws SnowflakeSQLException, IOException {
     boolean ocspEnabled = !(ocspAndProxyKey.getOcspMode().equals(OCSPMode.DISABLE_OCSP_CHECKS));
     logger.debug("Executing request with OCSP enabled: {}", ocspEnabled);
     execTimeData.setOCSPStatus(ocspEnabled);
@@ -742,7 +827,8 @@ public static String executeRequest(
         true, // include request GUID
         retryOnHTTP403,
         getHttpClient(ocspAndProxyKey),
-        execTimeData);
+        execTimeData,
+        retryContext);
   }
 
   /**
@@ -764,6 +850,7 @@ public static String executeRequest(
    * @param includeRequestGuid whether to include request_guid
    * @param retryOnHTTP403 whether to retry on HTTP 403
    * @param httpClient client object used to communicate with other machine
+   * @param retryContext RetryContext used to customize retry handling functionality
    * @return response in String
    * @throws SnowflakeSQLException if Snowflake error occurs
    * @throws IOException raises if a general IO error occurs
@@ -781,7 +868,8 @@ private static String executeRequestInternal(
       boolean includeRequestGuid,
       boolean retryOnHTTP403,
       CloseableHttpClient httpClient,
-      ExecTimeTelemetryData execTimeData)
+      ExecTimeTelemetryData execTimeData,
+      RetryContext retryContext)
       throws SnowflakeSQLException, IOException {
     // HttpRequest.toString() contains request URI. Scrub any credentials, if
     // present, before logging
@@ -815,7 +903,8 @@ private static String executeRequestInternal(
               includeRetryParameters,
               includeRequestGuid,
               retryOnHTTP403,
-              execTimeData);
+              execTimeData,
+              retryContext);
       if (logger.isDebugEnabled() && stopwatch != null) {
         stopwatch.stop();
       }

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

@@ -36,6 +36,7 @@
 import net.snowflake.client.core.auth.oauth.OAuthAccessTokenProviderFactory;
 import net.snowflake.client.core.auth.oauth.TokenResponseDTO;
 import net.snowflake.client.jdbc.ErrorCode;
+import net.snowflake.client.jdbc.RetryContext;
 import net.snowflake.client.jdbc.SnowflakeDriver;
 import net.snowflake.client.jdbc.SnowflakeReauthenticationRequest;
 import net.snowflake.client.jdbc.SnowflakeSQLException;
@@ -48,6 +49,7 @@
 import net.snowflake.client.log.SFLoggerFactory;
 import net.snowflake.client.util.SecretDetector;
 import net.snowflake.client.util.Stopwatch;
+import net.snowflake.client.util.ThrowingCallable;
 import net.snowflake.common.core.SqlState;
 import org.apache.http.HttpHeaders;
 import org.apache.http.client.config.RequestConfig;
@@ -74,6 +76,7 @@ public class SessionUtil {
   // Request path
   private static final String SF_PATH_LOGIN_REQUEST = "/session/v1/login-request";
   private static final String SF_PATH_TOKEN_REQUEST = "/session/token-request";
+  private static final String SF_PATH_OKTA_REQUEST_SUFFIX = "/api/v1/authn";
   public static final String SF_PATH_AUTHENTICATOR_REQUEST = "/session/authenticator-request";
   public static final String SF_PATH_CONSOLE_LOGIN_REQUEST = "/console/login";
 
@@ -1340,16 +1343,33 @@ static void closeSession(SFLoginInput loginInput) throws SFException, SnowflakeS
    *
    * @param loginInput Login Info for the request
    * @param ssoUrl URL to use for SSO
-   * @param oneTimeToken The token used for SSO
+   * @param oneTimeTokenSupplier The function returning token used for SSO
    * @return The response in HTML form
    * @throws SnowflakeSQLException Will be thrown if the destination URL in the SAML assertion does
    *     not match
    */
   private static String federatedFlowStep4(
-      SFLoginInput loginInput, String ssoUrl, String oneTimeToken) throws SnowflakeSQLException {
+      SFLoginInput loginInput,
+      String ssoUrl,
+      ThrowingCallable<String, SnowflakeSQLException> oneTimeTokenSupplier)
+      throws SnowflakeSQLException {
+    String oneTimeToken = oneTimeTokenSupplier.call();
     String responseHtml = "";
 
     try {
+      RetryContext retryWithNewOTTManager =
+          new RetryContext(RetryContext.RetryHook.ALWAYS_BEFORE_RETRY);
+      retryWithNewOTTManager.registerRetryCallback(
+          (HttpRequestBase retrieveSamlRequest) -> {
+            try {
+              String newOneTimeToken = oneTimeTokenSupplier.call();
+              prepareFederatedFlowStep4Request(retrieveSamlRequest, ssoUrl, newOneTimeToken);
+            } catch (Exception e) {
+              throw new RuntimeException(e);
+            }
+
+            return null;
+          });
 
       HttpGet httpGet = new HttpGet();
       prepareFederatedFlowStep4Request(httpGet, ssoUrl, oneTimeToken);
@@ -1481,7 +1501,8 @@ private static JsonNode federatedFlowStep1(SFLoginInput loginInput) throws Snowf
     JsonNode dataNode = null;
     try {
       StringEntity requestInput = prepareFederatedFlowStep1RequestInput(loginInput);
-      HttpPost postRequest = prepareFederatedFlowStep1PostRequest(loginInput, requestInput);
+      HttpPost postRequest = new HttpPost();
+      prepareFederatedFlowStep1PostRequest(postRequest, loginInput, requestInput);
 
       final String gsResponse =
           HttpUtil.executeGeneralRequest(
@@ -1559,8 +1580,10 @@ private static String getSamlResponseUsingOkta(SFLoginInput loginInput)
         String tokenUrl = dataNode.path("tokenUrl").asText();
         String ssoUrl = dataNode.path("ssoUrl").asText();
         federatedFlowStep2(loginInput, tokenUrl, ssoUrl);
-        final String oneTimeToken = federatedFlowStep3(loginInput, tokenUrl);
-        return federatedFlowStep4(loginInput, ssoUrl, oneTimeToken);
+        ThrowingCallable<String, SnowflakeSQLException> oneTimeTokenSupplier =
+            () -> federatedFlowStep3(loginInput, tokenUrl);
+
+        return federatedFlowStep4(loginInput, ssoUrl, oneTimeTokenSupplier);
       } catch (SnowflakeSQLException ex) {
         // This error gets thrown if the okta request encountered a retry-able error that
         // requires getting a new one-time token.
@@ -1889,7 +1912,8 @@ public static boolean isNewRetryStrategyRequest(HttpRequestBase request) {
     if (requestPath != null) {
       return requestPath.equals(SF_PATH_LOGIN_REQUEST)
           || requestPath.equals(SF_PATH_AUTHENTICATOR_REQUEST)
-          || requestPath.equals(SF_PATH_TOKEN_REQUEST);
+          || requestPath.equals(SF_PATH_TOKEN_REQUEST)
+          || requestPath.contains(SF_PATH_OKTA_REQUEST_SUFFIX);
     }
     return false;
   }
@@ -1899,25 +1923,22 @@ public static boolean isNewRetryStrategyRequest(HttpRequestBase request) {
    *
    * @param loginInput The login information for the request.
    * @param inputData The JSON input data to include in the request.
-   * @return An {@link HttpPost} object ready to execute the federated flow request.
    * @throws URISyntaxException If the constructed URI is invalid.
    */
-  private static HttpPost prepareFederatedFlowStep1PostRequest(
-      SFLoginInput loginInput, StringEntity inputData) throws URISyntaxException {
+  private static void prepareFederatedFlowStep1PostRequest(
+          HttpPost postRequest, SFLoginInput loginInput, StringEntity inputData) throws URISyntaxException {
     URIBuilder fedUriBuilder = new URIBuilder(loginInput.getServerUrl());
     // TODO: if loginInput.serverUrl contains port or additional segments - it will be ignored and
     // overwritten here - to be fixed in SNOW-1922872
     fedUriBuilder.setPath(SF_PATH_AUTHENTICATOR_REQUEST);
     URI fedUrlUri = fedUriBuilder.build();
+    postRequest.setURI(fedUrlUri);
 
-    HttpPost postRequest = new HttpPost(fedUrlUri);
     postRequest.setEntity(inputData);
     postRequest.addHeader("accept", "application/json");
 
     postRequest.addHeader(SF_HEADER_CLIENT_APP_ID, loginInput.getAppId());
     postRequest.addHeader(SF_HEADER_CLIENT_APP_VERSION, loginInput.getAppVersion());
-
-    return postRequest;
   }
 
   /**