From 42f2ccedfa14a809a3a315a2f81561332bb44ef5 Mon Sep 17 00:00:00 2001
From: Alec Huang <alec.huang@snowflake.com>
Date: Wed, 27 Mar 2024 11:23:54 -0700
Subject: [PATCH] Reduce unnecessary OAuth token refresh retry.

---
 .../ingest/connection/OAuthClient.java        | 36 +++++++++----------
 .../ingest/connection/OAuthManager.java       |  4 +++
 2 files changed, 20 insertions(+), 20 deletions(-)

diff --git a/src/main/java/net/snowflake/ingest/connection/OAuthClient.java b/src/main/java/net/snowflake/ingest/connection/OAuthClient.java
index bc2b64089..61c736a42 100644
--- a/src/main/java/net/snowflake/ingest/connection/OAuthClient.java
+++ b/src/main/java/net/snowflake/ingest/connection/OAuthClient.java
@@ -71,28 +71,24 @@ public AtomicReference<OAuthCredential> getOAuthCredentialRef() {
   }
 
   /** Refresh access token using a valid refresh token */
-  public void refreshToken() {
-    String respBodyString = null;
-    try (CloseableHttpResponse httpResponse = httpClient.execute(makeRefreshTokenRequest())) {
-      respBodyString = EntityUtils.toString(httpResponse.getEntity());
-
-      if (httpResponse.getStatusLine().getStatusCode() == HttpStatusCodes.STATUS_CODE_OK) {
-        JsonObject respBody = JsonParser.parseString(respBodyString).getAsJsonObject();
-
-        if (respBody.has(ACCESS_TOKEN) && respBody.has(EXPIRES_IN)) {
-          // Trim surrounding quotation marks
-          String newAccessToken = respBody.get(ACCESS_TOKEN).toString().replaceAll("^\"|\"$", "");
-          oAuthCredential.get().setAccessToken(newAccessToken);
-          oAuthCredential.get().setExpiresIn(respBody.get(EXPIRES_IN).getAsInt());
-          return;
-        }
+  public void refreshToken() throws IOException {
+    CloseableHttpResponse httpResponse = httpClient.execute(makeRefreshTokenRequest());
+    String respBodyString = EntityUtils.toString(httpResponse.getEntity());
+
+    if (httpResponse.getStatusLine().getStatusCode() == HttpStatusCodes.STATUS_CODE_OK) {
+      JsonObject respBody = JsonParser.parseString(respBodyString).getAsJsonObject();
+
+      if (respBody.has(ACCESS_TOKEN) && respBody.has(EXPIRES_IN)) {
+        // Trim surrounding quotation marks
+        String newAccessToken = respBody.get(ACCESS_TOKEN).toString().replaceAll("^\"|\"$", "");
+        oAuthCredential.get().setAccessToken(newAccessToken);
+        oAuthCredential.get().setExpiresIn(respBody.get(EXPIRES_IN).getAsInt());
+        return;
       }
-      throw new SFException(
-          ErrorCode.OAUTH_REFRESH_TOKEN_ERROR,
-          "Refresh access token fail with response: " + respBodyString);
-    } catch (IOException e) {
-      throw new SFException(ErrorCode.OAUTH_REFRESH_TOKEN_ERROR, e.getMessage());
     }
+    throw new SFException(
+        ErrorCode.OAUTH_REFRESH_TOKEN_ERROR,
+        "Refresh access token fail with response: " + respBodyString);
   }
 
   /** Helper method for making refresh request */
diff --git a/src/main/java/net/snowflake/ingest/connection/OAuthManager.java b/src/main/java/net/snowflake/ingest/connection/OAuthManager.java
index a30396e7b..f9e614a27 100644
--- a/src/main/java/net/snowflake/ingest/connection/OAuthManager.java
+++ b/src/main/java/net/snowflake/ingest/connection/OAuthManager.java
@@ -4,6 +4,7 @@
 
 package net.snowflake.ingest.connection;
 
+import java.io.IOException;
 import java.util.concurrent.TimeUnit;
 import net.snowflake.client.jdbc.internal.apache.http.client.utils.URIBuilder;
 import net.snowflake.ingest.utils.Constants;
@@ -154,6 +155,9 @@ void refreshToken() {
             "Refresh access token, next refresh is scheduled after {} seconds", nextRefreshDelay);
 
         return;
+      } catch (IOException e) {
+        // Http client already retried on IO exception, skip retires
+        break;
       } catch (SFException e1) {
         // Exponential backoff retries
         try {