snowflake-sdk
Version:
Node.js driver for Snowflake
402 lines • 15.6 kB
JavaScript
"use strict";
const { NodeHttpHandler } = require('@smithy/node-http-handler');
const fs = require('fs');
const EncryptionMetadata = require('./encrypt_util').EncryptionMetadata;
const FileHeader = require('../file_util').FileHeader;
const expandTilde = require('expand-tilde');
const getProxyAgent = require('../http/node').getProxyAgent;
const ProxyUtil = require('../proxy_util');
const Logger = require('../logger').default;
const GlobalConfigTyped = require('../global_config_typed').default;
const { MULTIPART_THRESHOLD_BYTES, MULTIPART_PART_SIZE_BYTES, readChunk } = require('./multipart');
const AMZ_IV = 'x-amz-iv';
const AMZ_KEY = 'x-amz-key';
const AMZ_MATDESC = 'x-amz-matdesc';
const SFC_DIGEST = 'sfc-digest';
const EXPIRED_TOKEN = 'ExpiredToken';
const NOT_FOUND = 'NotFound';
const SNOWFLAKE_S3_DESTINATION = 's3.amazonaws.com';
const ERRORNO_WSAECONNABORTED = 10053; // network connection was aborted
const DATA_SIZE_THRESHOLD = 67108864; // magic number, given from error message.
const resultStatus = require('../file_util').resultStatus;
const HTTP_HEADER_VALUE_OCTET_STREAM = 'application/octet-stream';
// S3 Location: S3 bucket name + path
function S3Location(bucketName, s3path) {
return {
bucketName: bucketName, // S3 bucket name
s3path: s3path, // S3 path name
};
}
/**
* Creates an S3 utility object.
*
* @param connectionConfig
*
* @param s3 - used for tests, mock can be supplied
* @returns {Object}
* @constructor
*/
function S3Util(connectionConfig, s3) {
const AWS = typeof s3 !== 'undefined' ? s3 : require('@aws-sdk/client-s3');
/**
* Create an AWS S3 client using an AWS token.
*/
this.createClient = function (stageInfo, useAccelerateEndpoint) {
const stageCredentials = stageInfo['creds'];
const securityToken = stageCredentials['AWS_TOKEN'];
const isRegionalUrlEnabled = stageInfo.useRegionalUrl || stageInfo.useS3RegionalUrl;
// if GS sends us an endpoint, it's likely for FIPS. Use it.
let endPoint = null;
if (stageInfo['endPoint']) {
endPoint = `https://${stageInfo['endPoint']}`;
}
else {
if (stageInfo.region && isRegionalUrlEnabled) {
const domainSuffixForRegionalUrl = stageInfo.region.toLowerCase().startsWith('cn-')
? 'amazonaws.com.cn'
: 'amazonaws.com';
endPoint = `https://s3.${stageInfo.region}.${domainSuffixForRegionalUrl}`;
}
}
const config = {
apiVersion: '2006-03-01',
region: stageInfo['region'],
credentials: {
accessKeyId: stageCredentials['AWS_KEY_ID'],
secretAccessKey: stageCredentials['AWS_SECRET_KEY'],
sessionToken: securityToken,
},
endpoint: endPoint,
useAccelerateEndpoint: useAccelerateEndpoint,
};
const proxy = ProxyUtil.getProxy(connectionConfig.getProxy(), 'S3 Util');
if (proxy) {
const proxyAgent = getProxyAgent({
proxyOptions: proxy,
connectionConfig,
parsedUrl: new URL(connectionConfig.accessUrl),
destination: endPoint || SNOWFLAKE_S3_DESTINATION,
});
config.requestHandler = new NodeHttpHandler({
httpAgent: proxyAgent,
httpsAgent: proxyAgent,
});
}
return new AWS.S3(config);
};
/**
* Get file header based on file being uploaded or not.
*
* @param {Object} meta
* @param {String} filename
*
* @returns {Object}
*/
this.getFileHeader = async function (meta, filename) {
const stageInfo = meta['stageInfo'];
const client = this.createClient(stageInfo);
try {
const s3location = extractBucketNameAndPath(stageInfo['location']);
const params = {
Bucket: s3location.bucketName,
Key: s3location.s3path + filename,
};
let akey;
try {
await client.headObject(params).then(function (data) {
akey = data;
});
}
catch (err) {
if (err['Code'] === EXPIRED_TOKEN) {
meta['resultStatus'] = resultStatus.RENEW_TOKEN;
return null;
}
else if (err['name'] === NOT_FOUND) {
meta['resultStatus'] = resultStatus.NOT_FOUND_FILE;
return FileHeader(null, null, null);
}
else if (err['Code'] === '400') {
meta['resultStatus'] = resultStatus.RENEW_TOKEN;
return null;
}
else {
meta['resultStatus'] = resultStatus.ERROR;
return null;
}
}
meta['resultStatus'] = resultStatus.UPLOADED;
let encryptionMetadata;
if (akey && akey.Metadata[AMZ_KEY]) {
encryptionMetadata = EncryptionMetadata(akey.Metadata[AMZ_KEY], akey.Metadata[AMZ_IV], akey.Metadata[AMZ_MATDESC]);
}
return FileHeader(akey.Metadata[SFC_DIGEST], akey.ContentLength, encryptionMetadata);
}
finally {
client.destroy();
}
};
/**
* Create the file metadata then upload the file.
*
* @param {String} dataFile
* @param {Object} meta
* @param {Object} encryptionMetadata
*/
this.uploadFile = async function (dataFile, meta, encryptionMetadata) {
if (GlobalConfigTyped.getValue('enableExperimentalMultipartUploads')) {
const fileSize = (await fs.promises.stat(dataFile)).size;
if (fileSize > MULTIPART_THRESHOLD_BYTES) {
await this.uploadFileMultipart(dataFile, fileSize, meta, encryptionMetadata);
}
else {
// TODO: temporarily use readFile instead of createReadStream to work around a Bun bug
// where createReadStream fails the upload in compiled apps. Investigate the root cause
// and fix it properly.
const buffer = await fs.promises.readFile(dataFile);
await this.uploadFileStream(buffer, meta, encryptionMetadata);
}
}
else {
const fileStream = fs.createReadStream(dataFile);
await this.uploadFileStream(fileStream, meta, encryptionMetadata);
}
};
/**
* Multipart upload of `dataFile` to S3, reading `MULTIPART_PART_SIZE_BYTES`
* bytes per chunk into a fresh Buffer and uploading each as a separate
* UploadPart call.
*
* On any error (token expiry, transient network failure, etc.), best-effort
* `AbortMultipartUpload` is issued so we don't leak partial-upload storage.
* The caller's existing retry semantics (RENEW_TOKEN / NEED_RETRY result
* statuses) layer cleanly on top: a failed upload leaves the stage clean
* for the retry to start over.
*
* @param {String} dataFile - Local file path to upload.
* @param {Number} fileSize - Pre-stat'd file size in bytes.
* @param {Object} meta
* @param {Object} encryptionMetadata
*/
this.uploadFileMultipart = async function (dataFile, fileSize, meta, encryptionMetadata) {
const stageInfo = meta['stageInfo'];
const client = this.createClient(stageInfo);
const s3location = extractBucketNameAndPath(stageInfo['location']);
const Bucket = s3location.bucketName;
const Key = s3location.s3path + meta['dstFileName'];
let uploadId;
try {
const createResp = await client.createMultipartUpload({
Bucket,
Key,
Metadata: buildS3Metadata(meta, encryptionMetadata),
});
uploadId = createResp.UploadId;
const fd = await fs.promises.open(dataFile, 'r');
const parts = [];
try {
let position = 0;
let partNumber = 1;
while (position < fileSize) {
const thisPartSize = Math.min(MULTIPART_PART_SIZE_BYTES, fileSize - position);
const buf = await readChunk(fd, position, thisPartSize);
const partResp = await client.uploadPart({
Bucket,
Key,
UploadId: uploadId,
PartNumber: partNumber,
Body: buf,
});
parts.push({ ETag: partResp.ETag, PartNumber: partNumber });
position += thisPartSize;
partNumber++;
}
}
finally {
await fd.close();
}
await client.completeMultipartUpload({
Bucket,
Key,
UploadId: uploadId,
MultipartUpload: { Parts: parts },
});
meta['dstFileSize'] = meta['uploadSize'];
meta['resultStatus'] = resultStatus.UPLOADED;
}
catch (err) {
// If we have an in-flight upload, best-effort abort to release storage.
if (uploadId) {
try {
await client.abortMultipartUpload({ Bucket, Key, UploadId: uploadId });
}
catch (_abortErr) {
// Suppress
}
}
applyS3UploadError(err, meta);
}
finally {
client.destroy();
}
};
/**
* Create the file metadata then upload the file stream.
*
* @param {Buffer|string|stream.Readable} fileStream
* @param {Object} meta
* @param {Object} encryptionMetadata
*/
this.uploadFileStream = async function (fileStream, meta, encryptionMetadata) {
const s3Metadata = buildS3Metadata(meta, encryptionMetadata);
const stageInfo = meta['stageInfo'];
const client = this.createClient(stageInfo);
try {
const s3location = extractBucketNameAndPath(meta['stageInfo']['location']);
const params = {
Bucket: s3location.bucketName,
Body: fileStream,
Key: s3location.s3path + meta['dstFileName'],
Metadata: s3Metadata,
};
// call S3 to upload file to specified bucket
try {
await client.putObject(params);
}
catch (err) {
applyS3UploadError(err, meta);
return;
}
meta['dstFileSize'] = meta['uploadSize'];
meta['resultStatus'] = resultStatus.UPLOADED;
}
finally {
client.destroy();
}
};
/**
* Download the file.
*
* @param {String} dataFile
* @param {Object} meta
* @param {Object} encryptionMetadata
*/
this.nativeDownloadFile = async function (meta, fullDstPath) {
const stageInfo = meta['stageInfo'];
const client = this.createClient(stageInfo);
try {
const s3location = extractBucketNameAndPath(meta['stageInfo']['location']);
const params = {
Bucket: s3location.bucketName,
Key: s3location.s3path + meta['dstFileName'],
};
// call S3 to download file to specified bucket
try {
Logger().debug(`Send Get Request to the Bucket: ${params.Bucket}, GET request: ${params.Key}`);
await client
.getObject(params)
.then((data) => {
Logger().debug(`Http Status for the GET request: ${params.Key} : ${data.$metadata.httpStatusCode}`);
return data.Body.transformToByteArray();
})
.then((data) => {
return new Promise((resolve, reject) => {
fs.writeFile(fullDstPath, data, 'binary', (err) => {
if (err) {
reject(err);
}
resolve();
});
});
});
}
catch (err) {
if (err['Code'] === EXPIRED_TOKEN) {
meta['resultStatus'] = resultStatus.RENEW_TOKEN;
}
else {
meta['lastError'] = err;
if (err['Code'] === ERRORNO_WSAECONNABORTED.toString()) {
meta['resultStatus'] = resultStatus.NEED_RETRY_WITH_LOWER_CONCURRENCY;
}
else {
meta['resultStatus'] = resultStatus.NEED_RETRY;
}
}
return;
}
meta['resultStatus'] = resultStatus.DOWNLOADED;
}
finally {
client.destroy();
}
};
}
/**
* Extract the bucket name and path from the metadata's stage location.
*
* @param {String} stageLocation
*
* @returns {Object}
*/
function extractBucketNameAndPath(stageLocation) {
// expand '~' and '~user' expressions
if (process.platform !== 'win32') {
stageLocation = expandTilde(stageLocation);
}
let bucketName = stageLocation;
let s3path;
// split stage location as bucket name and path
if (stageLocation.includes('/')) {
bucketName = stageLocation.substring(0, stageLocation.indexOf('/'));
s3path = stageLocation.substring(stageLocation.indexOf('/') + 1, stageLocation.length);
if (s3path && !s3path.endsWith('/')) {
s3path += '/';
}
}
return S3Location(bucketName, s3path);
}
/**
* Build the S3 object metadata, including the encryption envelope when the
* upload is encrypted, so every upload path persists the same metadata shape.
*/
function buildS3Metadata(meta, encryptionMetadata) {
const s3Metadata = {
HTTP_HEADER_CONTENT_TYPE: HTTP_HEADER_VALUE_OCTET_STREAM,
SFC_DIGEST: meta['SHA256_DIGEST'],
};
if (encryptionMetadata) {
s3Metadata[AMZ_IV] = encryptionMetadata.iv;
s3Metadata[AMZ_KEY] = encryptionMetadata.key;
s3Metadata[AMZ_MATDESC] = encryptionMetadata.matDesc;
}
return s3Metadata;
}
/**
* Classify an upload failure into the file-transfer-agent's `meta.resultStatus`
* convention so retry semantics stay consistent across every catch site. The
* error may be an S3 SDK error (mapped to RENEW_TOKEN /
* NEED_RETRY_WITH_LOWER_CONCURRENCY) or any other JS error (e.g. a short read),
* which falls back to NEED_RETRY.
*/
function applyS3UploadError(err, meta) {
if (err && err['Code'] === EXPIRED_TOKEN) {
meta['resultStatus'] = resultStatus.RENEW_TOKEN;
}
else {
meta['lastError'] = err;
if (err && err['Code'] === ERRORNO_WSAECONNABORTED.toString()) {
meta['resultStatus'] = resultStatus.NEED_RETRY_WITH_LOWER_CONCURRENCY;
}
else {
meta['resultStatus'] = resultStatus.NEED_RETRY;
}
}
}
module.exports = {
S3Util,
SNOWFLAKE_S3_DESTINATION,
DATA_SIZE_THRESHOLD,
extractBucketNameAndPath,
};
//# sourceMappingURL=s3_util.js.map