SNOW-913746 Honor useS3RegionalUrl from JsonNode in getStageInfo (#1514)

* SNOW-913746 Honor useS3RegionalUrl from JsonNode

* Address Ilesh feedback

src/main/java/net/snowflake/client/jdbc/SnowflakeFileTransferAgent.java

@@ -929,21 +929,16 @@ private void parseCommand() throws SnowflakeSQLException {
     }
   }
 
+  /**
+   * Construct Stage Info object from JsonNode.
+   *
+   * @param jsonNode JsonNode to use serialize into StageInfo Object
+   * @param session can be null.
+   * @return StageInfo constructed from JsonNode and session params.
+   * @throws SnowflakeSQLException
+   */
   static StageInfo getStageInfo(JsonNode jsonNode, SFSession session) throws SnowflakeSQLException {
 
-    StageInfo stageInfo = getStageInfo(jsonNode);
-
-    // Update StageInfo to reflect use of S3 regional URL.
-    // This is required for connecting to S3 over privatelink when the
-    // target stage is in us-east-1
-    if (stageInfo.getStageType() == StageInfo.StageType.S3)
-      stageInfo.setUseS3RegionalUrl(session.getUseRegionalS3EndpointsForPresignedURL());
-
-    return stageInfo;
-  }
-
-  static StageInfo getStageInfo(JsonNode jsonNode) throws SnowflakeSQLException {
-
     // more parameters common to upload/download
     String stageLocation = jsonNode.path("data").path("stageInfo").path("location").asText();
 
@@ -1033,6 +1028,24 @@ static StageInfo getStageInfo(JsonNode jsonNode) throws SnowflakeSQLException {
       }
     }
 
+    if (stageInfo.getStageType() == StageInfo.StageType.S3) {
+      if (session == null) {
+        // This node's value is set if PUT is used without Session. (For Snowpipe Streaming, we rely
+        // on a response from a server to have this field set to use S3RegionalURL)
+        JsonNode useS3RegionalURLNode =
+            jsonNode.path("data").path("stageInfo").path("useS3RegionalUrl");
+        if (!useS3RegionalURLNode.isMissingNode()) {
+          boolean useS3RegionalUrl = useS3RegionalURLNode.asBoolean(false);
+          stageInfo.setUseS3RegionalUrl(useS3RegionalUrl);
+        }
+      } else {
+        // Update StageInfo to reflect use of S3 regional URL.
+        // This is required for connecting to S3 over privatelink when the
+        // target stage is in us-east-1
+        stageInfo.setUseS3RegionalUrl(session.getUseRegionalS3EndpointsForPresignedURL());
+      }
+    }
+
     return stageInfo;
   }
   /**
@@ -1229,6 +1242,8 @@ public List<SnowflakeFileTransferMetadata> getFileTransferMetadatas()
    *
    * <p>NOTE: It only supports PUT on S3/AZURE/GCS (i.e. NOT LOCAL_FS)
    *
+   * <p>It also assumes there is no active SFSession
+   *
    * @param jsonNode JSON doc returned by GS from PUT call
    * @return The file transfer metadatas for to-be-transferred files.
    * @throws SnowflakeSQLException if any error occurs
@@ -1272,7 +1287,7 @@ public static List<SnowflakeFileTransferMetadata> getFileTransferMetadatas(JsonN
 
     final Set<String> sourceFiles = expandFileNames(srcLocations);
 
-    StageInfo stageInfo = getStageInfo(jsonNode);
+    StageInfo stageInfo = getStageInfo(jsonNode, null /*SFSession*/);
 
     List<SnowflakeFileTransferMetadata> result = new ArrayList<>();
     if (stageInfo.getStageType() != StageInfo.StageType.GCS

src/test/java/net/snowflake/client/jdbc/FileUploaderPrepIT.java

@@ -120,6 +120,7 @@ abstract class FileUploaderPrepIT extends BaseJDBCTest {
           + "      \"region\": \"us-west-2\",\n"
           + "      \"storageAccount\": null,\n"
           + "      \"isClientSideEncrypted\": true,\n"
+          + "      \"useS3RegionalUrl\": true,\n"
           + "      \"creds\": {\n"
           + "        \"AWS_KEY_ID\": \"EXAMPLE_AWS_KEY_ID\",\n"
           + "        \"AWS_SECRET_KEY\": \"EXAMPLE_AWS_SECRET_KEY\",\n"

src/test/java/net/snowflake/client/jdbc/FileUploaderSessionlessTest.java

@@ -56,7 +56,7 @@ public void testGetEncryptionMaterial() throws Exception {
 
   @Test
   public void testGetS3StageData() throws Exception {
-    StageInfo stageInfo = SnowflakeFileTransferAgent.getStageInfo(exampleS3JsonNode);
+    StageInfo stageInfo = SnowflakeFileTransferAgent.getStageInfo(exampleS3JsonNode, null);
     Map<String, String> expectedCreds = new HashMap<>();
     expectedCreds.put("AWS_ID", "EXAMPLE_AWS_ID");
     expectedCreds.put("AWS_KEY", "EXAMPLE_AWS_KEY");
@@ -71,11 +71,13 @@ public void testGetS3StageData() throws Exception {
     Assert.assertEquals("null", stageInfo.getEndPoint());
     Assert.assertEquals(null, stageInfo.getStorageAccount());
     Assert.assertEquals(true, stageInfo.getIsClientSideEncrypted());
+    Assert.assertEquals(true, stageInfo.getUseS3RegionalUrl());
   }
 
   @Test
   public void testGetS3StageDataWithStageEndpoint() throws Exception {
-    StageInfo stageInfo = SnowflakeFileTransferAgent.getStageInfo(exampleS3StageEndpointJsonNode);
+    StageInfo stageInfo =
+        SnowflakeFileTransferAgent.getStageInfo(exampleS3StageEndpointJsonNode, null);
     Map<String, String> expectedCreds = new HashMap<>();
     expectedCreds.put("AWS_ID", "EXAMPLE_AWS_ID");
     expectedCreds.put("AWS_KEY", "EXAMPLE_AWS_KEY");
@@ -94,7 +96,7 @@ public void testGetS3StageDataWithStageEndpoint() throws Exception {
 
   @Test
   public void testGetAzureStageData() throws Exception {
-    StageInfo stageInfo = SnowflakeFileTransferAgent.getStageInfo(exampleAzureJsonNode);
+    StageInfo stageInfo = SnowflakeFileTransferAgent.getStageInfo(exampleAzureJsonNode, null);
     Map<String, String> expectedCreds = new HashMap<>();
     expectedCreds.put("AZURE_SAS_TOKEN", "EXAMPLE_AZURE_SAS_TOKEN");
 
@@ -109,7 +111,7 @@ public void testGetAzureStageData() throws Exception {
 
   @Test
   public void testGetGCSStageData() throws Exception {
-    StageInfo stageInfo = SnowflakeFileTransferAgent.getStageInfo(exampleGCSJsonNode);
+    StageInfo stageInfo = SnowflakeFileTransferAgent.getStageInfo(exampleGCSJsonNode, null);
     Map<String, String> expectedCreds = new HashMap<>();
 
     Assert.assertEquals(StageInfo.StageType.GCS, stageInfo.getStageType());