diff --git a/src/main/java/net/snowflake/client/jdbc/SnowflakeFileTransferAgent.java b/src/main/java/net/snowflake/client/jdbc/SnowflakeFileTransferAgent.java index 2b660eb08..a573a7466 100644 --- a/src/main/java/net/snowflake/client/jdbc/SnowflakeFileTransferAgent.java +++ b/src/main/java/net/snowflake/client/jdbc/SnowflakeFileTransferAgent.java @@ -1111,8 +1111,15 @@ static StageInfo getStageInfo(JsonNode jsonNode, SFSession session) throws Snowf // specifically // for FIPS or VPCE S3 endpoint. SNOW-652696 String endPoint = null; - if ("AZURE".equalsIgnoreCase(stageLocationType) || "S3".equalsIgnoreCase(stageLocationType)) { + if ("AZURE".equalsIgnoreCase(stageLocationType) + || "S3".equalsIgnoreCase(stageLocationType) + || "GCS".equalsIgnoreCase(stageLocationType)) { endPoint = jsonNode.path("data").path("stageInfo").findValue("endPoint").asText(); + if ("GCS".equalsIgnoreCase(stageLocationType) + && (endPoint.trim().isEmpty() || "null".equals(endPoint))) { + // setting to null to preserve previous behaviour for GCS + endPoint = null; + } } String stgAcct = null; @@ -1179,6 +1186,8 @@ static StageInfo getStageInfo(JsonNode jsonNode, SFSession session) throws Snowf } } + setupUseRegionalUrl(jsonNode, stageInfo); + 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 @@ -1200,6 +1209,18 @@ static StageInfo getStageInfo(JsonNode jsonNode, SFSession session) throws Snowf return stageInfo; } + private static void setupUseRegionalUrl(JsonNode jsonNode, StageInfo stageInfo) { + if (stageInfo.getStageType() != StageInfo.StageType.GCS + && stageInfo.getStageType() != StageInfo.StageType.S3) { + return; + } + JsonNode useRegionalURLNode = jsonNode.path("data").path("stageInfo").path("useRegionalUrl"); + if (!useRegionalURLNode.isMissingNode()) { + boolean useRegionalURL = useRegionalURLNode.asBoolean(false); + stageInfo.setUseRegionalUrl(useRegionalURL); + } + } + /** * A helper method to verify if the local file path from GS matches what's parsed locally. This is * for security purpose as documented in SNOW-15153. 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 index 188ba40d4..d6bf6ba84 100644 --- a/src/main/java/net/snowflake/client/jdbc/cloud/storage/SnowflakeGCSClient.java +++ b/src/main/java/net/snowflake/client/jdbc/cloud/storage/SnowflakeGCSClient.java @@ -18,9 +18,11 @@ import com.google.api.gax.rpc.FixedHeaderProvider; import com.google.auth.oauth2.AccessToken; import com.google.auth.oauth2.GoogleCredentials; +import com.google.cloud.NoCredentials; import com.google.cloud.storage.Blob; import com.google.cloud.storage.BlobId; import com.google.cloud.storage.BlobInfo; +import com.google.cloud.storage.HttpStorageOptions; import com.google.cloud.storage.Storage; import com.google.cloud.storage.Storage.BlobListOption; import com.google.cloud.storage.StorageException; @@ -1312,6 +1314,8 @@ private void setupGCSClient( if (accessToken != null) { // We are authenticated with an oauth access token. StorageOptions.Builder builder = StorageOptions.newBuilder(); + stage.gcsCustomEndpoint().ifPresent(builder::setHost); + if (areDisabledGcsDefaultCredentials(session)) { logger.debug( "Adding explicit credentials to avoid default credential lookup by the GCS client"); @@ -1329,7 +1333,10 @@ private void setupGCSClient( .getService(); } else { // Use anonymous authentication. - this.gcsClient = StorageOptions.getUnauthenticatedInstance().getService(); + HttpStorageOptions.Builder builder = + HttpStorageOptions.newBuilder().setCredentials(NoCredentials.getInstance()); + stage.gcsCustomEndpoint().ifPresent(builder::setHost); + this.gcsClient = builder.build().getService(); } if (encMat != null) { diff --git a/src/main/java/net/snowflake/client/jdbc/cloud/storage/StageInfo.java b/src/main/java/net/snowflake/client/jdbc/cloud/storage/StageInfo.java index 7a8bf4d36..3a14b8fa0 100644 --- a/src/main/java/net/snowflake/client/jdbc/cloud/storage/StageInfo.java +++ b/src/main/java/net/snowflake/client/jdbc/cloud/storage/StageInfo.java @@ -2,10 +2,17 @@ import java.io.Serializable; import java.util.Map; +import java.util.Optional; import java.util.Properties; +import net.snowflake.client.core.SnowflakeJdbcInternalApi; -/** Encapsulates all the required stage properties used by GET/PUT for Azure and S3 stages */ +/** Encapsulates all the required stage properties used by GET/PUT for Azure, GCS and S3 stages */ public class StageInfo implements Serializable { + + // me-central2 GCS region always use regional urls + // TODO SNOW-1818804: the value is hardcoded now, but it should be server driven + private static final String GCS_REGION_ME_CENTRAL_2 = "me-central2"; + public enum StageType { S3, AZURE, @@ -17,12 +24,18 @@ public enum StageType { private StageType stageType; // The stage type private String location; // The container or bucket private Map credentials; // the credentials required for the stage - private String region; // AWS/S3/GCS region (S3/GCS only) - private String endPoint; // The Azure Storage endpoint (Azure only) + private String region; // S3/GCS region + // An endpoint (Azure, AWS FIPS and GCS custom endpoint override) + private String endPoint; private String storageAccount; // The Azure Storage account (Azure only) private String presignedUrl; // GCS gives us back a presigned URL instead of a cred private boolean isClientSideEncrypted; // whether to encrypt/decrypt files on the stage - private boolean useS3RegionalUrl; // whether to use s3 regional URL (AWS Only) + // whether to use s3 regional URL (AWS Only) + // TODO SNOW-1818804: this field will be deprecated when the server returns {@link + // #useRegionalUrl} + private boolean useS3RegionalUrl; + // whether to use regional URL (AWS and GCS only) + private boolean useRegionalUrl; private Properties proxyProperties; /* @@ -166,6 +179,16 @@ public boolean getUseS3RegionalUrl() { return useS3RegionalUrl; } + @SnowflakeJdbcInternalApi + public void setUseRegionalUrl(boolean useRegionalUrl) { + this.useRegionalUrl = useRegionalUrl; + } + + @SnowflakeJdbcInternalApi + public boolean getUseRegionalUrl() { + return useRegionalUrl; + } + private static boolean isSpecified(String arg) { return !(arg == null || arg.equalsIgnoreCase("")); } @@ -173,9 +196,22 @@ private static boolean isSpecified(String arg) { public void setProxyProperties(Properties proxyProperties) { this.proxyProperties = proxyProperties; } - ; public Properties getProxyProperties() { return proxyProperties; } + + @SnowflakeJdbcInternalApi + public Optional gcsCustomEndpoint() { + if (stageType != StageType.GCS) { + return Optional.empty(); + } + if (endPoint != null && !endPoint.trim().isEmpty() && !"null".equals(endPoint)) { + return Optional.of(endPoint); + } + if (GCS_REGION_ME_CENTRAL_2.equalsIgnoreCase(region) || useRegionalUrl) { + return Optional.of(String.format("storage.%s.rep.googleapis.com", region.toLowerCase())); + } + return Optional.empty(); + } } diff --git a/src/main/java/net/snowflake/client/jdbc/cloud/storage/StorageClientFactory.java b/src/main/java/net/snowflake/client/jdbc/cloud/storage/StorageClientFactory.java index a321b6ebd..69d56e195 100644 --- a/src/main/java/net/snowflake/client/jdbc/cloud/storage/StorageClientFactory.java +++ b/src/main/java/net/snowflake/client/jdbc/cloud/storage/StorageClientFactory.java @@ -59,8 +59,9 @@ public SnowflakeStorageClient createClient( switch (stage.getStageType()) { case S3: boolean useS3RegionalUrl = - (stage.getUseS3RegionalUrl() - || (session != null && session.getUseRegionalS3EndpointsForPresignedURL())); + stage.getUseS3RegionalUrl() + || stage.getUseRegionalUrl() + || session != null && session.getUseRegionalS3EndpointsForPresignedURL(); return createS3Client( stage.getCredentials(), parallel, diff --git a/src/test/java/net/snowflake/client/jdbc/FileUploaderPrep.java b/src/test/java/net/snowflake/client/jdbc/FileUploaderPrep.java index d8aa8143f..90eb76866 100644 --- a/src/test/java/net/snowflake/client/jdbc/FileUploaderPrep.java +++ b/src/test/java/net/snowflake/client/jdbc/FileUploaderPrep.java @@ -241,11 +241,111 @@ abstract class FileUploaderPrep extends BaseJDBCTest { + " \"message\": null,\n" + " \"success\": true\n" + "}"; + private final String exampleGCSJsonStringWithUseRegionalUrl = + "{\n" + + " \"data\": {\n" + + " \"uploadInfo\": {\n" + + " \"locationType\": \"GCS\",\n" + + " \"useRegionalUrl\": true,\n" + + " \"location\": \"foo/tables/9224/\",\n" + + " \"path\": \"tables/9224/\",\n" + + " \"region\": \"US-WEST1\",\n" + + " \"storageAccount\": \"\",\n" + + " \"isClientSideEncrypted\": true,\n" + + " \"creds\": {},\n" + + " \"presignedUrl\": \"EXAMPLE_PRESIGNED_URL\",\n" + + " \"endPoint\": \"\"\n" + + " },\n" + + " \"src_locations\": [\n" + + " \"/foo/bart/orders_100.csv\"\n" + + " ],\n" + + " \"parallel\": 4,\n" + + " \"threshold\": 209715200,\n" + + " \"autoCompress\": true,\n" + + " \"overwrite\": false,\n" + + " \"sourceCompression\": \"auto_detect\",\n" + + " \"clientShowEncryptionParameter\": false,\n" + + " \"queryId\": \"EXAMPLE_QUERY_ID\",\n" + + " \"encryptionMaterial\": {\n" + + " \"queryStageMasterKey\": \"EXAMPLE_QUERY_STAGE_MASTER_KEY\",\n" + + " \"queryId\": \"EXAMPLE_QUERY_ID\",\n" + + " \"smkId\": 123\n" + + " },\n" + + " \"stageInfo\": {\n" + + " \"locationType\": \"GCS\",\n" + + " \"useRegionalUrl\": true,\n" + + " \"location\": \"foo/tables/9224/\",\n" + + " \"path\": \"tables/9224/\",\n" + + " \"region\": \"US-WEST1\",\n" + + " \"storageAccount\": \"\",\n" + + " \"isClientSideEncrypted\": true,\n" + + " \"creds\": {},\n" + + " \"presignedUrl\": \"EXAMPLE_PRESIGNED_URL\",\n" + + " \"endPoint\": \"\"\n" + + " },\n" + + " \"command\": \"UPLOAD\",\n" + + " \"kind\": null,\n" + + " \"operation\": \"Node\"\n" + + " },\n" + + " \"code\": null,\n" + + " \"message\": null,\n" + + " \"success\": true\n" + + "}"; + private final String exampleGCSJsonStringWithEndpoint = + "{\n" + + " \"data\": {\n" + + " \"uploadInfo\": {\n" + + " \"locationType\": \"GCS\",\n" + + " \"location\": \"foo/tables/9224/\",\n" + + " \"path\": \"tables/9224/\",\n" + + " \"region\": \"US-WEST1\",\n" + + " \"storageAccount\": \"\",\n" + + " \"isClientSideEncrypted\": true,\n" + + " \"creds\": {},\n" + + " \"presignedUrl\": \"EXAMPLE_PRESIGNED_URL\",\n" + + " \"endPoint\": \"example.com\"\n" + + " },\n" + + " \"src_locations\": [\n" + + " \"/foo/bart/orders_100.csv\"\n" + + " ],\n" + + " \"parallel\": 4,\n" + + " \"threshold\": 209715200,\n" + + " \"autoCompress\": true,\n" + + " \"overwrite\": false,\n" + + " \"sourceCompression\": \"auto_detect\",\n" + + " \"clientShowEncryptionParameter\": false,\n" + + " \"queryId\": \"EXAMPLE_QUERY_ID\",\n" + + " \"encryptionMaterial\": {\n" + + " \"queryStageMasterKey\": \"EXAMPLE_QUERY_STAGE_MASTER_KEY\",\n" + + " \"queryId\": \"EXAMPLE_QUERY_ID\",\n" + + " \"smkId\": 123\n" + + " },\n" + + " \"stageInfo\": {\n" + + " \"locationType\": \"GCS\",\n" + + " \"location\": \"foo/tables/9224/\",\n" + + " \"path\": \"tables/9224/\",\n" + + " \"region\": \"US-WEST1\",\n" + + " \"storageAccount\": \"\",\n" + + " \"isClientSideEncrypted\": true,\n" + + " \"creds\": {},\n" + + " \"presignedUrl\": \"EXAMPLE_PRESIGNED_URL\",\n" + + " \"endPoint\": \"example.com\"\n" + + " },\n" + + " \"command\": \"UPLOAD\",\n" + + " \"kind\": null,\n" + + " \"operation\": \"Node\"\n" + + " },\n" + + " \"code\": null,\n" + + " \"message\": null,\n" + + " \"success\": true\n" + + "}"; protected JsonNode exampleS3JsonNode; protected JsonNode exampleS3StageEndpointJsonNode; protected JsonNode exampleAzureJsonNode; protected JsonNode exampleGCSJsonNode; + protected JsonNode exampleGCSJsonNodeWithUseRegionalUrl; + protected JsonNode exampleGCSJsonNodeWithEndPoint; protected List exampleNodes; @Before @@ -254,6 +354,8 @@ public void setup() throws Exception { exampleS3StageEndpointJsonNode = mapper.readTree(exampleS3JsonStringWithStageEndpoint); exampleAzureJsonNode = mapper.readTree(exampleAzureJsonString); exampleGCSJsonNode = mapper.readTree(exampleGCSJsonString); + exampleGCSJsonNodeWithUseRegionalUrl = mapper.readTree(exampleGCSJsonStringWithUseRegionalUrl); + exampleGCSJsonNodeWithEndPoint = mapper.readTree(exampleGCSJsonStringWithEndpoint); exampleNodes = Arrays.asList(exampleS3JsonNode, exampleAzureJsonNode, exampleGCSJsonNode); } } diff --git a/src/test/java/net/snowflake/client/jdbc/FileUploaderSessionlessTest.java b/src/test/java/net/snowflake/client/jdbc/FileUploaderSessionlessTest.java index e23800e4e..94fd734f2 100644 --- a/src/test/java/net/snowflake/client/jdbc/FileUploaderSessionlessTest.java +++ b/src/test/java/net/snowflake/client/jdbc/FileUploaderSessionlessTest.java @@ -3,6 +3,10 @@ */ package net.snowflake.client.jdbc; +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertFalse; +import static org.junit.Assert.assertTrue; + import com.fasterxml.jackson.core.JsonProcessingException; import com.fasterxml.jackson.databind.JsonNode; import com.fasterxml.jackson.databind.ObjectMapper; @@ -11,6 +15,7 @@ import java.util.HashMap; import java.util.List; import java.util.Map; +import java.util.Optional; import net.snowflake.client.jdbc.cloud.storage.StageInfo; import net.snowflake.common.core.RemoteStoreFileEncryptionMaterial; import org.junit.Assert; @@ -265,6 +270,7 @@ public void testGetFileTransferMetadatasGCS() throws Exception { Assert.assertEquals(null, stageInfo.getEndPoint()); Assert.assertEquals(null, stageInfo.getStorageAccount()); Assert.assertEquals(true, stageInfo.getIsClientSideEncrypted()); + assertEquals(Optional.empty(), stageInfo.gcsCustomEndpoint()); // EncryptionMaterial check Assert.assertEquals("EXAMPLE_QUERY_ID", metadata.getEncryptionMaterial().getQueryId()); @@ -279,12 +285,42 @@ public void testGetFileTransferMetadatasGCS() throws Exception { Assert.assertEquals("orders_100.csv", metadata.getPresignedUrlFileName()); } + @Test + public void testGetFileTransferMetadataGCSWithUseRegionalUrl() throws Exception { + List metadataList = + SnowflakeFileTransferAgent.getFileTransferMetadatas(exampleGCSJsonNodeWithUseRegionalUrl); + Assert.assertEquals(1, metadataList.size()); + + SnowflakeFileTransferMetadataV1 metadata = + (SnowflakeFileTransferMetadataV1) metadataList.get(0); + + StageInfo stageInfo = metadata.getStageInfo(); + + assertTrue(stageInfo.getUseRegionalUrl()); + assertEquals(Optional.of("storage.us-west1.rep.googleapis.com"), stageInfo.gcsCustomEndpoint()); + } + + @Test + public void testGetFileTransferMetadataGCSWithEndPoint() throws Exception { + List metadataList = + SnowflakeFileTransferAgent.getFileTransferMetadatas(exampleGCSJsonNodeWithEndPoint); + Assert.assertEquals(1, metadataList.size()); + + SnowflakeFileTransferMetadataV1 metadata = + (SnowflakeFileTransferMetadataV1) metadataList.get(0); + + StageInfo stageInfo = metadata.getStageInfo(); + + assertFalse(stageInfo.getUseRegionalUrl()); + assertEquals(Optional.of("example.com"), stageInfo.gcsCustomEndpoint()); + } + @Test public void testGetFileTransferMetadatasUploadError() throws Exception { JsonNode downloadNode = mapper.readTree("{\"data\": {\"command\": \"DOWNLOAD\"}}"); try { SnowflakeFileTransferAgent.getFileTransferMetadatas(downloadNode); - Assert.assertTrue(false); + assertTrue(false); } catch (SnowflakeSQLException err) { Assert.assertEquals((long) ErrorCode.INTERNAL_ERROR.getMessageCode(), err.getErrorCode()); Assert.assertEquals( @@ -297,10 +333,10 @@ public void testGetFileTransferMetadatasEncryptionMaterialError() throws Excepti JsonNode garbageNode = mapper.readTree("{\"data\": {\"src_locations\": [1, 2]}}"); try { SnowflakeFileTransferAgent.getFileTransferMetadatas(garbageNode); - Assert.assertTrue(false); + assertTrue(false); } catch (SnowflakeSQLException err) { Assert.assertEquals((long) ErrorCode.INTERNAL_ERROR.getMessageCode(), err.getErrorCode()); - Assert.assertTrue( + assertTrue( err.getMessage().contains("JDBC driver internal error: Failed to parse the credentials")); } } @@ -312,11 +348,10 @@ public void testGetFileTransferMetadatasUnsupportedLocationError() throws Except foo.put("locationType", "LOCAL_FS"); try { SnowflakeFileTransferAgent.getFileTransferMetadatas(modifiedNode); - Assert.assertTrue(false); + assertTrue(false); } catch (SnowflakeSQLException err) { Assert.assertEquals((long) ErrorCode.INTERNAL_ERROR.getMessageCode(), err.getErrorCode()); - Assert.assertTrue( - err.getMessage().contains("JDBC driver internal error: This API only supports")); + assertTrue(err.getMessage().contains("JDBC driver internal error: This API only supports")); } } @@ -325,10 +360,10 @@ public void testGetFileTransferMetadatasSrcLocationsArrayError() throws JsonProc JsonNode garbageNode = mapper.readTree("{\"data\": {\"src_locations\": \"abc\"}}"); try { SnowflakeFileTransferAgent.getFileTransferMetadatas(garbageNode); - Assert.assertTrue(false); + assertTrue(false); } catch (SnowflakeSQLException err) { Assert.assertEquals((long) ErrorCode.INTERNAL_ERROR.getMessageCode(), err.getErrorCode()); - Assert.assertTrue( + assertTrue( err.getMessage().contains("JDBC driver internal error: src_locations must be an array")); } } @@ -340,10 +375,10 @@ public void testGetFileMetadatasEncryptionMaterialsException() { foo.put("encryptionMaterial", "[1, 2, 3]]"); try { SnowflakeFileTransferAgent.getFileTransferMetadatas(modifiedNode); - Assert.assertTrue(false); + assertTrue(false); } catch (SnowflakeSQLException err) { Assert.assertEquals((long) ErrorCode.INTERNAL_ERROR.getMessageCode(), err.getErrorCode()); - Assert.assertTrue(err.getMessage().contains("Failed to parse encryptionMaterial")); + assertTrue(err.getMessage().contains("Failed to parse encryptionMaterial")); } } } diff --git a/src/test/java/net/snowflake/client/jdbc/cloud/storage/StageInfoGcsCustomEndpointTest.java b/src/test/java/net/snowflake/client/jdbc/cloud/storage/StageInfoGcsCustomEndpointTest.java new file mode 100644 index 000000000..f8e00d7eb --- /dev/null +++ b/src/test/java/net/snowflake/client/jdbc/cloud/storage/StageInfoGcsCustomEndpointTest.java @@ -0,0 +1,57 @@ +/* + * Copyright (c) 2024 Snowflake Computing Inc. All rights reserved. + */ +package net.snowflake.client.jdbc.cloud.storage; + +import static org.junit.Assert.assertEquals; + +import java.util.HashMap; +import java.util.Optional; +import org.junit.Test; +import org.junit.runner.RunWith; +import org.junit.runners.Parameterized; + +@RunWith(Parameterized.class) +public class StageInfoGcsCustomEndpointTest { + private final String region; + private final boolean useRegionalUrl; + private final String endPoint; + private final Optional expectedHost; + + public StageInfoGcsCustomEndpointTest( + String region, boolean useRegionalUrl, String endPoint, Optional expectedHost) { + this.region = region; + this.useRegionalUrl = useRegionalUrl; + this.endPoint = endPoint; + this.expectedHost = expectedHost; + } + + @Test + public void shouldReturnEmptyGCSRegionalUrlWhenNotMeCentral1AndNotUseRegionalUrl() { + StageInfo stageInfo = + StageInfo.createStageInfo("GCS", "bla", new HashMap<>(), region, endPoint, "account", true); + stageInfo.setUseRegionalUrl(useRegionalUrl); + assertEquals(expectedHost, stageInfo.gcsCustomEndpoint()); + } + + @Parameterized.Parameters() + public static Object[][] data() { + return new Object[][] { + {"US-CENTRAL1", false, null, Optional.empty()}, + {"US-CENTRAL1", false, "", Optional.empty()}, + {"US-CENTRAL1", false, "null", Optional.empty()}, + {"US-CENTRAL1", false, " ", Optional.empty()}, + {"US-CENTRAL1", false, "example.com", Optional.of("example.com")}, + {"ME-CENTRAL2", false, null, Optional.of("storage.me-central2.rep.googleapis.com")}, + {"ME-CENTRAL2", true, null, Optional.of("storage.me-central2.rep.googleapis.com")}, + {"ME-CENTRAL2", true, "", Optional.of("storage.me-central2.rep.googleapis.com")}, + {"ME-CENTRAL2", true, " ", Optional.of("storage.me-central2.rep.googleapis.com")}, + {"ME-CENTRAL2", true, "example.com", Optional.of("example.com")}, + {"US-CENTRAL1", true, null, Optional.of("storage.us-central1.rep.googleapis.com")}, + {"US-CENTRAL1", true, "", Optional.of("storage.us-central1.rep.googleapis.com")}, + {"US-CENTRAL1", true, " ", Optional.of("storage.us-central1.rep.googleapis.com")}, + {"US-CENTRAL1", true, "null", Optional.of("storage.us-central1.rep.googleapis.com")}, + {"US-CENTRAL1", true, "example.com", Optional.of("example.com")}, + }; + } +}