diff --git a/src/main/java/net/snowflake/client/jdbc/SnowflakeFileTransferAgent.java b/src/main/java/net/snowflake/client/jdbc/SnowflakeFileTransferAgent.java index 94c3a323f..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,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 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 @@ -1272,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/FileUploaderPrepIT.java b/src/test/java/net/snowflake/client/jdbc/FileUploaderPrepIT.java index 823929319..dc47d2725 100644 --- a/src/test/java/net/snowflake/client/jdbc/FileUploaderPrepIT.java +++ b/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" diff --git a/src/test/java/net/snowflake/client/jdbc/FileUploaderSessionlessTest.java b/src/test/java/net/snowflake/client/jdbc/FileUploaderSessionlessTest.java index 08aa01c6f..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"); @@ -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 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 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 expectedCreds = new HashMap<>(); Assert.assertEquals(StageInfo.StageType.GCS, stageInfo.getStageType());