SNOW-3745081: fix retry in GCS storage client

CHANGELOG.md

@@ -2,6 +2,7 @@
 
 # Changelog
 - v4.3.2-SNAPSHOT
+  - Fixed GCS PUT operations not retrying on transient errors (e.g. HTTP 503) despite `putGetMaxRetries` being configured (snowflakedb/snowflake-jdbc#2688).
   - Bumped jackson-databind to 2.18.8 from 2.18.7 (snowflakedb/snowflake-jdbc#2669).
   - Fixed snowflake-jdbc writing a `snowflake-minicore-*` temp directory and loading the native library at driver class-load time even when the driver was never used (e.g. when present on the classpath only as a transitive dependency). Minicore now loads lazily when the first Snowflake connection is created (`ConnectionFactory.createConnection`) instead of during `DriverInitializer.initialize()` (snowflakedb/snowflake-jdbc#2670).
   - Restored `GetCallerIdentity` as the default AWS Workload Identity Federation attestation method to avoid breaking existing users who have not configured the `ISSUER` in their Snowflake WIF setup. The `GetWebIdentityToken` (outbound JWT) flow introduced in v4.3.0 is now opt-in via the new `workloadIdentityAwsUseOutboundToken` connection property.

src/main/java/net/snowflake/client/internal/jdbc/SnowflakeUtil.java

@@ -649,6 +649,22 @@ public static Throwable getRootCause(Exception ex) {
     return cause;
   }
 
+  /**
+   * Walks the exception cause chain and returns the first {@link Throwable} that is an instance of
+   * {@code type}, or {@code null} if none is found. Unlike {@link #getRootCause(Exception)}, which
+   * always returns the deepest cause, this method stops at the first match of the requested type.
+   */
+  public static <T extends Throwable> T findFirstCauseOfType(Throwable ex, Class<T> type) {
+    Throwable cause = ex;
+    while (cause != null) {
+      if (type.isInstance(cause)) {
+        return type.cast(cause);
+      }
+      cause = cause.getCause();
+    }
+    return null;
+  }
+
   public static boolean isBlank(String input) {
     if ("".equals(input) || input == null) {
       return true;

src/main/java/net/snowflake/client/internal/jdbc/cloud/storage/GCSAccessStrategyAwsSdk.java

@@ -32,6 +32,7 @@
 import software.amazon.awssdk.core.client.config.ClientOverrideConfiguration;
 import software.amazon.awssdk.core.client.config.SdkAdvancedClientOption;
 import software.amazon.awssdk.core.exception.SdkException;
+import software.amazon.awssdk.core.retry.RetryPolicy;
 import software.amazon.awssdk.http.nio.netty.NettyNioAsyncHttpClient;
 import software.amazon.awssdk.http.nio.netty.ProxyConfiguration;
 import software.amazon.awssdk.regions.Region;
@@ -85,6 +86,13 @@ class GCSAccessStrategyAwsSdk implements GCSAccessStrategy {
     overrideConfiguration.putAdvancedOption(
         SdkAdvancedClientOption.SIGNER, new AwsSdkGCPSigner(accessToken));
 
+    // Disable the SDK's own internal retries so the outer JDBC retry loop in
+    // SnowflakeGCSClient.upload() is the sole controller of retry count and backoff.
+    // Starting from the default policy (rather than an empty builder) preserves the AWS
+    // token-bucket circuit breaker and throttling-aware retry conditions.
+    overrideConfiguration.retryPolicy(
+        RetryPolicy.defaultRetryPolicy().toBuilder().numRetries(0).build());
+
     ProxyConfiguration proxyConfiguration;
     if (session != null) {
       proxyConfiguration =
@@ -261,17 +269,31 @@ public boolean handleStorageException(
       String queryId,
       SnowflakeGCSClient gcsClient)
       throws SnowflakeSQLException {
-    Throwable cause = ex.getCause();
-    if (cause instanceof SdkException) {
-      logger.debug("GCSAccessStrategyAwsSdk: " + cause.getMessage());
+    // Walk the full cause chain to find an SdkException. The exception may arrive wrapped
+    // (e.g. inside SnowflakeSQLLoggedException → CompletionException) when it originates from the
+    // inner uploadWithDownScopedToken wrapper. Walking the chain ensures that transient errors
+    // (503, connection refused, etc.) are always treated as retryable regardless of wrapping depth,
+    // consistent with how GCSDefaultAccessStrategy handles StorageException.
+    SdkException sdkEx = SnowflakeUtil.findFirstCauseOfType(ex, SdkException.class);
+    if (sdkEx != null) {
+      if (sdkEx != ex) {
+        logger.debug(
+            "GCSAccessStrategyAwsSdk: found SdkException ({}) wrapped inside {} during {};"
+                + " treating as retryable",
+            sdkEx.getMessage(),
+            ex.getClass().getSimpleName(),
+            operation);
+      } else {
+        logger.debug("GCSAccessStrategyAwsSdk: {}", sdkEx.getMessage());
+      }
 
       if (retryCount > gcsClient.getMaxRetries()
-          || S3ErrorHandler.isClientException400Or404(cause)) {
+          || S3ErrorHandler.isClientException400Or404(sdkEx)) {
         throwIfClientExceptionOrMaxRetryReached(
-            operation, session, command, queryId, gcsClient, cause);
+            operation, session, command, queryId, gcsClient, sdkEx);
       } else {
         retryRequestWithExponentialBackoff(
-            ex, retryCount, operation, session, command, gcsClient, queryId, cause);
+            ex, retryCount, operation, session, command, gcsClient, queryId, sdkEx);
       }
       return true;
     } else {

src/main/java/net/snowflake/client/internal/jdbc/cloud/storage/GCSDefaultAccessStrategy.java

@@ -18,6 +18,7 @@
 import com.google.cloud.storage.StorageOptions;
 import java.io.File;
 import java.io.InputStream;
+import java.net.SocketTimeoutException;
 import java.nio.channels.Channels;
 import java.util.Map;
 import net.snowflake.client.api.exception.SnowflakeSQLException;
@@ -29,6 +30,7 @@
 import net.snowflake.client.internal.log.SFLoggerFactory;
 import net.snowflake.client.internal.util.SFPair;
 import net.snowflake.common.core.SqlState;
+import org.apache.http.HttpStatus;
 
 class GCSDefaultAccessStrategy implements GCSAccessStrategy {
   private static final SFLogger logger = SFLoggerFactory.getLogger(GCSDefaultAccessStrategy.class);
@@ -203,14 +205,71 @@ public boolean handleStorageException(
       String queryId,
       SnowflakeGCSClient gcsClient)
       throws SnowflakeSQLException {
-    if (ex instanceof StorageException) {
-      // NOTE: this code path only handle Access token based operation,
-      // presigned URL is not covered. Presigned Url do not raise
-      // StorageException
+    // Find a StorageException either directly or anywhere in the cause chain.
+    // The exception may arrive wrapped (e.g. in SnowflakeSQLException) when it originates from the
+    // inner uploadWithDownScopedToken method, which catches the original StorageException and
+    // re-wraps it before re-throwing. Walking the chain ensures that a transient 503 or similar
+    // GCS error is always treated as retryable regardless of how many layers of wrapping were
+    // added.
+    StorageException se = SnowflakeUtil.findFirstCauseOfType(ex, StorageException.class);
+    if (se != null && se != ex) {
+      logger.debug(
+          "GCSDefaultAccessStrategy: found StorageException (HTTP {}) wrapped inside {} during {};"
+              + " treating as retryable",
+          se.getCode(),
+          ex.getClass().getSimpleName(),
+          operation);
+    }
+
+    if (se != null) {
+      // NOTE: this code path only handles Access token based operations.
+      // Presigned URL operations do not raise StorageException.
 
-      StorageException se = (StorageException) ex;
-      // If we have exceeded the max number of retries, propagate the error
+      // Non-transient client errors must not be retried during upload/download, mirroring
+      // S3ErrorHandler.isClientException400Or404 (400 and 404 only).
+      // 401 is excluded here because it is handled below via token refresh.
+      //
+      // WHY THE SCOPE IS LIMITED TO UPLOAD AND DOWNLOAD:
+      // When OVERWRITE=FALSE, SnowflakeFileTransferAgent.filterExistingFiles() calls
+      // storageClient.listObjects() (returns an empty collection on miss — never a 404) and
+      // then storageClient.getObjectMetadata() on each listed object to compare digests.
+      // If the object disappears between the list and the metadata fetch (a race condition),
+      // GCSDefaultAccessStrategy.getObjectMetadata() throws StorageProviderException wrapping a
+      // synthetic StorageException(404). The SnowflakeFileTransferAgent.compareAndSkipRemoteFiles()
+      // SNOW-14521 handler attempts to recognise this 404 via isServiceException404(), but that
+      // method only checks SdkServiceException (AWS SDK); it does not recognise the GCS-native
+      // StorageException, so the exception is re-thrown and eventually reaches
+      // handleStorageException() with operation="compareRemoteFiles". Applying the
+      // immediate-throw rule for 404 here would break that retry path — the correct behaviour is
+      // to retry the list/compare cycle via the normal max-retries logic.
+      // NOTE: for the upload path itself, OVERWRITE=FALSE produces only a LIST (GET ?prefix=...)
+      // call and, if the file is found, skips the upload before getObjectMetadata() is ever
+      // reached. No 404 is generated by the upload-path existence check.
+      if ((StorageHelper.UPLOAD.equals(operation) || StorageHelper.DOWNLOAD.equals(operation))
+          && (se.getCode() == HttpStatus.SC_BAD_REQUEST
+              || se.getCode() == HttpStatus.SC_NOT_FOUND)) {
+        throw new SnowflakeSQLLoggedException(
+            queryId,
+            session,
+            SqlState.SYSTEM_ERROR,
+            StorageHelper.getOperationException(operation).getMessageCode(),
+            se,
+            operation,
+            se.getCode(),
+            se.getMessage(),
+            se.getReason());
+      }
+
+      // If we have exceeded the max number of retries, propagate the error.
       if (retryCount > gcsClient.getMaxRetries()) {
+        logger.error(
+            "GCSDefaultAccessStrategy: max retries ({}) exceeded for StorageException"
+                + " (HTTP {}, message: {}) during {}, giving up",
+            gcsClient.getMaxRetries(),
+            se.getCode(),
+            se.getMessage(),
+            operation,
+            ex);
         throw new SnowflakeSQLLoggedException(
             queryId,
             session,
@@ -224,10 +283,10 @@ public boolean handleStorageException(
       } else {
         logger.debug(
             "Encountered exception ({}) during {}, retry count: {}",
-            ex.getMessage(),
+            se.getMessage(),
             operation,
-            retryCount);
-        logger.debug("Stack trace: ", ex);
+            retryCount,
+            ex);
 
         // exponential backoff up to a limit
         int backoffInMillis = gcsClient.getRetryBackoffMin();
@@ -260,7 +319,29 @@ public boolean handleStorageException(
       }
       return true;
     } else {
-      return false;
+      // InterruptedException and SocketTimeoutException are handled by the outer retry loop in
+      // SnowflakeGCSClient.handleStorageException (retry within max retries, throw when exhausted).
+      // Return false so that logic is preserved rather than bypassed.
+      if (ex instanceof InterruptedException
+          || SnowflakeUtil.getRootCause(ex) instanceof SocketTimeoutException) {
+        return false;
+      }
+      logger.error(
+          "GCSDefaultAccessStrategy: unhandled exception type {} during {}, not retrying",
+          ex.getClass().getName(),
+          operation,
+          ex);
+      // Re-throw as-is to avoid double-wrapping an already well-formed SQL exception.
+      if (ex instanceof SnowflakeSQLException) {
+        throw (SnowflakeSQLException) ex;
+      }
+      throw new SnowflakeSQLLoggedException(
+          queryId,
+          session,
+          SqlState.SYSTEM_ERROR,
+          StorageHelper.getOperationException(operation).getMessageCode(),
+          ex,
+          "Encountered exception during " + operation + ": " + ex.getMessage());
     }
   }
 

src/test/java/net/snowflake/client/internal/jdbc/cloud/storage/GCSDefaultAccessStrategyTest.java

@@ -0,0 +1,103 @@
+package net.snowflake.client.internal.jdbc.cloud.storage;
+
+import static org.junit.jupiter.api.Assertions.assertDoesNotThrow;
+import static org.junit.jupiter.api.Assertions.assertThrows;
+import static org.mockito.Mockito.CALLS_REAL_METHODS;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.when;
+import static org.mockito.Mockito.withSettings;
+
+import com.google.cloud.storage.StorageException;
+import java.util.HashMap;
+import net.snowflake.client.api.exception.SnowflakeSQLException;
+import org.junit.jupiter.api.BeforeEach;
+import org.junit.jupiter.api.Test;
+
+class GCSDefaultAccessStrategyTest {
+
+  private GCSDefaultAccessStrategy strategy;
+  private SnowflakeGCSClient mockGcsClient;
+
+  @BeforeEach
+  void setUp() throws SnowflakeSQLException {
+    // Use empty credentials (no GCS_ACCESS_TOKEN) so the constructor falls back to anonymous auth,
+    // which does not probe any real GCS endpoint.
+    StageInfo stage =
+        StageInfo.createStageInfo(
+            "GCS", "test-bucket/path", new HashMap<>(), "US-CENTRAL1", null, null, true);
+    strategy = new GCSDefaultAccessStrategy(stage, /* session= */ null);
+
+    // Mock SnowflakeGCSClient to avoid real network calls and to control retry parameters.
+    // getRetryBackoffMin is set to 0 so tests are not slowed down by exponential backoff.
+    mockGcsClient =
+        mock(SnowflakeGCSClient.class, withSettings().defaultAnswer(CALLS_REAL_METHODS));
+    when(mockGcsClient.getMaxRetries()).thenReturn(1);
+    when(mockGcsClient.getRetryBackoffMin()).thenReturn(0);
+    when(mockGcsClient.getRetryBackoffMaxExponent()).thenReturn(0);
+  }
+
+  // A direct StorageException(503) below the max retry threshold must not be thrown — the handler
+  // should perform a backoff and signal to the caller that a retry should proceed.
+  @Test
+  void testHandleStorageExceptionDirect503BelowMaxRetries() {
+    StorageException storageException = new StorageException(503, "Service Unavailable");
+
+    assertDoesNotThrow(
+        () ->
+            strategy.handleStorageException(
+                storageException,
+                /* retryCount= */ 1,
+                "upload",
+                /* session= */ null,
+                /* command= */ null,
+                /* queryId= */ null,
+                mockGcsClient));
+  }
+
+  // A StorageException(503) wrapped inside another exception must also be detected via cause-chain
+  // walking and treated as retryable when below the max retry threshold.  This is the key
+  // regression scenario: uploadWithDownScopedToken wraps the original StorageException inside a
+  // SnowflakeSQLException before re-throwing it to the outer upload() retry loop.
+  @Test
+  void testHandleStorageExceptionWrapped503BelowMaxRetries() {
+    StorageException storageException = new StorageException(503, "Service Unavailable");
+    RuntimeException wrappedException =
+        new RuntimeException(
+            "Encountered exception during upload: " + storageException.getMessage(),
+            storageException);
+
+    assertDoesNotThrow(
+        () ->
+            strategy.handleStorageException(
+                wrappedException,
+                /* retryCount= */ 1,
+                "upload",
+                /* session= */ null,
+                /* command= */ null,
+                /* queryId= */ null,
+                mockGcsClient));
+  }
+
+  // When a wrapped StorageException(503) arrives after the maximum number of retries has been
+  // exceeded the handler must surface it as a SnowflakeSQLException so the caller stops retrying.
+  @Test
+  void testHandleStorageExceptionWrapped503OverMaxRetries() {
+    StorageException storageException = new StorageException(503, "Service Unavailable");
+    RuntimeException wrappedException =
+        new RuntimeException(
+            "Encountered exception during upload: " + storageException.getMessage(),
+            storageException);
+
+    assertThrows(
+        SnowflakeSQLException.class,
+        () ->
+            strategy.handleStorageException(
+                wrappedException,
+                /* retryCount= */ 2,
+                "upload",
+                /* session= */ null,
+                /* command= */ null,
+                /* queryId= */ null,
+                mockGcsClient));
+  }
+}