Snow 1853185 jdbc driver v3.16 native okta http retry storm

TestOnly/pom.xml

@@ -250,6 +250,19 @@
       <artifactId>jna-platform</artifactId>
       <version>${jna.version}</version>
     </dependency>
+    <dependency>
+      <groupId>org.apache.httpcomponents</groupId>
+      <artifactId>httpclient</artifactId>
+      <version>4.5.14</version>
+      <scope>test</scope>
+    </dependency>
+
+    <dependency>
+      <groupId>org.apache.httpcomponents</groupId>
+      <artifactId>httpcore</artifactId>
+      <version>4.4.16</version>
+      <scope>test</scope>
+    </dependency>
   </dependencies>
 
   <build>

src/main/java/net/snowflake/client/CallableThrowingSnowflakeSqlException.java

@@ -0,0 +1,20 @@
+package net.snowflake.client;
+
+import net.snowflake.client.jdbc.SnowflakeSQLException;
+
+/**
+ * A functional interface similar to {@code Callable<V>}, but whose {@code call()} method is
+ * permitted to throw a {@link SnowflakeSQLException}.
+ *
+ * @param <V> the result type of the call method
+ */
+@FunctionalInterface
+public interface CallableThrowingSnowflakeSqlException<V> {
+  /**
+   * Computes a result, potentially throwing a {@code SnowflakeSQLException}.
+   *
+   * @return the computed result
+   * @throws SnowflakeSQLException if a SQL exception occurs during execution
+   */
+  V call() throws SnowflakeSQLException;
+}

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

@@ -40,6 +40,7 @@
 import net.snowflake.client.log.SFLogger;
 import net.snowflake.client.log.SFLoggerFactory;
 import net.snowflake.client.log.SFLoggerUtil;
+import net.snowflake.client.util.RetryContextManager;
 import net.snowflake.client.util.SecretDetector;
 import net.snowflake.client.util.Stopwatch;
 import net.snowflake.common.core.SqlState;
@@ -621,6 +622,24 @@ static String executeRequestWithoutCookies(
         new ExecTimeTelemetryData());
   }
 
+  public static String executeGeneralRequest(
+      HttpRequestBase httpRequest,
+      int retryTimeout,
+      int authTimeout,
+      int socketTimeout,
+      int retryCount,
+      HttpClientSettingsKey ocspAndProxyAndGzipKey)
+      throws SnowflakeSQLException, IOException {
+    return executeGeneralRequest(
+        httpRequest,
+        retryTimeout,
+        authTimeout,
+        socketTimeout,
+        retryCount,
+        ocspAndProxyAndGzipKey,
+        null);
+  }
+
   /**
    * Executes a HTTP request for Snowflake.
    *
@@ -630,6 +649,7 @@ static String executeRequestWithoutCookies(
    * @param socketTimeout socket timeout (in ms)
    * @param retryCount retry count for the request
    * @param ocspAndProxyAndGzipKey OCSP mode and proxy settings for httpclient
+   * @param retryContextManager RetryContextManager passed to RestRequest
    * @return response
    * @throws SnowflakeSQLException if Snowflake error occurs
    * @throws IOException raises if a general IO error occurs
@@ -640,7 +660,8 @@ public static String executeGeneralRequest(
       int authTimeout,
       int socketTimeout,
       int retryCount,
-      HttpClientSettingsKey ocspAndProxyAndGzipKey)
+      HttpClientSettingsKey ocspAndProxyAndGzipKey,
+      RetryContextManager retryContextManager)
       throws SnowflakeSQLException, IOException {
     logger.debug("Executing general request");
     return executeRequest(
@@ -654,7 +675,8 @@ public static String executeGeneralRequest(
         false, // no retry parameter
         false, // no retry on HTTP 403
         ocspAndProxyAndGzipKey,
-        new ExecTimeTelemetryData());
+        new ExecTimeTelemetryData(),
+        retryContextManager);
   }
 
   /**
@@ -695,6 +717,34 @@ public static String executeGeneralRequest(
         new ExecTimeTelemetryData());
   }
 
+  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)
+      throws SnowflakeSQLException, IOException {
+    return executeRequest(
+        httpRequest,
+        retryTimeout,
+        authTimeout,
+        socketTimeout,
+        maxRetries,
+        injectSocketTimeout,
+        canceling,
+        includeRetryParameters,
+        retryOnHTTP403,
+        ocspAndProxyKey,
+        execTimeData,
+        null);
+  }
+
   /**
    * Executes a HTTP request for Snowflake.
    *
@@ -709,6 +759,7 @@ public static String executeGeneralRequest(
    * @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 retryContextManager RetryContextManager passed to RestRequest
    * @return response
    * @throws SnowflakeSQLException if Snowflake error occurs
    * @throws IOException raises if a general IO error occurs
@@ -724,7 +775,8 @@ public static String executeRequest(
       boolean includeRetryParameters,
       boolean retryOnHTTP403,
       HttpClientSettingsKey ocspAndProxyKey,
-      ExecTimeTelemetryData execTimeData)
+      ExecTimeTelemetryData execTimeData,
+      RetryContextManager retryContextManager)
       throws SnowflakeSQLException, IOException {
     boolean ocspEnabled = !(ocspAndProxyKey.getOcspMode().equals(OCSPMode.DISABLE_OCSP_CHECKS));
     logger.debug("Executing request with OCSP enabled: {}", ocspEnabled);
@@ -742,7 +794,40 @@ public static String executeRequest(
         true, // include request GUID
         retryOnHTTP403,
         getHttpClient(ocspAndProxyKey),
-        execTimeData);
+        execTimeData,
+        retryContextManager);
+  }
+
+  private static String executeRequestInternal(
+      HttpRequestBase httpRequest,
+      int retryTimeout,
+      int authTimeout,
+      int socketTimeout,
+      int maxRetries,
+      int injectSocketTimeout,
+      AtomicBoolean canceling,
+      boolean withoutCookies,
+      boolean includeRetryParameters,
+      boolean includeRequestGuid,
+      boolean retryOnHTTP403,
+      CloseableHttpClient httpClient,
+      ExecTimeTelemetryData execTimeData)
+      throws SnowflakeSQLException, IOException {
+    return executeRequestInternal(
+        httpRequest,
+        retryTimeout,
+        authTimeout,
+        socketTimeout,
+        maxRetries,
+        injectSocketTimeout,
+        canceling,
+        withoutCookies,
+        includeRetryParameters,
+        includeRequestGuid,
+        retryOnHTTP403,
+        httpClient,
+        execTimeData,
+        null);
   }
 
   /**
@@ -764,6 +849,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 retryContextManager RetryContextManager passed to RestRequest
    * @return response in String
    * @throws SnowflakeSQLException if Snowflake error occurs
    * @throws IOException raises if a general IO error occurs
@@ -781,7 +867,8 @@ private static String executeRequestInternal(
       boolean includeRequestGuid,
       boolean retryOnHTTP403,
       CloseableHttpClient httpClient,
-      ExecTimeTelemetryData execTimeData)
+      ExecTimeTelemetryData execTimeData,
+      RetryContextManager retryContextManager)
       throws SnowflakeSQLException, IOException {
     // HttpRequest.toString() contains request URI. Scrub any credentials, if
     // present, before logging
@@ -815,7 +902,8 @@ private static String executeRequestInternal(
               includeRetryParameters,
               includeRequestGuid,
               retryOnHTTP403,
-              execTimeData);
+              execTimeData,
+              retryContextManager);
       if (logger.isDebugEnabled() && stopwatch != null) {
         stopwatch.stop();
       }