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/cloud/storage/GCSAccessStrategyAwsSdk.java

@@ -9,6 +9,7 @@
 import java.io.InputStream;
 import java.net.URI;
 import java.net.URISyntaxException;
+import java.time.Duration;
 import java.util.List;
 import java.util.Map;
 import java.util.Optional;
@@ -32,6 +33,8 @@
 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.core.retry.backoff.FullJitterBackoffStrategy;
 import software.amazon.awssdk.http.nio.netty.NettyNioAsyncHttpClient;
 import software.amazon.awssdk.http.nio.netty.ProxyConfiguration;
 import software.amazon.awssdk.regions.Region;
@@ -54,7 +57,13 @@ class GCSAccessStrategyAwsSdk implements GCSAccessStrategy {
   private static final SFLogger logger = SFLoggerFactory.getLogger(GCSAccessStrategyAwsSdk.class);
   private final S3AsyncClient amazonClient;
 
-  GCSAccessStrategyAwsSdk(StageInfo stage, SFBaseSession session) throws SnowflakeSQLException {
+  GCSAccessStrategyAwsSdk(
+      StageInfo stage,
+      SFBaseSession session,
+      int maxRetries,
+      int retryBackoffMin,
+      int retryBackoffMaxExponent)
+      throws SnowflakeSQLException {
     String accessToken = (String) stage.getCredentials().get("GCS_ACCESS_TOKEN");
 
     Optional<String> oEndpoint = stage.gcsCustomEndpoint();
@@ -85,6 +94,21 @@ class GCSAccessStrategyAwsSdk implements GCSAccessStrategy {
     overrideConfiguration.putAdvancedOption(
         SdkAdvancedClientOption.SIGNER, new AwsSdkGCPSigner(accessToken));
 
+    // Mirror the JDBC driver's retry settings so putGetMaxRetries and backoff parameters
+    // apply consistently to this strategy. maxBackoffMs = retryBackoffMin <<
+    // retryBackoffMaxExponent
+    // (e.g. 1000ms << 4 = 16 000ms), matching S3ErrorHandler.retryRequestWithExponentialBackoff.
+    long maxBackoffMs = (long) retryBackoffMin << retryBackoffMaxExponent;
+    overrideConfiguration.retryPolicy(
+        RetryPolicy.builder()
+            .numRetries(maxRetries)
+            .backoffStrategy(
+                FullJitterBackoffStrategy.builder()
+                    .baseDelay(Duration.ofMillis(retryBackoffMin))
+                    .maxBackoffTime(Duration.ofMillis(maxBackoffMs))
+                    .build())
+            .build());
+
     ProxyConfiguration proxyConfiguration;
     if (session != null) {
       proxyConfiguration =

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

@@ -203,14 +203,48 @@ public boolean handleStorageException(
       String queryId,
       SnowflakeGCSClient gcsClient)
       throws SnowflakeSQLException {
+    // 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 = null;
     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
+      se = (StorageException) ex;
+    } else {
+      Throwable cause = ex.getCause();
+      while (cause != null) {
+        if (cause instanceof StorageException) {
+          se = (StorageException) cause;
+          break;
+        }
+        cause = cause.getCause();
+      }
+      if (se != null) {
+        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
+      // 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);
+        logger.error("Stack trace: ", ex);
         throw new SnowflakeSQLLoggedException(
             queryId,
             session,
@@ -224,7 +258,7 @@ public boolean handleStorageException(
       } else {
         logger.debug(
             "Encountered exception ({}) during {}, retry count: {}",
-            ex.getMessage(),
+            se.getMessage(),
             operation,
             retryCount);
         logger.debug("Stack trace: ", ex);
@@ -260,6 +294,12 @@ public boolean handleStorageException(
       }
       return true;
     } else {
+      logger.error(
+          "GCSDefaultAccessStrategy: unhandled exception type {} during {}, not retrying: {}",
+          ex.getClass().getName(),
+          operation,
+          ex.getMessage());
+      logger.error("Stack trace: ", ex);
       return false;
     }
   }

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

@@ -1244,7 +1244,13 @@ private void setupGCSClient(
       boolean overrideAwsAccessStrategy =
           Boolean.valueOf(systemGetEnv("SNOWFLAKE_GCS_FORCE_VIRTUAL_STYLE_DOMAINS"));
       if (stage.getUseVirtualUrl() || overrideAwsAccessStrategy) {
-        this.gcsAccessStrategy = new GCSAccessStrategyAwsSdk(stage, session);
+        this.gcsAccessStrategy =
+            new GCSAccessStrategyAwsSdk(
+                stage,
+                session,
+                this.getMaxRetries(),
+                this.getRetryBackoffMin(),
+                this.getRetryBackoffMaxExponent());
       } else {
         this.gcsAccessStrategy = new GCSDefaultAccessStrategy(stage, session);
       }

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));
+  }
+}