diff --git a/src/main/java/net/snowflake/client/jdbc/SnowflakeFileTransferAgent.java b/src/main/java/net/snowflake/client/jdbc/SnowflakeFileTransferAgent.java index 216d87c74..543e8a7f5 100644 --- 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 getFileTransferMetadatas() * *

NOTE: It only supports PUT on S3/AZURE/GCS (i.e. NOT LOCAL_FS) * + *

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 getFileTransferMetadatas(JsonN final Set sourceFiles = expandFileNames(srcLocations); - StageInfo stageInfo = getStageInfo(jsonNode); + StageInfo stageInfo = getStageInfo(jsonNode, null /*SFSession*/); List 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 index 4d67cbf8c..9bf765da3 100644 --- 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 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 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 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 expectedCreds = new HashMap<>(); Assert.assertEquals(StageInfo.StageType.GCS, stageInfo.getStageType());