SNOW-999358 Handle GCP token expiration correctly when using connecti…

…onless mode (#1611)
snowflakedb · Jan 23, 2024 · 600d4c2 · 600d4c2
1 parent 697073f
commit 600d4c2
Showing 1 changed file with 13 additions and 4 deletions.
diff --git a/src/main/java/net/snowflake/client/jdbc/cloud/storage/SnowflakeGCSClient.java b/src/main/java/net/snowflake/client/jdbc/cloud/storage/SnowflakeGCSClient.java
@@ -575,7 +575,8 @@ public void uploadWithPresignedUrlWithoutConnection(
           destFileName,
           meta.getContentEncoding(),
           meta.getUserMetadata(),
-          uploadStreamInfo.left);
+          uploadStreamInfo.left,
+          queryId);
       logger.debug("Upload successfully with downscoped token");
     } else {
       logger.debug("Starting upload with presigned url");
@@ -682,7 +683,8 @@ public void upload(
             destFileName,
             meta.getContentEncoding(),
             meta.getUserMetadata(),
-            uploadStreamInfo.left);
+            uploadStreamInfo.left,
+            queryId);
 
         logger.debug("Upload successful", false);
 
@@ -742,7 +744,9 @@ private void uploadWithDownScopedToken(
       String destFileName,
       String contentEncoding,
       Map<String, String> metadata,
-      InputStream content) {
+      InputStream content,
+      String queryId)
+      throws SnowflakeSQLException {
     logger.debug("Uploading file {} to bucket {}", destFileName, remoteStorageLocation);
     BlobId blobId = BlobId.of(remoteStorageLocation, destFileName);
     BlobInfo blobInfo =
@@ -751,7 +755,12 @@ private void uploadWithDownScopedToken(
             .setMetadata(metadata)
             .build();
 
-    gcsClient.create(blobInfo, content);
+    try {
+      gcsClient.create(blobInfo, content);
+    } catch (Exception e) {
+      handleStorageException(e, 0, "upload", session, queryId);
+      throw e;
+    }
   }
 
   /**