Many resources are needed to download a project. Please understand that we have to compensate our server costs. Thank you in advance. Project price only 1 $
You can buy this project and download/modify it how often you want.
/*
* Copyright (c) 2012-2019 Snowflake Computing Inc. All rights reserved.
*/
package net.snowflake.client.jdbc.cloud.storage;
import com.amazonaws.util.Base64;
import com.fasterxml.jackson.core.JsonFactory;
import com.fasterxml.jackson.core.JsonParser;
import com.fasterxml.jackson.databind.JsonNode;
import com.fasterxml.jackson.databind.ObjectMapper;
import com.google.api.gax.paging.Page;
import com.google.auth.oauth2.AccessToken;
import com.google.auth.oauth2.GoogleCredentials;
import com.google.cloud.storage.Blob;
import com.google.cloud.storage.BlobId;
import com.google.cloud.storage.BlobInfo;
import com.google.cloud.storage.Storage;
import com.google.cloud.storage.Storage.BlobListOption;
import com.google.cloud.storage.StorageException;
import com.google.cloud.storage.StorageOptions;
import com.google.common.base.Strings;
import java.io.ByteArrayInputStream;
import java.io.File;
import java.io.FileInputStream;
import java.io.FileNotFoundException;
import java.io.FileOutputStream;
import java.io.IOException;
import java.io.InputStream;
import java.io.OutputStream;
import java.net.SocketTimeoutException;
import java.net.URISyntaxException;
import java.security.InvalidKeyException;
import java.util.AbstractMap;
import java.util.ArrayList;
import java.util.List;
import java.util.Map;
import java.util.Map.Entry;
import net.snowflake.client.core.HttpUtil;
import net.snowflake.client.core.OCSPMode;
import net.snowflake.client.core.ObjectMapperFactory;
import net.snowflake.client.core.SFSession;
import net.snowflake.client.jdbc.ErrorCode;
import net.snowflake.client.jdbc.FileBackedOutputStream;
import net.snowflake.client.jdbc.MatDesc;
import net.snowflake.client.jdbc.RestRequest;
import net.snowflake.client.jdbc.SnowflakeFileTransferAgent;
import net.snowflake.client.jdbc.SnowflakeSQLException;
import net.snowflake.client.jdbc.SnowflakeUtil;
import net.snowflake.client.log.ArgSupplier;
import net.snowflake.client.log.SFLogger;
import net.snowflake.client.log.SFLoggerFactory;
import net.snowflake.client.util.SFPair;
import net.snowflake.common.core.RemoteStoreFileEncryptionMaterial;
import net.snowflake.common.core.SqlState;
import org.apache.commons.io.IOUtils;
import org.apache.http.Header;
import org.apache.http.HttpResponse;
import org.apache.http.client.methods.HttpGet;
import org.apache.http.client.methods.HttpPut;
import org.apache.http.client.utils.URIBuilder;
import org.apache.http.entity.InputStreamEntity;
import org.apache.http.impl.client.CloseableHttpClient;
import static net.snowflake.client.jdbc.SnowflakeUtil.systemGetProperty;
/**
* Encapsulates the GCS Storage client and all GCS operations and logic
*
* @author ppaulus
*/
public class SnowflakeGCSClient implements SnowflakeStorageClient
{
private final static String GCS_ENCRYPTIONDATAPROP = "encryptiondata";
private final static String localFileSep = systemGetProperty("file.separator");
private final static String GCS_METADATA_PREFIX = "x-goog-meta-";
private int encryptionKeySize = 0; // used for PUTs
private StageInfo stageInfo;
private RemoteStoreFileEncryptionMaterial encMat;
private Storage gcsClient = null;
private final static SFLogger logger =
SFLoggerFactory.getLogger(SnowflakeGCSClient.class);
private SnowflakeGCSClient()
{
}
/*
* Factory method for a SnowflakeGCSClient object
* @param stage The stage information that the client will operate on
* @param encMat The encryption material
* required to decrypt/encrypt content in stage
*/
public static SnowflakeGCSClient createSnowflakeGCSClient(StageInfo stage,
RemoteStoreFileEncryptionMaterial encMat)
throws SnowflakeSQLException
{
SnowflakeGCSClient sfGcsClient = new SnowflakeGCSClient();
sfGcsClient.setupGCSClient(stage, encMat);
return sfGcsClient;
}
// Returns the Max number of retry attempts
@Override
public int getMaxRetries()
{
return 25;
}
// Returns the max exponent for multiplying backoff with the power of 2, the value
// of 4 will give us 16secs as the max number of time to sleep before retry
@Override
public int getRetryBackoffMaxExponent()
{
return 4;
}
// Returns the min number of milliseconds to sleep before retry
@Override
public int getRetryBackoffMin()
{
return 1000;
}
/**
* @return Returns true if encryption is enabled
*/
@Override
public boolean isEncrypting()
{
return encryptionKeySize > 0;
}
/**
* @return Returns the size of the encryption key
*/
@Override
public int getEncryptionKeySize()
{
return encryptionKeySize;
}
/**
* @return Whether this client requires the use of presigned URLs for upload
* and download instead of credentials that work for all files uploaded/
* downloaded to a stage path. True for GCS.
*/
@Override
public boolean requirePresignedUrl()
{
return true;
}
@Override
public void renew(Map stageCredentials) throws SnowflakeSQLException
{
stageInfo.setCredentials(stageCredentials);
setupGCSClient(stageInfo, encMat);
}
@Override
public void shutdown()
{
// nothing to do here
}
/**
* listObjects gets all the objects in a path
*
* @param remoteStorageLocation bucket name
* @param prefix Path
* @return
* @throws StorageProviderException
*/
@Override
public StorageObjectSummaryCollection listObjects(String remoteStorageLocation,
String prefix) throws StorageProviderException
{
try
{
Page blobs = this.gcsClient.list(remoteStorageLocation, BlobListOption.prefix(prefix));
return new StorageObjectSummaryCollection(blobs);
}
catch (Exception e)
{
logger.debug("Failed to list objects");
throw new StorageProviderException(e);
}
}
@Override
public StorageObjectMetadata getObjectMetadata(String remoteStorageLocation,
String prefix)
throws StorageProviderException
{
try
{
BlobId blobId = BlobId.of(remoteStorageLocation, prefix);
Blob blob = gcsClient.get(blobId);
// GCS returns null if the blob was not found
// By design, our storage platform expects to see a "blob not found" situation
// as a RemoteStorageProviderException
// Hence, we throw a RemoteStorageProviderException
if (blob == null)
{
throw new StorageProviderException(
new StorageException(
404, // because blob not found
"Blob" + blobId.getName() + " not found in bucket "
+ blobId.getBucket())
);
}
return new CommonObjectMetadata(blob.getSize(),
blob.getContentEncoding(),
blob.getMetadata());
}
catch (StorageException ex)
{
throw new StorageProviderException(ex);
}
}
/**
* Download a file from remote storage.
*
* @param connection connection object
* @param command command to download file
* @param localLocation local file path
* @param destFileName destination file name
* @param parallelism [ not used by the GCP implementation ]
* @param remoteStorageLocation remote storage location, i.e. bucket for S3
* @param stageFilePath stage file path
* @param stageRegion region name where the stage persists
* @param presignedUrl Credential to use for download
* @throws SnowflakeSQLException download failure
**/
@Override
public void download(SFSession connection,
String command,
String localLocation,
String destFileName,
int parallelism,
String remoteStorageLocation,
String stageFilePath,
String stageRegion,
String presignedUrl) throws SnowflakeSQLException
{
int retryCount = 0;
String localFilePath = localLocation + localFileSep + destFileName;
File localFile = new File(localFilePath);
do
{
try
{
String key = null;
String iv = null;
if (!Strings.isNullOrEmpty(presignedUrl))
{
logger.debug("Starting download with presigned URL");
URIBuilder uriBuilder = new URIBuilder(presignedUrl);
HttpGet httpRequest = new HttpGet(uriBuilder.build());
httpRequest.addHeader("accept-encoding", "GZIP");
logger.debug("Fetching result: {}", scrubPresignedUrl(presignedUrl));
CloseableHttpClient httpClient = HttpUtil.getHttpClientWithoutDecompression(
connection.getOCSPMode());
// Put the file on storage using the presigned url
HttpResponse response =
RestRequest.execute(httpClient,
httpRequest,
connection.getNetworkTimeoutInMilli() / 1000, // retry timeout
0, // no socketime injection
null, // no canceling
false, // no cookie
false, // no retry
false, // no request_guid
true // retry on HTTP 403
);
logger.debug("Call returned for URL: {}",
(ArgSupplier) () -> scrubPresignedUrl(this.stageInfo.getPresignedUrl()));
if (response.getStatusLine().getStatusCode() == 200)
{
try
{
InputStream bodyStream = response.getEntity().getContent();
byte[] buffer = new byte[8 * 1024];
int bytesRead;
OutputStream outStream = new FileOutputStream(localFile);
while ((bytesRead = bodyStream.read(buffer)) != -1)
{
outStream.write(buffer, 0, bytesRead);
}
outStream.flush();
outStream.close();
bodyStream.close();
for (Header header : response.getAllHeaders())
{
if (header.getName().equalsIgnoreCase(GCS_METADATA_PREFIX + GCS_ENCRYPTIONDATAPROP))
{
AbstractMap.SimpleEntry encryptionData =
parseEncryptionData(header.getValue());
key = encryptionData.getKey();
iv = encryptionData.getValue();
break;
}
}
logger.debug("Download successful");
}
catch (IOException ex)
{
logger.debug("Download unsuccessful {}", ex);
handleStorageException(ex, ++retryCount, "download", connection, command);
}
}
}
else
{
BlobId blobId = BlobId.of(remoteStorageLocation, stageFilePath);
Blob blob = gcsClient.get(blobId);
if (blob == null)
{
throw new StorageProviderException(
new StorageException(
404, // because blob not found
"Blob" + blobId.getName() + " not found in bucket "
+ blobId.getBucket())
);
}
logger.debug("Starting download without presigned URL");
blob.downloadTo(localFile.toPath());
logger.debug("Download successful");
// Get the user-defined BLOB metadata
Map userDefinedMetadata = blob.getMetadata();
if (userDefinedMetadata != null)
{
AbstractMap.SimpleEntry encryptionData =
parseEncryptionData(userDefinedMetadata.get(GCS_ENCRYPTIONDATAPROP));
key = encryptionData.getKey();
iv = encryptionData.getValue();
}
}
if (!Strings.isNullOrEmpty(iv) && !Strings.isNullOrEmpty(key) &&
this.isEncrypting() && this.getEncryptionKeySize() <= 256)
{
if (key == null || iv == null)
{
throw new SnowflakeSQLException(SqlState.INTERNAL_ERROR,
ErrorCode.INTERNAL_ERROR.getMessageCode(),
"File metadata incomplete");
}
// Decrypt file
try
{
EncryptionProvider.decrypt(localFile, key, iv, this.encMat);
}
catch (Exception ex)
{
logger.error("Error decrypting file", ex);
throw new SnowflakeSQLException(SqlState.INTERNAL_ERROR,
ErrorCode.INTERNAL_ERROR.getMessageCode(),
"Cannot decrypt file");
}
}
return;
}
catch (Exception ex)
{
logger.debug("Download unsuccessful {}", ex);
handleStorageException(ex, ++retryCount, "download", connection, command);
}
}
while (retryCount <= getMaxRetries());
throw new SnowflakeSQLException(SqlState.INTERNAL_ERROR,
ErrorCode.INTERNAL_ERROR.getMessageCode(),
"Unexpected: download unsuccessful without exception!");
}
/**
* Download a file from remote storage
*
* @param connection connection object
* @param command command to download file
* @param parallelism number of threads for parallel downloading
* @param remoteStorageLocation remote storage location, i.e. bucket for s3
* @param stageFilePath stage file path
* @param stageRegion region name where the stage persists
* @param presignedUrl Signed credential for download
* @return input file stream
* @throws SnowflakeSQLException when download failure
*/
@Override
public InputStream downloadToStream(SFSession connection, String command,
int parallelism,
String remoteStorageLocation,
String stageFilePath, String stageRegion,
String presignedUrl)
throws SnowflakeSQLException
{
int retryCount = 0;
InputStream inputStream = null;
do
{
try
{
String key = null;
String iv = null;
if (!Strings.isNullOrEmpty(presignedUrl))
{
logger.debug("Starting download with presigned URL");
URIBuilder uriBuilder = new URIBuilder(presignedUrl);
HttpGet httpRequest = new HttpGet(uriBuilder.build());
httpRequest.addHeader("accept-encoding", "GZIP");
logger.debug("Fetching result: {}", scrubPresignedUrl(presignedUrl));
CloseableHttpClient httpClient = HttpUtil.getHttpClientWithoutDecompression(
connection.getOCSPMode());
// Put the file on storage using the presigned url
HttpResponse response =
RestRequest.execute(httpClient,
httpRequest,
connection.getNetworkTimeoutInMilli() / 1000, // retry timeout
0, // no socketime injection
null, // no canceling
false, // no cookie
false, // no retry
false, // no request_guid
true // retry on HTTP 403
);
logger.debug("Call returned for URL: {}",
(ArgSupplier) () -> scrubPresignedUrl(this.stageInfo.getPresignedUrl()));
if (response.getStatusLine().getStatusCode() == 200)
{
try
{
inputStream = response.getEntity().getContent();
for (Header header : response.getAllHeaders())
{
if (header.getName().equalsIgnoreCase(GCS_METADATA_PREFIX + GCS_ENCRYPTIONDATAPROP))
{
AbstractMap.SimpleEntry encryptionData =
parseEncryptionData(header.getValue());
key = encryptionData.getKey();
iv = encryptionData.getValue();
break;
}
}
logger.debug("Download successful");
}
catch (IOException ex)
{
logger.debug("Download unsuccessful {}", ex);
handleStorageException(ex, ++retryCount, "download", connection, command);
}
}
}
else
{
BlobId blobId = BlobId.of(remoteStorageLocation, stageFilePath);
Blob blob = gcsClient.get(blobId);
if (blob == null)
{
throw new StorageProviderException(
new StorageException(
404, // because blob not found
"Blob" + blobId.getName() + " not found in bucket "
+ blobId.getBucket())
);
}
inputStream = new ByteArrayInputStream(blob.getContent());
// Get the user-defined BLOB metadata
Map userDefinedMetadata = blob.getMetadata();
AbstractMap.SimpleEntry encryptionData =
parseEncryptionData(userDefinedMetadata.get(GCS_ENCRYPTIONDATAPROP));
key = encryptionData.getKey();
iv = encryptionData.getValue();
}
if (this.isEncrypting() && this.getEncryptionKeySize() <= 256)
{
if (key == null || iv == null)
{
throw new SnowflakeSQLException(SqlState.INTERNAL_ERROR,
ErrorCode.INTERNAL_ERROR.getMessageCode(),
"File metadata incomplete");
}
// Decrypt file
try
{
if (inputStream != null)
{
inputStream = EncryptionProvider.decryptStream(inputStream, key, iv, this.encMat);
return inputStream;
}
}
catch (Exception ex)
{
logger.error("Error decrypting file", ex);
throw new SnowflakeSQLException(SqlState.INTERNAL_ERROR,
ErrorCode.INTERNAL_ERROR.getMessageCode(),
"Cannot decrypt file");
}
}
}
catch (Exception ex)
{
logger.debug("Download unsuccessful {}", ex);
handleStorageException(ex, ++retryCount, "download", connection, command);
}
}
while (retryCount <= getMaxRetries());
throw new SnowflakeSQLException(SqlState.INTERNAL_ERROR,
ErrorCode.INTERNAL_ERROR.getMessageCode(),
"Unexpected: download unsuccessful without exception!");
}
/**
* Upload a file (-stream) to remote storage with Pre-signed URL without
* JDBC connection.
*
* @param networkTimeoutInMilli Network timeout for the upload
* @param ocspMode OCSP mode for the upload.
* @param parallelism number of threads do parallel uploading
* @param uploadFromStream true if upload source is stream
* @param remoteStorageLocation s3 bucket name
* @param srcFile source file if not uploading from a stream
* @param destFileName file name on remote storage after upload
* @param inputStream stream used for uploading if
* fileBackedOutputStream is null
* @param fileBackedOutputStream stream used for uploading if not null
* @param meta object meta data
* @param stageRegion region name where the stage persists
* @param presignedUrl presigned URL for upload. Used by GCP.
* @throws SnowflakeSQLException if upload failed
*/
@Override
public void uploadWithPresignedUrlWithoutConnection(
int networkTimeoutInMilli, OCSPMode ocspMode,
int parallelism, boolean uploadFromStream,
String remoteStorageLocation, File srcFile,
String destFileName, InputStream inputStream,
FileBackedOutputStream fileBackedOutputStream,
StorageObjectMetadata meta, String stageRegion,
String presignedUrl)
throws SnowflakeSQLException
{
final List toClose = new ArrayList<>();
long originalContentLength = meta.getContentLength();
SFPair uploadStreamInfo = createUploadStream(
srcFile, uploadFromStream, inputStream, meta,
originalContentLength, fileBackedOutputStream, toClose);
if (!(meta instanceof CommonObjectMetadata))
{
throw new IllegalArgumentException("Unexpected metadata object type");
}
if (Strings.isNullOrEmpty(presignedUrl))
{
throw new IllegalArgumentException("pre-signed URL has to be specified");
}
logger.debug("Starting upload");
uploadWithPresignedUrl(networkTimeoutInMilli,
meta.getContentEncoding(),
meta.getUserMetadata(),
uploadStreamInfo.left,
presignedUrl,
ocspMode);
logger.debug("Upload successful");
// close any open streams in the "toClose" list and return
for (FileInputStream is : toClose)
{
IOUtils.closeQuietly(is);
}
}
/**
* Upload a file/stream to remote storage
*
* @param connection connection object
* @param command upload command
* @param parallelism [ not used by the GCP implementation ]
* @param uploadFromStream true if upload source is stream
* @param remoteStorageLocation storage container name
* @param srcFile source file if not uploading from a stream
* @param destFileName file name on remote storage after upload
* @param inputStream stream used for uploading if fileBackedOutputStream is null
* @param fileBackedOutputStream stream used for uploading if not null
* @param meta object meta data
* @param stageRegion region name where the stage persists
* @param presignedUrl Credential used for upload of a file
* @throws SnowflakeSQLException if upload failed even after retry
*/
@Override
public void upload(SFSession connection,
String command,
int parallelism,
boolean uploadFromStream,
String remoteStorageLocation,
File srcFile,
String destFileName,
InputStream inputStream,
FileBackedOutputStream fileBackedOutputStream,
StorageObjectMetadata meta,
String stageRegion,
String presignedUrl) throws SnowflakeSQLException
{
final List toClose = new ArrayList<>();
long originalContentLength = meta.getContentLength();
SFPair uploadStreamInfo = createUploadStream(
srcFile, uploadFromStream, inputStream, meta, originalContentLength,
fileBackedOutputStream, toClose);
if (!(meta instanceof CommonObjectMetadata))
{
throw new IllegalArgumentException("Unexpected metadata object type");
}
if (!Strings.isNullOrEmpty(presignedUrl))
{
logger.debug("Starting upload");
uploadWithPresignedUrl(connection.getNetworkTimeoutInMilli(),
meta.getContentEncoding(),
meta.getUserMetadata(),
uploadStreamInfo.left,
presignedUrl,
connection.getOCSPMode());
logger.debug("Upload successful");
// close any open streams in the "toClose" list and return
for (FileInputStream is : toClose)
IOUtils.closeQuietly(is);
return;
}
// No presigned URL. This codepath is for when we have a token instead.
int retryCount = 0;
do
{
try
{
logger.debug("Starting upload");
InputStream fileInputStream = uploadStreamInfo.left;
BlobId blobId = BlobId.of(remoteStorageLocation, destFileName);
BlobInfo blobInfo = BlobInfo
.newBuilder(blobId)
.setContentEncoding(meta.getContentEncoding())
.setMetadata(meta.getUserMetadata())
.build();
gcsClient.create(blobInfo, fileInputStream);
logger.debug("Upload successful");
// close any open streams in the "toClose" list and return
for (FileInputStream is : toClose)
IOUtils.closeQuietly(is);
return;
}
catch (Exception ex)
{
handleStorageException(ex, ++retryCount, "upload", connection, command);
if (uploadFromStream && fileBackedOutputStream == null)
{
throw new SnowflakeSQLException(ex, SqlState.SYSTEM_ERROR,
ErrorCode.IO_ERROR.getMessageCode(),
"Encountered exception during upload: " +
ex.getMessage() + "\nCannot retry upload from stream.");
}
uploadStreamInfo = createUploadStream(srcFile, uploadFromStream,
inputStream, meta, originalContentLength,
fileBackedOutputStream, toClose);
}
}
while (retryCount <= getMaxRetries());
for (FileInputStream is : toClose)
IOUtils.closeQuietly(is);
throw new SnowflakeSQLException(SqlState.INTERNAL_ERROR,
ErrorCode.INTERNAL_ERROR.getMessageCode(),
"Unexpected: upload unsuccessful without exception!");
}
/**
* Performs upload using a presigned URL
*
* @param networkTimeoutInMilli Network timeout
* @param contentEncoding Object's content encoding. We do special things for "gzip"
* @param metadata Custom metadata to be uploaded with the object
* @param content File content
* @param presignedUrl Credential to upload the object
* @param ocspMode OCSP mode
* @throws SnowflakeSQLException
*/
private void uploadWithPresignedUrl(int networkTimeoutInMilli,
String contentEncoding,
Map metadata,
InputStream content,
String presignedUrl,
OCSPMode ocspMode)
throws SnowflakeSQLException
{
try
{
URIBuilder uriBuilder = new URIBuilder(presignedUrl);
HttpPut httpRequest = new HttpPut(uriBuilder.build());
logger.debug("Fetching result: {}", scrubPresignedUrl(presignedUrl));
// We set the contentEncoding to blank for GZIP files. We don't want GCS to think
// our gzip files are gzips because it makes them download uncompressed, and
// none of the other providers do that. There's essentially no way for us
// to prevent that behavior. Bad Google.
if ("gzip".equals(contentEncoding))
{
contentEncoding = "";
}
httpRequest.addHeader("content-encoding", contentEncoding);
for (Entry entry : metadata.entrySet())
{
httpRequest.addHeader(GCS_METADATA_PREFIX + entry.getKey(), entry.getValue());
}
InputStreamEntity contentEntity = new InputStreamEntity(content, -1);
httpRequest.setEntity(contentEntity);
CloseableHttpClient httpClient = HttpUtil.getHttpClient(ocspMode);
// Put the file on storage using the presigned url
HttpResponse response =
RestRequest.execute(httpClient,
httpRequest,
networkTimeoutInMilli / 1000, // retry timeout
0, // no socketime injection
null, // no canceling
false, // no cookie
false, // no retry
false, // no request_guid
true // retry on HTTP 403
);
logger.debug("Call returned for URL: {}",
(ArgSupplier) () -> scrubPresignedUrl(this.stageInfo.getPresignedUrl()));
}
catch (URISyntaxException e)
{
throw new SnowflakeSQLException(SqlState.INTERNAL_ERROR,
ErrorCode.INTERNAL_ERROR.getMessageCode(),
"Unexpected: upload presigned URL invalid");
}
catch (Exception e)
{
throw new SnowflakeSQLException(SqlState.INTERNAL_ERROR,
ErrorCode.INTERNAL_ERROR.getMessageCode(),
"Unexpected: upload with presigned url failed");
}
}
/**
* When we log the URL, make sure we don't log the credential
*
* @param presignedUrl Presigned URL with full signature
* @return Just the object path
*/
private String scrubPresignedUrl(String presignedUrl)
{
if (Strings.isNullOrEmpty(presignedUrl))
{
return "";
}
int indexOfQueryString = presignedUrl.lastIndexOf("?");
indexOfQueryString = indexOfQueryString > 0 ?
indexOfQueryString :
presignedUrl.length() - 1;
return presignedUrl.substring(0, indexOfQueryString);
}
private SFPair createUploadStream(
File srcFile,
boolean uploadFromStream,
InputStream inputStream,
StorageObjectMetadata meta,
long originalContentLength,
FileBackedOutputStream fileBackedOutputStream,
List toClose)
throws SnowflakeSQLException
{
logger.debug(
"createUploadStream({}, {}, {}, {}, {}, {})",
this, srcFile, uploadFromStream, inputStream, fileBackedOutputStream, toClose);
final InputStream stream;
FileInputStream srcFileStream = null;
try
{
if (isEncrypting() && getEncryptionKeySize() < 256)
{
try
{
final InputStream uploadStream = uploadFromStream ?
(fileBackedOutputStream != null ?
fileBackedOutputStream.asByteSource().openStream() :
inputStream) :
(srcFileStream = new FileInputStream(srcFile));
toClose.add(srcFileStream);
// Encrypt
stream = EncryptionProvider.encrypt(meta, originalContentLength,
uploadStream, this.encMat, this);
uploadFromStream = true;
}
catch (Exception ex)
{
logger.error("Failed to encrypt input", ex);
throw new SnowflakeSQLException(ex, SqlState.INTERNAL_ERROR,
ErrorCode.INTERNAL_ERROR.getMessageCode(),
"Failed to encrypt input", ex.getMessage());
}
}
else
{
if (uploadFromStream)
{
if (fileBackedOutputStream != null)
{
stream = fileBackedOutputStream.asByteSource().openStream();
}
else
{
stream = inputStream;
}
}
else
{
srcFileStream = new FileInputStream(srcFile);
toClose.add(srcFileStream);
stream = srcFileStream;
}
}
}
catch (FileNotFoundException ex)
{
logger.error("Failed to open input file", ex);
throw new SnowflakeSQLException(ex, SqlState.INTERNAL_ERROR,
ErrorCode.INTERNAL_ERROR.getMessageCode(),
"Failed to open input file", ex.getMessage());
}
catch (IOException ex)
{
logger.error("Failed to open input stream", ex);
throw new SnowflakeSQLException(ex, SqlState.INTERNAL_ERROR,
ErrorCode.INTERNAL_ERROR.getMessageCode(),
"Failed to open input stream", ex.getMessage());
}
return SFPair.of(stream, uploadFromStream);
}
@Override
public void handleStorageException(Exception ex,
int retryCount,
String operation,
SFSession connection,
String command)
throws SnowflakeSQLException
{
// no need to retry if it is invalid key exception
if (ex.getCause() instanceof InvalidKeyException)
{
// Most likely cause is that the unlimited strength policy files are not installed
// Log the error and throw a message that explains the cause
SnowflakeFileTransferAgent.throwJCEMissingError(operation, ex);
}
if (ex instanceof StorageException)
{
StorageException se = (StorageException) ex;
if (se.getCode() == 403 && connection != null && command != null)
{
// A 403 indicates that the access token has expired,
// we need to refresh the GCS client with the new token
SnowflakeFileTransferAgent.renewExpiredToken(connection, command, this);
}
// If we have exceeded the max number of retries, propagate the error
if (retryCount > getMaxRetries())
{
throw new SnowflakeSQLException(se, SqlState.SYSTEM_ERROR,
ErrorCode.GCP_SERVICE_ERROR.getMessageCode(),
operation,
se.getCode(),
se.getMessage(),
se.getReason());
}
else
{
logger.debug("Encountered exception ({}) during {}, retry count: {}",
ex.getMessage(), operation, retryCount);
logger.debug("Stack trace: ", ex);
// exponential backoff up to a limit
int backoffInMillis = getRetryBackoffMin();
if (retryCount > 1)
{
backoffInMillis <<= (Math.min(retryCount - 1,
getRetryBackoffMaxExponent()));
}
try
{
logger.debug("Sleep for {} milliseconds before retry", backoffInMillis);
Thread.sleep(backoffInMillis);
}
catch (InterruptedException ex1)
{
// ignore
}
if (se.getCode() == 403 && connection != null && command != null)
{
// A 403 indicates that the access token has expired,
// we need to refresh the GCS client with the new token
SnowflakeFileTransferAgent.renewExpiredToken(connection, command, this);
}
}
}
else
{
if (ex instanceof InterruptedException ||
SnowflakeUtil.getRootCause(ex) instanceof SocketTimeoutException)
{
if (retryCount > getMaxRetries())
{
throw new SnowflakeSQLException(ex, SqlState.SYSTEM_ERROR,
ErrorCode.IO_ERROR.getMessageCode(),
"Encountered exception during " + operation + ": " +
ex.getMessage());
}
else
{
logger.debug("Encountered exception ({}) during {}, retry count: {}",
ex.getMessage(), operation, retryCount);
}
}
else
{
throw new SnowflakeSQLException(ex, SqlState.SYSTEM_ERROR,
ErrorCode.IO_ERROR.getMessageCode(),
"Encountered exception during " + operation + ": " +
ex.getMessage());
}
}
}
/**
* Returns the material descriptor key
*/
@Override
public String getMatdescKey()
{
return "matdesc";
}
/**
* Adds encryption metadata to the StorageObjectMetadata object
*/
@Override
public void addEncryptionMetadata(StorageObjectMetadata meta,
MatDesc matDesc,
byte[] ivData,
byte[] encKeK,
long contentLength)
{
meta.addUserMetadata(getMatdescKey(),
matDesc.toString());
meta.addUserMetadata(GCS_ENCRYPTIONDATAPROP, buildEncryptionMetadataJSON(
Base64.encodeAsString(ivData),
Base64.encodeAsString(encKeK))
);
meta.setContentLength(contentLength);
}
/*
* buildEncryptionMetadataJSON
* Takes the base64-encoded iv and key and creates the JSON block to be
* used as the encryptiondata metadata field on the blob.
*/
private String buildEncryptionMetadataJSON(String iv64, String key64)
{
return String.format("{\"EncryptionMode\":\"FullBlob\",\"WrappedContentKey\""
+ ":{\"KeyId\":\"symmKey1\",\"EncryptedKey\":\"%s\""
+ ",\"Algorithm\":\"AES_CBC_256\"},\"EncryptionAgent\":"
+ "{\"Protocol\":\"1.0\",\"EncryptionAlgorithm\":"
+ "\"AES_CBC_256\"},\"ContentEncryptionIV\":\"%s\""
+ ",\"KeyWrappingMetadata\":{\"EncryptionLibrary\":"
+ "\"Java 5.3.0\"}}", key64, iv64);
}
/*
* parseEncryptionData
* Takes the json string in the encryptiondata metadata field of the encrypted
* blob and parses out the key and iv. Returns the pair as key = key, iv = value.
*/
private AbstractMap.SimpleEntry parseEncryptionData(String jsonEncryptionData)
throws SnowflakeSQLException
{
ObjectMapper mapper = ObjectMapperFactory.getObjectMapper();
JsonFactory factory = mapper.getFactory();
try
{
JsonParser parser = factory.createParser(jsonEncryptionData);
JsonNode encryptionDataNode = mapper.readTree(parser);
String iv = encryptionDataNode.get("ContentEncryptionIV").asText();
String key = encryptionDataNode.get("WrappedContentKey").get("EncryptedKey").asText();
return new AbstractMap.SimpleEntry<>(key, iv);
}
catch (Exception ex)
{
throw new SnowflakeSQLException(ex, SqlState.SYSTEM_ERROR,
ErrorCode.IO_ERROR.getMessageCode(),
"Error parsing encryption data as json" + ": " +
ex.getMessage());
}
}
/**
* Adds digest metadata to the StorageObjectMetadata object
*/
@Override
public void addDigestMetadata(StorageObjectMetadata meta, String digest)
{
if (!SnowflakeUtil.isBlank(digest))
{
meta.addUserMetadata("sfc-digest", digest);
}
}
/**
* Gets digest metadata to the StorageObjectMetadata object
*/
@Override
public String getDigestMetadata(StorageObjectMetadata meta)
{
return meta.getUserMetadata().get("sfc-digest");
}
/*
* Initializes the GCS client
* This method is used during the object contruction, but also to
* reset/recreate the encapsulated CloudBlobClient object with new
* credentials (after token expiration)
* @param stage The stage information that the client will operate on
* @param encMat The encryption material
* required to decrypt/encrypt content in stage
* @throws IllegalArgumentException when invalid credentials are used
*/
private void setupGCSClient(StageInfo stage, RemoteStoreFileEncryptionMaterial encMat)
throws IllegalArgumentException, SnowflakeSQLException
{
// Save the client creation parameters so that we can reuse them,
// to reset the GCS client.
this.stageInfo = stage;
this.encMat = encMat;
logger.debug("Setting up the GCS client ");
try
{
String accessToken = (String) stage.getCredentials().get("GCS_ACCESS_TOKEN");
GoogleCredentials googleCreds;
if (accessToken != null)
{
AccessToken googleAccessToken = new AccessToken(accessToken, null);
googleCreds = GoogleCredentials.create(googleAccessToken);
// We are authenticated with an oauth access token.
this.gcsClient = StorageOptions.newBuilder()
.setCredentials(googleCreds)
.build()
.getService();
}
else
{
// Use anonymous authentication.
this.gcsClient = StorageOptions.getUnauthenticatedInstance()
.getService();
}
if (encMat != null)
{
byte[] decodedKey = Base64.decode(encMat.getQueryStageMasterKey());
encryptionKeySize = decodedKey.length * 8;
if (encryptionKeySize != 128 &&
encryptionKeySize != 192 &&
encryptionKeySize != 256)
{
throw new SnowflakeSQLException(SqlState.INTERNAL_ERROR,
ErrorCode.INTERNAL_ERROR.getMessageCode(),
"unsupported key size", encryptionKeySize);
}
}
}
catch (Exception ex)
{
throw new IllegalArgumentException("invalid_gcs_credentials");
}
}
}
