Skip to content

Commit

Permalink
SNOW-1789749: Support regional GCS endpoints
Browse files Browse the repository at this point in the history
  • Loading branch information
@sfc-gh-dprzybysz
sfc-gh-dprzybysz committed Nov 22, 2024
1 parent 1286a3e commit 1793940
Show file tree
Hide file tree
Showing 7 changed files with 278 additions and 19 deletions.
23 changes: 22 additions & 1 deletion src/main/java/net/snowflake/client/jdbc/SnowflakeFileTransferAgent.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -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
Expand All @@ -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.
Expand Down
9 changes: 8 additions & 1 deletion src/main/java/net/snowflake/client/jdbc/cloud/storage/SnowflakeGCSClient.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -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");
Expand All @@ -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) {
Expand Down
46 changes: 41 additions & 5 deletions src/main/java/net/snowflake/client/jdbc/cloud/storage/StageInfo.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand All @@ -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;

/*
Expand Down Expand Up @@ -166,16 +179,39 @@ 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(""));
}

public void setProxyProperties(Properties proxyProperties) {
this.proxyProperties = proxyProperties;
}
;

public Properties getProxyProperties() {
return proxyProperties;
}

@SnowflakeJdbcInternalApi
public Optional<String> 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();
}
}
5 changes: 3 additions & 2 deletions src/main/java/net/snowflake/client/jdbc/cloud/storage/StorageClientFactory.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down
102 changes: 102 additions & 0 deletions src/test/java/net/snowflake/client/jdbc/FileUploaderPrep.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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<JsonNode> exampleNodes;

@Before
Expand All @@ -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);
}
}
Loading

0 comments on commit 1793940

Please sign in to comment.