Address Ilesh feedback

snowflakedb · Sep 21, 2023 · fe8263f · fe8263f
1 parent ea0b701
commit fe8263f
Show file tree

Hide file tree

Showing 2 changed files with 30 additions and 26 deletions.
diff --git a/src/main/java/net/snowflake/client/jdbc/SnowflakeFileTransferAgent.java b/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,15 +1028,21 @@ static StageInfo getStageInfo(JsonNode jsonNode) throws SnowflakeSQLException {
       }
     }
 
-    boolean useS3RegionalUrl;
     if (stageInfo.getStageType() == StageInfo.StageType.S3) {
-      // 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()) {
-        useS3RegionalUrl = useS3RegionalURLNode.asBoolean(false);
-        stageInfo.setUseS3RegionalUrl(useS3RegionalUrl);
+      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());
       }
     }
 
@@ -1241,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
@@ -1284,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

diff --git a/src/test/java/net/snowflake/client/jdbc/FileUploaderSessionlessTest.java b/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");
@@ -76,7 +76,8 @@ public void testGetS3StageData() throws Exception {
 
   @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");
@@ -95,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");
 
@@ -110,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());