diff --git a/index.d.ts b/index.d.ts index b598633b4..4c2dc3ce3 100644 --- a/index.d.ts +++ b/index.d.ts @@ -4,6 +4,8 @@ import { WIP_ConnectionConfig } from './lib/connection/types'; * The snowflake-sdk module provides an instance to connect to the Snowflake server * @see [source] {@link https://docs.snowflake.com/en/developer-guide/node-js/nodejs-driver} */ +import { HttpHeadersCustomizer } from './lib/connection/types'; + declare module 'snowflake-sdk' { export const ErrorCode: typeof import('./lib/error_code').default; @@ -343,6 +345,12 @@ declare module 'snowflake-sdk' { * The option to pass passcode from DUO. */ passcode?: string; + + /** + * Customizes the HTTP headers sent with each request. + * The customizer functions are called with the HTTP method and URL. + */ + httpHeadersCustomizer?: Array; } export interface Connection { diff --git a/lib/connection/connection_config.js b/lib/connection/connection_config.js index e2454ed74..84756f2c3 100644 --- a/lib/connection/connection_config.js +++ b/lib/connection/connection_config.js @@ -2,6 +2,7 @@ const os = require('os'); const url = require('url'); const Util = require('../util'); const ProxyUtil = require('../proxy_util'); +const { isValidHTTPHeadersCustomizer } = require('../http/request_util'); const Errors = require('../errors'); const ConnectionConstants = require('../constants/connection_constants'); const path = require('path'); @@ -74,6 +75,7 @@ const DEFAULT_PARAMS = 'oauthScope', 'oauthChallengeMethod', 'oauthHttpAllowed', //only for tests + 'httpHeadersCustomizer', 'workloadIdentityProvider', 'workloadIdentityAzureEntraIdResource' ]; @@ -591,6 +593,14 @@ function ConnectionConfig(options, validateCredentials, qaMode, clientInfo) { passcode = options.passcode; } + let httpHeadersCustomizer = []; + if (Util.exists(options.httpHeadersCustomizer)) { + Errors.checkArgumentValid(isValidHTTPHeadersCustomizer(options.httpHeadersCustomizer), + ErrorCodes.ERR_CONN_CREATE_INVALID_HTTP_HEADER_CUSTOMIZERS); + + httpHeadersCustomizer = options.httpHeadersCustomizer; + } + if (validateDefaultParameters) { for (const [key] of Object.entries(options)) { if (!DEFAULT_PARAMS.includes(key)) { @@ -964,6 +974,10 @@ function ConnectionConfig(options, validateCredentials, qaMode, clientInfo) { return oauthHttpAllowed || false; }; + this.getHttpHeadersCustomizer = function () { + return httpHeadersCustomizer; + }; + /** * Returns attributes of Connection Config object that can be used to identify * the connection, when ID is not available in the scope. This is not sufficient set, diff --git a/lib/connection/statement.js b/lib/connection/statement.js index a36854eef..42fb61a18 100644 --- a/lib/connection/statement.js +++ b/lib/connection/statement.js @@ -1473,15 +1473,17 @@ function sendSfRequest(statementContext, options, appendQueryParamOnRetry) { const sendRequest = function () { // if this is a retry and a query parameter should be appended to the url on // retry, update the url - if ((numRetries > 0) && appendQueryParamOnRetry) { - const retryOption = { - url: urlOrig, - retryCount: numRetries, - retryReason: lastStatusCodeForRetry, - includeRetryReason: connectionConfig.getIncludeRetryReason(), - }; - - options.url = Util.url.appendRetryParam(retryOption); + if (numRetries > 0) { + options.isRetry = true; + if (appendQueryParamOnRetry) { + const retryOption = { + url: urlOrig, + retryCount: numRetries, + retryReason: lastStatusCodeForRetry, + includeRetryReason: connectionConfig.getIncludeRetryReason(), + }; + options.url = Util.url.appendRetryParam(retryOption); + } } sf.request(options); @@ -1497,6 +1499,7 @@ function sendSfRequest(statementContext, options, appendQueryParamOnRetry) { )) { // increment the retry count numRetries++; + options.isRetry = true; lastStatusCodeForRetry = err.response ? err.response.statusCode : 0; // use exponential backoff with decorrelated jitter to compute the diff --git a/lib/connection/types.ts b/lib/connection/types.ts index c4b2a7aa2..8d82eac62 100644 --- a/lib/connection/types.ts +++ b/lib/connection/types.ts @@ -33,3 +33,8 @@ export interface WIP_ConnectionConfig { */ workloadIdentityAzureEntraIdResource?: string; } + +export interface HttpHeadersCustomizer { + applies(method: string, url: string): boolean; + newHeaders() : Record; +} \ No newline at end of file diff --git a/lib/constants/error_messages.js b/lib/constants/error_messages.js index cad7970f6..b817fdb6b 100644 --- a/lib/constants/error_messages.js +++ b/lib/constants/error_messages.js @@ -86,6 +86,7 @@ exports[404059] = 'Invalid oauth client id. The specified value must not be an e exports[404060] = 'Invalid oauth client secret. The specified value must not be an empty string'; exports[404061] = 'Invalid oauth token request URL. The specified value must be a valid URL starting with the https or http protocol.'; exports[404062] = 'No workload identity credentials were found. Provider: %s'; +exports[404063] = 'Invalid Http headers customizer. The specified value must contain the following functions: \'applices\', \'newHeaders\' and \'invokeOnce\''; // 405001 exports[405001] = 'Invalid callback. The specified value must be a function.'; diff --git a/lib/error_code.ts b/lib/error_code.ts index c41da0b08..42650f79e 100644 --- a/lib/error_code.ts +++ b/lib/error_code.ts @@ -85,6 +85,7 @@ enum ErrorCode { ERR_CONN_CREATE_INVALID_OUATH_CLIENT_SECRET = 404060, ERR_CONN_CREATE_INVALID_OUATH_TOKEN_REQUEST_URL = 404061, ERR_CONN_CREATE_MISSING_WORKLOAD_IDENTITY_CREDENTIALS = 404062, + ERR_CONN_CREATE_INVALID_HTTP_HEADER_CUSTOMIZERS = 404063, // 405001 ERR_CONN_CONNECT_INVALID_CALLBACK = 405001, diff --git a/lib/file_transfer_agent/gcs_util.js b/lib/file_transfer_agent/gcs_util.js index d266b687a..82ba0eb60 100644 --- a/lib/file_transfer_agent/gcs_util.js +++ b/lib/file_transfer_agent/gcs_util.js @@ -4,6 +4,7 @@ const FileHeader = require('../file_util').FileHeader; const getProxyAgent = require('../http/node').getProxyAgent; const ProxyUtil = require('../proxy_util'); const Util = require('../util'); +const RequestUtil = require('../http/request_util'); const { shouldPerformGCPBucket, lstrip } = require('../util'); const GCS_METADATA_PREFIX = 'x-goog-meta-'; @@ -144,7 +145,9 @@ function GCSUtil(connectionConfig, httpClient) { let matDescKey; try { - if (shouldPerformGCPBucket(accessToken) && !isProxyEnabled) { + const customHeaders = RequestUtil.getCustomHeaders(connectionConfig, 'HEAD', url); + RequestUtil.cleanOverridingHeaders(gcsHeaders, customHeaders); + if (shouldPerformGCPBucket(accessToken) && !isProxyEnabled && Object.keys(customHeaders).length === 0) { const gcsLocation = this.extractBucketNameAndPath(meta['stageInfo']['location']); const metadata = await meta['client'].gcsClient .bucket(gcsLocation.bucketName) @@ -156,7 +159,7 @@ function GCSUtil(connectionConfig, httpClient) { encryptionDataProp = metadata[0].metadata[ENCRYPTIONDATAPROP]; matDescKey = metadata[0].metadata[MATDESC_KEY]; } else { - const response = await axios.head(url, { headers: gcsHeaders }); + const response = await axios.head(url, { headers: meta.isRetry ? gcsHeaders : { ...gcsHeaders, ...customHeaders } }); digest = response.headers[GCS_METADATA_SFC_DIGEST]; contentLength = response.headers['content-length']; @@ -188,6 +191,7 @@ function GCSUtil(connectionConfig, httpClient) { if ([403, 408, 429, 500, 503].includes(errCode)) { meta['lastError'] = err; meta['resultStatus'] = resultStatus.NEED_RETRY; + meta.isRetry = true; return; } if (errCode === 404) { @@ -283,7 +287,10 @@ function GCSUtil(connectionConfig, httpClient) { } try { - if (shouldPerformGCPBucket(accessToken) && !isProxyEnabled) { + const customHeaders = RequestUtil.getCustomHeaders(connectionConfig, 'PUT', uploadUrl); + RequestUtil.cleanOverridingHeaders(gcsHeaders, customHeaders); + + if (shouldPerformGCPBucket(accessToken) && !isProxyEnabled && Object.keys(customHeaders).length === 0) { const gcsLocation = this.extractBucketNameAndPath(meta['stageInfo']['location']); await meta['client'].gcsClient @@ -295,18 +302,19 @@ function GCSUtil(connectionConfig, httpClient) { metadata: { [ENCRYPTIONDATAPROP]: gcsHeaders[GCS_METADATA_ENCRYPTIONDATAPROP], [MATDESC_KEY]: gcsHeaders[GCS_METADATA_MATDESC_KEY], - [SFC_DIGEST]: gcsHeaders[GCS_METADATA_SFC_DIGEST] + [SFC_DIGEST]: gcsHeaders[GCS_METADATA_SFC_DIGEST], } } }); } else { // Set maxBodyLength to allow large file uploading - await axios.put(uploadUrl, fileStream, { maxBodyLength: Infinity, headers: gcsHeaders }); + await axios.put(uploadUrl, fileStream, { maxBodyLength: Infinity, headers: meta.isRetry ? gcsHeaders : { ...gcsHeaders, ...customHeaders } }); } } catch (err) { if ([403, 408, 429, 500, 503].includes(err['code'])) { meta['lastError'] = err; - meta['resultStatus'] = resultStatus.NEED_RETRY; + meta['resultStatus'] = resultStatus.NEED_RETRY, + meta.isRetry = true; } else if (!accessToken && err['code'] === 400 && (!meta['lastError'] || meta['lastError']['code'] !== 400)) { // Only attempt to renew urls if this isn't the second time this happens @@ -356,7 +364,10 @@ function GCSUtil(connectionConfig, httpClient) { let size; try { - if (shouldPerformGCPBucket(accessToken) && !isProxyEnabled) { + const customHeaders = RequestUtil.getCustomHeaders(connectionConfig, 'GET', downloadUrl); + RequestUtil.cleanOverridingHeaders(gcsHeaders, customHeaders); + + if (shouldPerformGCPBucket(accessToken) && !isProxyEnabled && Object.keys(customHeaders).length === 0) { const gcsLocation = this.extractBucketNameAndPath(meta['stageInfo']['location']); await meta['client'].gcsClient .bucket(gcsLocation.bucketName) @@ -377,7 +388,7 @@ function GCSUtil(connectionConfig, httpClient) { } else { let response; await axios.get(downloadUrl, { - headers: gcsHeaders, + headers: meta.isRetry ? gcsHeaders : { ...gcsHeaders, ...customHeaders }, responseType: 'stream' }).then(async (res) => { response = res; @@ -406,8 +417,10 @@ function GCSUtil(connectionConfig, httpClient) { meta['lastError'] = err; if (err['code'] === ERRORNO_WSAECONNABORTED) { meta['resultStatus'] = resultStatus.NEED_RETRY_WITH_LOWER_CONCURRENCY; + meta.isRetry = true; } else { meta['resultStatus'] = resultStatus.NEED_RETRY; + meta.isRetry = true; } } return; diff --git a/lib/file_transfer_agent/remote_storage_util.js b/lib/file_transfer_agent/remote_storage_util.js index da62ad8bf..4659760ef 100644 --- a/lib/file_transfer_agent/remote_storage_util.js +++ b/lib/file_transfer_agent/remote_storage_util.js @@ -159,7 +159,8 @@ function RemoteStorageUtil(connectionConfig) { } else { dataFile = meta['realSrcFilePath']; } - + meta.isRetry = false; + const utilClass = this.getForStorageType(meta['stageInfo']['locationType']); let maxConcurrency = meta['parallel']; @@ -198,6 +199,7 @@ function RemoteStorageUtil(connectionConfig) { const sleepingTime = Math.min(Math.pow(2, retry), 16); await new Promise(resolve => setTimeout(resolve, sleepingTime)); } + meta.isRetry = true; } else if (meta['resultStatus'] === resultStatus.NEED_RETRY_WITH_LOWER_CONCURRENCY) { lastErr = meta['lastError']; // Failed to upload file, retrying with max concurrency @@ -209,6 +211,7 @@ function RemoteStorageUtil(connectionConfig) { const sleepingTime = Math.min(Math.pow(2, retry), 16); await new Promise(resolve => setTimeout(resolve, sleepingTime)); } + meta.isRetry = true; } } if (lastErr) { @@ -377,6 +380,7 @@ function RemoteStorageUtil(connectionConfig) { const sleepingTime = Math.min(Math.pow(2, retry), 16); await new Promise(resolve => setTimeout(resolve, sleepingTime)); } + meta.isRetry = true; } else if (meta['resultStatus'] === resultStatus.NEED_RETRY) { lastErr = meta['lastError']; // Failed to download file, retrying @@ -385,6 +389,7 @@ function RemoteStorageUtil(connectionConfig) { await new Promise(resolve => setTimeout(resolve, sleepingTime)); } } + meta.isRetry = true; } if (lastErr) { throw new Error(lastErr); diff --git a/lib/file_transfer_agent/s3_util.js b/lib/file_transfer_agent/s3_util.js index 1cad36b86..2f380934d 100644 --- a/lib/file_transfer_agent/s3_util.js +++ b/lib/file_transfer_agent/s3_util.js @@ -4,6 +4,7 @@ const FileHeader = require('../file_util').FileHeader; const expandTilde = require('expand-tilde'); const getProxyAgent = require('../http/node').getProxyAgent; const ProxyUtil = require('../proxy_util'); +const RequestUtil = require('../http/request_util'); const AMZ_IV = 'x-amz-iv'; const AMZ_KEY = 'x-amz-key'; @@ -103,6 +104,9 @@ function S3Util(connectionConfig, s3, filestream) { Key: s3location.s3path + filename }; + if (!meta.isRetry) { + this.addCustomHeaders(params, 'GET', client?.config?.endPoint || SNOWFLAKE_S3_DESTINATION); + } let akey; try { @@ -187,6 +191,10 @@ function S3Util(connectionConfig, s3, filestream) { Metadata: s3Metadata }; + if (meta.isRetry) { + this.addCustomHeaders(params, 'PUT', client?.config?.endPoint || SNOWFLAKE_S3_DESTINATION); + } + // call S3 to upload file to specified bucket try { await client.putObject(params); @@ -197,8 +205,10 @@ function S3Util(connectionConfig, s3, filestream) { meta['lastError'] = err; if (err['Code'] === ERRORNO_WSAECONNABORTED.toString()) { meta['resultStatus'] = resultStatus.NEED_RETRY_WITH_LOWER_CONCURRENCY; + meta.isRetry = true; } else { meta['resultStatus'] = resultStatus.NEED_RETRY; + meta.isRetry = true; } } return; @@ -226,6 +236,9 @@ function S3Util(connectionConfig, s3, filestream) { Key: s3location.s3path + meta['dstFileName'], }; + if (meta.isRetry) { + this.addCustomHeaders(params, 'GET', client?.config?.endPoint || SNOWFLAKE_S3_DESTINATION); + } // call S3 to download file to specified bucket try { await client.getObject(params) @@ -247,14 +260,23 @@ function S3Util(connectionConfig, s3, filestream) { meta['lastError'] = err; if (err['Code'] === ERRORNO_WSAECONNABORTED.toString()) { meta['resultStatus'] = resultStatus.NEED_RETRY_WITH_LOWER_CONCURRENCY; + meta.isRetry = true; } else { meta['resultStatus'] = resultStatus.NEED_RETRY; + meta.isRetry = true; } } return; } meta['resultStatus'] = resultStatus.DOWNLOADED; }; + + this.addCustomHeaders = function (params, method, endPoint) { + const customHeaders = RequestUtil.getCustomHeaders(connectionConfig, method, endPoint); + if (Object.keys(customHeaders).length > 0) { + params['Metadata'] ? params['Metadata'] = { ...params['Metadata'], ...customHeaders } : params['Metadata'] = customHeaders; + } + }; } /** @@ -285,4 +307,6 @@ function extractBucketNameAndPath(stageLocation) { return S3Location(bucketName, s3path); } + + module.exports = { S3Util, SNOWFLAKE_S3_DESTINATION, DATA_SIZE_THRESHOLD, extractBucketNameAndPath }; diff --git a/lib/http/base.js b/lib/http/base.js index c4158b5b4..2b1b997a5 100644 --- a/lib/http/base.js +++ b/lib/http/base.js @@ -216,7 +216,6 @@ HttpClient.prototype.getAgent = function () { return null; }; -module.exports = HttpClient; function sanitizeAxiosResponse(response) { Logger.getInstance().trace('Request%s - sanitizing response data.', requestUtil.describeRequestFromResponse(response)); @@ -238,7 +237,17 @@ function sanitizeAxiosError(error) { function prepareRequestOptions(options, requestHandlers = {}) { Logger.getInstance().trace('Request%s - constructing options.', requestUtil.describeRequestFromOptions(options)); - const headers = normalizeHeaders(options.headers) || {}; + let headers = normalizeHeaders(options.headers) || {}; + if (!options.isRetry) { + //Save the original headers. + requestUtil.setOriginalHeader(headers); + const customHeaders = requestUtil.getCustomHeaders(this._connectionConfig, options.method, options.url); + requestUtil.cleanOverridingHeaders(headers, customHeaders); + Object.assign(headers, customHeaders); + } else { + Logger.getInstance().debug(`Customizer should only run on the first attempt and this is a ${options.retry} retry. Skipping.`); + headers = requestUtil.getOriginalHeader(); + } const timeout = options.timeout || this._connectionConfig.getTimeout() || @@ -261,7 +270,6 @@ function prepareRequestOptions(options, requestHandlers = {}) { } }); } - const params = options.params; let mock; @@ -371,3 +379,12 @@ function normalizeResponse(response) { return response; } + +//Testing purposes only +function getHttpRequestHeaders(connectionConfig, options) { + const httpClient = new HttpClient(connectionConfig); + httpClient.constructExponentialBackoffStrategy = () => 0; + return prepareRequestOptions.call(httpClient, options).headers; +} + +module.exports = { HttpClient, getHttpRequestHeaders }; diff --git a/lib/http/browser.js b/lib/http/browser.js index 6c0c865de..3b1a80450 100644 --- a/lib/http/browser.js +++ b/lib/http/browser.js @@ -1,6 +1,6 @@ const Util = require('../util'); const request = require('browser-request'); -const Base = require('./base'); +const { HttpClient } = require('./base'); const Logger = require('../logger'); /** @@ -12,10 +12,10 @@ const Logger = require('../logger'); function BrowserHttpClient(connectionConfig) { Logger.getInstance().trace('Initializing BrowserHttpClient with Connection Config[%s]', connectionConfig.describeIdentityAttributes()); - Base.apply(this, [connectionConfig]); + HttpClient.apply(this, [connectionConfig]); } -Util.inherits(BrowserHttpClient, Base); +Util.inherits(BrowserHttpClient, HttpClient); /** * @inheritDoc diff --git a/lib/http/node.js b/lib/http/node.js index 5a1531252..f3dc23677 100644 --- a/lib/http/node.js +++ b/lib/http/node.js @@ -1,6 +1,6 @@ const Util = require('../util'); const ProxyUtil = require('../proxy_util'); -const Base = require('./base'); +const { HttpClient } = require('./base'); const HttpsAgent = require('../agent/https_ocsp_agent'); const HttpsProxyAgent = require('../agent/https_proxy_agent'); const HttpAgent = require('http').Agent; @@ -37,10 +37,10 @@ NodeHttpClient.prototype.constructExponentialBackoffStrategy = function () { function NodeHttpClient(connectionConfig) { Logger.getInstance().trace('Initializing NodeHttpClient with Connection Config[%s]', connectionConfig.describeIdentityAttributes()); - Base.apply(this, [connectionConfig]); + HttpClient.apply(this, [connectionConfig]); } -Util.inherits(NodeHttpClient, Base); +Util.inherits(NodeHttpClient, HttpClient); const httpsAgentCache = new Map(); diff --git a/lib/http/request_util.js b/lib/http/request_util.js index 38f69222e..f738e0834 100644 --- a/lib/http/request_util.js +++ b/lib/http/request_util.js @@ -1,3 +1,4 @@ +const Logger = require('../logger'); const LoggingUtil = require('../logger/logging_util'); const sfParams = require('../constants/sf_params'); @@ -194,6 +195,68 @@ function describeURL( ); } +let originalHeader = {}; + +exports.getOriginalHeader = function () { + return originalHeader; +}; + +exports.setOriginalHeader = function (header) { + originalHeader = { ... header }; +}; + +exports.cleanOverridingHeaders = function (originalHeader, newHeader) { + const originalHeaderKeys = Object.keys(originalHeader); + const newHeaderKeys = Object.keys(newHeader); + + if (originalHeaderKeys.length === 0 || newHeaderKeys.length === 0) { + return; + } + + for (const headerKey of newHeaderKeys) { + if (originalHeaderKeys.find((key) => key.toLowerCase() === headerKey.toLowerCase())) { + Logger.getInstance().debug(`Customizer attempted to override existing driver header ${originalHeaderKeys} which is not allowed. Skipping.`); + delete newHeader[headerKey]; + } + } +}; + +exports.isValidHTTPHeadersCustomizer = function (customizers) { + const requireMethods = ['applies', 'newHeaders']; + for (const customizer of customizers) { + for (const method of requireMethods) { + if ( + typeof customizer !== 'object' || + customizer === null || + typeof customizer[method] !== 'function' + ) { + return false; + } + } + } + return true; +}; + +exports.getCustomHeaders = function (connectionConfig, method, endPoint) { + const customHeaders = {}; + + if (connectionConfig.getHttpHeadersCustomizer().length > 0) { + for (const customizer of connectionConfig.getHttpHeadersCustomizer()) { + if (customizer.applies(method, endPoint)) { + const headers = customizer.newHeaders(); + for (const [key, value] of Object.entries(headers)) { + if (typeof value === 'object' && value !== null) { + customHeaders[key] = JSON.stringify(value); + } else { + customHeaders[key] = value; + } + } + } + } + } + return customHeaders; +}; + exports.DEFAULT_ATTRIBUTES_DESCRIBING_REQUEST_WITH_VALUES = DEFAULT_ATTRIBUTES_DESCRIBING_REQUEST_WITH_VALUES; exports.DEFAULT_ATTRIBUTES_DESCRIBING_REQUEST_WITHOUT_VALUES = DEFAULT_ATTRIBUTES_DESCRIBING_REQUEST_WITHOUT_VALUES; diff --git a/lib/services/sf.js b/lib/services/sf.js index 91fd2cd3d..014c8ae32 100644 --- a/lib/services/sf.js +++ b/lib/services/sf.js @@ -598,6 +598,7 @@ function StateAbstract(options) { url: requestOptions.absoluteUrl, gzip: requestOptions.gzip, json: requestOptions.json, + isRetry: requestOptions.isRetry || false, params: params, callback: async function (err, response, body) { // if we got an error, wrap it into a network error @@ -665,11 +666,9 @@ function StateAbstract(options) { body.code === GSErrors.code.OAUTH_TOKEN_EXPIRED && data.authnMethod === 'OAUTH' ) { Logger.getInstance().debug('ID Token being used has expired. Reauthenticating'); - await auth.reauthenticate(requestOptions.json); return httpClient.request(realRequestOptions); } - err = Errors.createOperationFailedError( body.code, data, body.message, data && data.sqlState ? data.sqlState : undefined); @@ -689,6 +688,7 @@ function StateAbstract(options) { }; if (requestOptions.retry > 2) { + requestOptions.isRetry = true; const includesParam = requestOptions.url.includes('?'); realRequestOptions.url += (includesParam ? '&' : '?'); realRequestOptions.url += @@ -1256,6 +1256,7 @@ StateConnecting.prototype.continue = function () { scope: this, startTime: startTime, retry: numRetries, + isRetry: false, callback: requestCallback }); request.send(); diff --git a/lib/util.ts b/lib/util.ts index dd2713298..8f413dbaf 100644 --- a/lib/util.ts +++ b/lib/util.ts @@ -16,6 +16,11 @@ export { driverName, driverVersion }; export const userAgent = `JavaScript/${driverVersion} (${process.platform}-${process.arch}) NodeJS/${nodeJSVersion}`; +export interface HttpHeadersCustomizer { + applies: (url: string) => boolean; + newHeaders: () => Record; +} + /** * Note: A simple wrapper around util.inherits() for now, but this might change * in the future. @@ -677,4 +682,4 @@ export function escapeHTML(value: string) { */ export async function dynamicImportESMInTypescriptWithCommonJS(moduleName: string) { return Function(`return import("${moduleName}")`)() -} +} \ No newline at end of file diff --git a/package.json b/package.json index 198e3ffbf..cb0b06c2a 100644 --- a/package.json +++ b/package.json @@ -79,12 +79,12 @@ "lint:check:all": "eslint lib samples system_test test && check-dts index.d.ts", "lint:check:all:errorsOnly": "npm run lint:check:all -- --quiet", "lint:fix": "eslint --fix", - "test": "mocha 'test/unit/**/*.{js,ts}'", + "test": "mocha \"test/unit/**/*.{js,ts}\"", "test:authentication": "mocha 'test/authentication/**/*.{js,ts}'", "test:integration": "mocha 'test/integration/**/*.{js,ts}'", "test:single": "mocha", "test:system": "mocha 'system_test/**/*.{js,ts}'", - "test:unit": "mocha 'test/unit/**/*.{js,ts}'", + "test:unit": "mocha \"test/unit/**/*.{js,ts}'", "test:unit:coverage": "nyc npm run test:unit", "test:ci": "mocha 'test/{unit,integration}/**/*.{js,ts}'", "test:ci:coverage": "nyc npm run test:ci", diff --git a/test/unit/connection/connection_config_test.js b/test/unit/connection/connection_config_test.js index f5f218a5f..24b9e6a36 100644 --- a/test/unit/connection/connection_config_test.js +++ b/test/unit/connection/connection_config_test.js @@ -837,6 +837,23 @@ describe('ConnectionConfig: basic', function () { }, errorCode: ErrorCodes.ERR_CONN_CREATE_INVALID_OUATH_AUTHORIZATION_URL }, + { + name: 'invalid config - incorrect httpHeadersCustomizer', + + options: { + account: 'account', + username: 'username', + password: 'password', + httpHeadersCustomizer: [ + { + apply: function () { + return 'invalid'; + } + } + ] + }, + errorCode: ErrorCodes.ERR_CONN_CREATE_INVALID_HTTP_HEADER_CUSTOMIZERS + }, ]; const createNegativeITCallback = function (testCase) { @@ -1720,7 +1737,7 @@ describe('ConnectionConfig: basic', function () { oauthAuthorizationUrl: (config) => config.getOauthAuthorizationUrl() === 'http://host.snowflakecomputing.cn:8082/oauth/authorize', oauthTokenRequestUrl: (config) => config.getOauthTokenRequestUrl() === 'http://host.snowflakecomputing.cn:8082/oauth/token-request', } - } + }, ]; const createItCallback = function (testCase) { diff --git a/test/unit/file_transfer_agent/gcs_test.js b/test/unit/file_transfer_agent/gcs_test.js index 3c810ea60..44c9801f3 100644 --- a/test/unit/file_transfer_agent/gcs_test.js +++ b/test/unit/file_transfer_agent/gcs_test.js @@ -20,6 +20,7 @@ describe('GCS client', function () { getProxy: function () { return this.proxy; }, + getHttpHeadersCustomizer: () => [], accessUrl: 'http://fakeaccount.snowflakecomputing.com', }; @@ -352,4 +353,16 @@ describe('GCS client', function () { await GCS.uploadFile(dataFile, meta, encryptionMetadata); assert.strictEqual(meta['resultStatus'], resultStatus.RENEW_TOKEN); }); + + it('upload - isRetry flags', async function () { + httpClient.put = async () => { + const err = new Error(); + err.code = 403; + throw err; + }; + const GCS = new SnowflakeGCSUtil(connectionConfig, httpClient); + + await GCS.uploadFile(dataFile, meta, encryptionMetadata); + assert.strictEqual(meta.isRetry, true); + }); }); diff --git a/test/unit/file_transfer_agent/s3_test.js b/test/unit/file_transfer_agent/s3_test.js index 8dd88cb2a..b8076e46f 100644 --- a/test/unit/file_transfer_agent/s3_test.js +++ b/test/unit/file_transfer_agent/s3_test.js @@ -14,7 +14,8 @@ describe('S3 client', function () { const mockKey = 'mockKey'; const mockIv = 'mockIv'; const mockMatDesc = 'mockMatDesc'; - const noProxyConnectionConfig = { + const ConnectionConfig = { + getHttpHeadersCustomizer: () => [], getProxy: function () { return null; } @@ -77,7 +78,7 @@ describe('S3 client', function () { s3 = require('s3'); filesystem = require('filesystem'); - AWS = new SnowflakeS3Util(noProxyConnectionConfig, s3, filesystem); + AWS = new SnowflakeS3Util(ConnectionConfig, s3, filesystem); }); describe('AWS client endpoint testing', async function () { @@ -181,7 +182,7 @@ describe('S3 client', function () { } }); s3 = require('s3'); - const AWS = new SnowflakeS3Util(noProxyConnectionConfig, s3); + const AWS = new SnowflakeS3Util(ConnectionConfig, s3); await AWS.getFileHeader(meta, dataFile); assert.strictEqual(meta['resultStatus'], resultStatus.RENEW_TOKEN); }); @@ -208,7 +209,7 @@ describe('S3 client', function () { }); s3 = require('s3'); - const AWS = new SnowflakeS3Util(noProxyConnectionConfig, s3); + const AWS = new SnowflakeS3Util(ConnectionConfig, s3); await AWS.getFileHeader(meta, dataFile); assert.strictEqual(meta['resultStatus'], resultStatus.NOT_FOUND_FILE); }); @@ -234,7 +235,7 @@ describe('S3 client', function () { } }); s3 = require('s3'); - const AWS = new SnowflakeS3Util(noProxyConnectionConfig, s3, filesystem); + const AWS = new SnowflakeS3Util(ConnectionConfig, s3, filesystem); await AWS.getFileHeader(meta, dataFile); assert.strictEqual(meta['resultStatus'], resultStatus.RENEW_TOKEN); }); @@ -260,7 +261,7 @@ describe('S3 client', function () { } }); s3 = require('s3'); - const AWS = new SnowflakeS3Util(noProxyConnectionConfig, s3, filesystem); + const AWS = new SnowflakeS3Util(ConnectionConfig, s3, filesystem); await AWS.getFileHeader(meta, dataFile); assert.strictEqual(meta['resultStatus'], resultStatus.ERROR); }); @@ -297,7 +298,7 @@ describe('S3 client', function () { }); s3 = require('s3'); filesystem = require('filesystem'); - const AWS = new SnowflakeS3Util(noProxyConnectionConfig, s3, filesystem); + const AWS = new SnowflakeS3Util(ConnectionConfig, s3, filesystem); await AWS.uploadFile(dataFile, meta, encryptionMetadata); assert.strictEqual(meta['resultStatus'], resultStatus.RENEW_TOKEN); }); @@ -329,9 +330,10 @@ describe('S3 client', function () { }); s3 = require('s3'); filesystem = require('filesystem'); - const AWS = new SnowflakeS3Util(noProxyConnectionConfig, s3, filesystem); + const AWS = new SnowflakeS3Util(ConnectionConfig, s3, filesystem); await AWS.uploadFile(dataFile, meta, encryptionMetadata); assert.strictEqual(meta['resultStatus'], resultStatus.NEED_RETRY_WITH_LOWER_CONCURRENCY); + assert.strictEqual(meta.isRetry, true); }); it('upload - fail HTTP 400', async function () { @@ -361,9 +363,10 @@ describe('S3 client', function () { }); s3 = require('s3'); filesystem = require('filesystem'); - const AWS = new SnowflakeS3Util(noProxyConnectionConfig, s3, filesystem); + const AWS = new SnowflakeS3Util(ConnectionConfig, s3, filesystem); await AWS.uploadFile(dataFile, meta, encryptionMetadata); assert.strictEqual(meta['resultStatus'], resultStatus.NEED_RETRY); + assert.strictEqual(meta.isRetry, true); }); it('proxy configured', async function () { @@ -386,6 +389,7 @@ describe('S3 client', function () { protocol: 'https' }; const proxyConnectionConfig = { + getHttpHeadersCustomizer: () => [], accessUrl: 'http://snowflake.com', getProxy: function () { return proxyOptions; diff --git a/test/unit/http_headers_customizer_test.js b/test/unit/http_headers_customizer_test.js new file mode 100644 index 000000000..ef426d359 --- /dev/null +++ b/test/unit/http_headers_customizer_test.js @@ -0,0 +1,103 @@ +// import {getHttpRequestHeaders} from '../../lib/http/base'; +// import ConnectionConfig from '../../lib/connection/connection_config'; +// import assert from 'assert'; +// import * as Snowflake from '../../index'; + +const { getHttpRequestHeaders } = require('../../lib/http/base'); +const ConnectionConfig = require('../../lib/connection/connection_config'); +const assert = require('assert'); + +describe('customizer header tests', () => { + const customHeaders = [{ + applies: () => true, + newHeaders: function () { + return { + 'X-Custom-Header': 'CustomValue', + 'X-Another-Header': 'AnotherValue', + 'X-Array-Header': ['Value1', 'Value2'], + 'X-Object-Header': { key: 'value' }, + }; + } + }]; + + const result = { + 'X-Custom-Header': 'CustomValue', + 'X-Another-Header': 'AnotherValue', + 'X-Array-Header': JSON.stringify(['Value1', 'Value2']), + 'X-Object-Header': JSON.stringify({ key: 'value' }) + }; + + const testingHeaders = { + application: 'testing', + AppID: 'typescriptTest', + accept: 'application/json', + 'content-type': 'type', + }; + + // const mockConnectionOptions : Snowflake.ConnectionOptions = { + // accessUrl: 'http://fakeaccount.snowflakecomputing.com', + // username: 'fakeusername', + // password: 'fakepassword', + // account: 'fakeaccount', + // authenticator: 'DEFAULT_AUTHENTICATOR', + // } + + const mockConnectionOptions = { + accessUrl: 'http://fakeaccount.snowflakecomputing.com', + username: 'fakeusername', + password: 'fakepassword', + account: 'fakeaccount', + authenticator: 'DEFAULT_AUTHENTICATOR', + }; + + it('verify custom headers includes in the first attempt', () => { + mockConnectionOptions.httpHeadersCustomizer = customHeaders; + const connectionConfig = new ConnectionConfig(mockConnectionOptions); + const firstHeaders = getHttpRequestHeaders(connectionConfig, { isRetry: false, url: 'http://fakeaccount.snowflakecomputing.com', headers: testingHeaders }); + delete firstHeaders['user-agent']; + verifyHeaders(firstHeaders, { ...testingHeaders, ...result }); + + const retryHeader = getHttpRequestHeaders(connectionConfig, { isRetry: true, url: 'http://fakeaccount.snowflakecomputing.com', headers: testingHeaders }); + delete retryHeader['user-agent']; + verifyHeaders(retryHeader, { ...testingHeaders }); + }); + + it('verify custom headers do not overwrite the original headers', () => { + mockConnectionOptions.httpHeadersCustomizer = [ + ...customHeaders, + { + applies: () => true, + newHeaders: function () { + return { + 'accept': 'OverwrittenValue', + 'content-type': 'OverwrittenType', + }; + } + } + ]; + const connectionConfig = new ConnectionConfig(mockConnectionOptions); + const firstHeaders = getHttpRequestHeaders(connectionConfig, { isRetry: false, url: 'http://fakeaccount.snowflakecomputing.com', headers: testingHeaders }); + delete firstHeaders['user-agent']; + verifyHeaders(firstHeaders, { ...testingHeaders, ...result }); + }); +}); + +// function verifyHeaders(headers: Record, expectedHeaders: Record) { +// assert.strictEqual(Object.keys(headers).length, Object.keys(expectedHeaders).length, 'Headers length mismatch') + +// for (const key in expectedHeaders) { +// if (headers.hasOwnProperty(key)) { +// assert.deepStrictEqual(headers[key], expectedHeaders[key], `Header ${key} does not match expected value`); +// } +// } +// } + +function verifyHeaders(headers, expectedHeaders) { + assert.strictEqual(Object.keys(headers).length, Object.keys(expectedHeaders).length, 'Headers length mismatch'); + + for (const key in expectedHeaders) { + if (Object.prototype.hasOwnProperty.call(headers, key)) { + assert.deepStrictEqual(headers[key], expectedHeaders[key], `Header ${key} does not match expected value`); + } + } +} \ No newline at end of file diff --git a/test/unit/mock/mock_http_client.js b/test/unit/mock/mock_http_client.js index 6924e975f..822ee200b 100644 --- a/test/unit/mock/mock_http_client.js +++ b/test/unit/mock/mock_http_client.js @@ -212,6 +212,7 @@ function buildRequestOutputMappings(clientInfo) { { method: 'POST', url: 'http://fakeaccount.snowflakecomputing.com/session/v1/login-request', + isRetry: false, headers: { 'Accept': 'application/json', @@ -319,6 +320,7 @@ function buildRequestOutputMappings(clientInfo) { { method: 'POST', url: 'http://fakeaccount.snowflakecomputing.com/session?delete=true', + isRetry: false, headers: { 'Accept': 'application/json', @@ -348,6 +350,7 @@ function buildRequestOutputMappings(clientInfo) { { method: 'POST', url: 'http://fakeaccount.snowflakecomputing.com/session?delete=true', + isRetry: false, headers: { 'Accept': 'application/json', @@ -378,6 +381,7 @@ function buildRequestOutputMappings(clientInfo) { { method: 'POST', url: 'http://fakeaccount.snowflakecomputing.com/queries/v1/query-request?requestId=foobar', + isRetry: false, headers: { 'Accept': 'application/snowflake', @@ -482,6 +486,7 @@ function buildRequestOutputMappings(clientInfo) { { method: 'POST', url: 'http://fakeaccount.snowflakecomputing.com/queries/v1/query-request?requestId=foobar', + isRetry: false, headers: { 'Accept': 'application/snowflake', @@ -585,6 +590,7 @@ function buildRequestOutputMappings(clientInfo) { { method: 'POST', url: 'http://fakeaccount.snowflakecomputing.com/queries/v1/query-request?requestId=foobar', + isRetry: false, headers: { 'Accept': 'application/snowflake', @@ -706,6 +712,7 @@ function buildRequestOutputMappings(clientInfo) { { method: 'POST', url: 'http://fakeaccount.snowflakecomputing.com/queries/v1/query-request?requestId=foobar', + isRetry: false, headers: { 'Accept': 'application/snowflake', @@ -751,6 +758,7 @@ function buildRequestOutputMappings(clientInfo) { { method: 'POST', url: 'http://fakeaccount.snowflakecomputing.com/queries/v1/query-request?requestId=SNOW-728803-requestId', + isRetry: false, headers: { 'Accept': 'application/snowflake', @@ -793,6 +801,7 @@ function buildRequestOutputMappings(clientInfo) { { method: 'POST', url: 'http://fakeaccount.snowflakecomputing.com/queries/v1/query-request?requestId=SNOW-728803-requestId', + isRetry: false, headers: { 'Accept': 'application/snowflake', @@ -824,6 +833,7 @@ function buildRequestOutputMappings(clientInfo) { { method: 'GET', url: 'http://fakeaccount.snowflakecomputing.com/queries/df2852ef-e082-4bb3-94a4-e540bf0e70c6/result?disableOfflineChunks=false', + isRetry: false, headers: { 'Accept': 'application/snowflake', @@ -922,6 +932,7 @@ function buildRequestOutputMappings(clientInfo) { { method: 'GET', url: 'http://fakeaccount.snowflakecomputing.com/queries/13f12818-de4c-41d2-bf19-f115ee8a5cc1/result?disableOfflineChunks=false', + isRetry: false, headers: { 'Accept': 'application/snowflake', @@ -958,6 +969,7 @@ function buildRequestOutputMappings(clientInfo) { { method: 'POST', url: 'http://fakeaccount.snowflakecomputing.com/queries/v1/abort-request', + isRetry: false, headers: { 'Accept': 'application/json', @@ -998,6 +1010,7 @@ function buildRequestOutputMappings(clientInfo) { { method: 'POST', url: 'http://fakeaccount.snowflakecomputing.com/queries/df2852ef-e082-4bb3-94a4-e540bf0e70c6/abort-request', + isRetry: false, headers: { 'Accept': 'application/json', @@ -1021,6 +1034,7 @@ function buildRequestOutputMappings(clientInfo) { { method: 'POST', url: 'http://fakeaccount.snowflakecomputing.com/queries/13f12818-de4c-41d2-bf19-f115ee8a5cc1/abort-request', + isRetry: false, headers: { 'Accept': 'application/json', @@ -1075,6 +1089,7 @@ function buildRequestOutputMappings(clientInfo) { { method: 'GET', url: 'http://fakeaccount.snowflakecomputing.com/queries/foobar/result?disableOfflineChunks=false', + isRetry: false, headers: { 'Accept': 'application/snowflake', @@ -1098,6 +1113,7 @@ function buildRequestOutputMappings(clientInfo) { { method: 'POST', url: 'http://fakeaccount.snowflakecomputing.com/queries/foobar/abort-request', + isRetry: false, headers: { 'Accept': 'application/json', @@ -1121,6 +1137,7 @@ function buildRequestOutputMappings(clientInfo) { { method: 'POST', url: 'http://fakeaccount.snowflakecomputing.com/queries/v1/query-request?requestId=b97fee20-a805-11e5-a0ab-ddd3321ed586', + isRetry: false, headers: { 'Accept': 'application/snowflake', @@ -1164,6 +1181,7 @@ function buildRequestOutputMappings(clientInfo) { { method: 'POST', url: 'http://fakeaccount.snowflakecomputing.com/queries/v1/abort-request', + isRetry: false, headers: { 'Accept': 'application/json', @@ -1197,6 +1215,7 @@ function buildRequestOutputMappings(clientInfo) { { method: 'POST', url: 'http://fakeaccount.snowflakecomputing.com/queries/v1/query-request?requestId=foobar', + isRetry: false, headers: { 'Accept': 'application/snowflake', @@ -1232,6 +1251,7 @@ function buildRequestOutputMappings(clientInfo) { { method: 'POST', url: 'http://fakeaccount.snowflakecomputing.com/session/v1/login-request', + isRetry: false, headers: { 'Accept': 'application/json', @@ -1342,6 +1362,7 @@ function buildRequestOutputMappings(clientInfo) { { method: 'POST', url: 'http://fakeaccount.snowflakecomputing.com/queries/v1/query-request?requestId=foobar', + isRetry: false, headers: { 'Accept': 'application/snowflake', @@ -1452,6 +1473,7 @@ function buildRequestOutputMappings(clientInfo) { { method: 'POST', url: 'http://fakeaccount.snowflakecomputing.com/session/v1/login-request', + isRetry: false, headers: { 'Accept': 'application/json', @@ -1561,6 +1583,7 @@ function buildRequestOutputMappings(clientInfo) { { method: 'POST', url: 'http://fakeaccount.snowflakecomputing.com/session/v1/login-request', + isRetry: false, headers: { 'Accept': 'application/json', @@ -1670,6 +1693,7 @@ function buildRequestOutputMappings(clientInfo) { { method: 'POST', url: 'http://fakeaccount.snowflakecomputing.com/session/v1/login-request', + isRetry: false, headers: { 'Accept': 'application/json', @@ -1732,6 +1756,7 @@ function buildRequestOutputMappings(clientInfo) { { method: 'POST', url: 'http://fakeaccount.snowflakecomputing.com/session?delete=true', + isRetry: false, headers: { 'Accept': 'application/json', @@ -1763,6 +1788,7 @@ function buildRequestOutputMappings(clientInfo) { { method: 'POST', url: 'http://fakeaccount.snowflakecomputing.com/session/v1/login-request', + isRetry: false, headers: { 'Accept': 'application/json', @@ -1825,6 +1851,7 @@ function buildRequestOutputMappings(clientInfo) { { method: 'POST', url: 'http://fakeaccount.snowflakecomputing.com/session?delete=true', + isRetry: false, headers: { 'Accept': 'application/json', @@ -1855,6 +1882,7 @@ function buildRequestOutputMappings(clientInfo) { { method: 'POST', url: 'http://fake504.snowflakecomputing.com/session/v1/login-request', + isRetry: false, headers: { 'Accept': 'application/json', @@ -1896,6 +1924,7 @@ function buildRequestOutputMappings(clientInfo) { { method: 'POST', url: 'http://fake504.snowflakecomputing.com/session/v1/login-request?clientStartTime=FIXEDTIMESTAMP&retryCount=1', + isRetry: false, headers: { 'Accept': 'application/json', @@ -1938,6 +1967,7 @@ function buildRequestOutputMappings(clientInfo) { { method: 'POST', url: 'http://fake504.snowflakecomputing.com/session/v1/login-request?clientStartTime=FIXEDTIMESTAMP&retryCount=2', + isRetry: false, headers: { 'Accept': 'application/json', @@ -2000,6 +2030,7 @@ function buildRequestOutputMappings(clientInfo) { { method: 'POST', url: 'http://fake504.snowflakecomputing.com/session?delete=true', + isRetry: false, headers: { 'Accept': 'application/json', @@ -2030,6 +2061,7 @@ function buildRequestOutputMappings(clientInfo) { { method: 'POST', url: 'http://fakeaccount.snowflakecomputing.com/session/heartbeat', + isRetry: false, headers: { 'Accept': 'application/json',