From 702e560588b6b978057deafdc6da32133b806c0b Mon Sep 17 00:00:00 2001 From: Gabriel Lomba Date: Thu, 21 Dec 2023 09:49:48 +0100 Subject: [PATCH] Reuse storage clients + destroy when no longer needed (fixes #734) (#735) --- lib/file_transfer_agent/azure_util.js | 6 ++-- lib/file_transfer_agent/s3_util.js | 8 ++--- test/unit/file_transfer_agent/azure_test.js | 28 +++++++++++----- test/unit/file_transfer_agent/s3_test.js | 36 ++++++++++++++++----- 4 files changed, 54 insertions(+), 24 deletions(-) diff --git a/lib/file_transfer_agent/azure_util.js b/lib/file_transfer_agent/azure_util.js index 6a54f2381..41ab158a8 100644 --- a/lib/file_transfer_agent/azure_util.js +++ b/lib/file_transfer_agent/azure_util.js @@ -89,7 +89,7 @@ function AzureUtil(azure, filestream) { */ this.getFileHeader = async function (meta, filename) { const stageInfo = meta['stageInfo']; - const client = this.createClient(stageInfo); + const client = meta['client']; const azureLocation = this.extractContainerNameAndPath(stageInfo['location']); const containerClient = client.getContainerClient(azureLocation.containerName); @@ -188,7 +188,7 @@ function AzureUtil(azure, filestream) { } const stageInfo = meta['stageInfo']; - const client = this.createClient(stageInfo); + const client = meta['client']; const azureLocation = this.extractContainerNameAndPath(stageInfo['location']); const blobName = azureLocation.path + meta['dstFileName']; @@ -230,7 +230,7 @@ function AzureUtil(azure, filestream) { */ this.nativeDownloadFile = async function (meta, fullDstPath) { const stageInfo = meta['stageInfo']; - const client = this.createClient(stageInfo); + const client = meta['client']; const azureLocation = this.extractContainerNameAndPath(stageInfo['location']); const blobName = azureLocation.path + meta['srcFileName']; diff --git a/lib/file_transfer_agent/s3_util.js b/lib/file_transfer_agent/s3_util.js index 2b134bf8c..0dc72e7af 100644 --- a/lib/file_transfer_agent/s3_util.js +++ b/lib/file_transfer_agent/s3_util.js @@ -114,7 +114,7 @@ function S3Util(s3, filestream) { */ this.getFileHeader = async function (meta, filename) { const stageInfo = meta['stageInfo']; - const client = this.createClient(stageInfo); + const client = meta['client']; const s3location = this.extractBucketNameAndPath(stageInfo['location']); const params = { @@ -194,8 +194,7 @@ function S3Util(s3, filestream) { s3Metadata[AMZ_MATDESC] = encryptionMetadata.matDesc; } - const stageInfo = meta['stageInfo']; - const client = this.createClient(stageInfo); + const client = meta['client']; const s3location = this.extractBucketNameAndPath(meta['stageInfo']['location']); @@ -235,8 +234,7 @@ function S3Util(s3, filestream) { * @param {Object} encryptionMetadata */ this.nativeDownloadFile = async function (meta, fullDstPath) { - const stageInfo = meta['stageInfo']; - const client = this.createClient(stageInfo); + const client = meta['client']; const s3location = this.extractBucketNameAndPath(meta['stageInfo']['location']); diff --git a/test/unit/file_transfer_agent/azure_test.js b/test/unit/file_transfer_agent/azure_test.js index 0a621de0e..3f9c409b9 100644 --- a/test/unit/file_transfer_agent/azure_test.js +++ b/test/unit/file_transfer_agent/azure_test.js @@ -20,15 +20,8 @@ describe('Azure client', function () { let Azure = null; let client = null; let filestream = null; + let meta = null; const dataFile = mockDataFile; - const meta = { - stageInfo: { - location: mockLocation, - path: mockTable + '/' + mockPath + '/', - creds: {} - }, - SHA256_DIGEST: mockDigest, - }; const encryptionMetadata = { key: mockKey, iv: mockIv, @@ -108,6 +101,18 @@ describe('Azure client', function () { filestream = require('filestream'); Azure = new SnowflakeAzureUtil(client, filestream); }); + beforeEach(function () { + const stageInfo = { + location: mockLocation, + path: mockTable + '/' + mockPath + '/', + creds: {} + }; + meta = { + stageInfo, + SHA256_DIGEST: mockDigest, + client: Azure.createClient(stageInfo), + }; + }); it('extract bucket name and path', async function () { verifyNameAndPath('sfc-eng-regression/test_sub_dir/', 'sfc-eng-regression', 'test_sub_dir/'); @@ -132,6 +137,7 @@ describe('Azure client', function () { client = require('client'); Azure = new SnowflakeAzureUtil(client); + meta['client'] = Azure.createClient(meta['stageInfo']); await Azure.getFileHeader(meta, dataFile); assert.strictEqual(meta['resultStatus'], resultStatus.RENEW_TOKEN); @@ -147,6 +153,7 @@ describe('Azure client', function () { client = require('client'); const Azure = new SnowflakeAzureUtil(client); + meta['client'] = Azure.createClient(meta['stageInfo']); await Azure.getFileHeader(meta, dataFile); assert.strictEqual(meta['resultStatus'], resultStatus.NOT_FOUND_FILE); @@ -162,6 +169,7 @@ describe('Azure client', function () { client = require('client'); Azure = new SnowflakeAzureUtil(client); + meta['client'] = Azure.createClient(meta['stageInfo']); await Azure.getFileHeader(meta, dataFile); assert.strictEqual(meta['resultStatus'], resultStatus.RENEW_TOKEN); @@ -177,6 +185,7 @@ describe('Azure client', function () { client = require('client'); Azure = new SnowflakeAzureUtil(client); + meta['client'] = Azure.createClient(meta['stageInfo']); await Azure.getFileHeader(meta, dataFile); assert.strictEqual(meta['resultStatus'], resultStatus.ERROR); @@ -193,6 +202,7 @@ describe('Azure client', function () { client = require('client'); filestream = require('filestream'); Azure = new SnowflakeAzureUtil(client, filestream); + meta['client'] = Azure.createClient(meta['stageInfo']); await Azure.uploadFile(dataFile, meta, encryptionMetadata); assert.strictEqual(meta['resultStatus'], resultStatus.UPLOADED); @@ -213,6 +223,7 @@ describe('Azure client', function () { client = require('client'); filestream = require('filestream'); Azure = new SnowflakeAzureUtil(client, filestream); + meta['client'] = Azure.createClient(meta['stageInfo']); await Azure.uploadFile(dataFile, meta, encryptionMetadata); assert.strictEqual(meta['resultStatus'], resultStatus.RENEW_TOKEN); @@ -233,6 +244,7 @@ describe('Azure client', function () { client = require('client'); filestream = require('filestream'); Azure = new SnowflakeAzureUtil(client, filestream); + meta['client'] = Azure.createClient(meta['stageInfo']); await Azure.uploadFile(dataFile, meta, encryptionMetadata); assert.strictEqual(meta['resultStatus'], resultStatus.NEED_RETRY); diff --git a/test/unit/file_transfer_agent/s3_test.js b/test/unit/file_transfer_agent/s3_test.js index 2750bfb70..2c1aa29bb 100644 --- a/test/unit/file_transfer_agent/s3_test.js +++ b/test/unit/file_transfer_agent/s3_test.js @@ -20,15 +20,8 @@ describe('S3 client', function () { let AWS; let s3; let filesystem; + let meta; const dataFile = mockDataFile; - const meta = { - stageInfo: { - location: mockLocation, - path: mockTable + '/' + mockPath + '/', - creds: {} - }, - SHA256_DIGEST: mockDigest, - }; const encryptionMetadata = { key: mockKey, iv: mockIv, @@ -59,6 +52,7 @@ describe('S3 client', function () { return new putObject; }; + this.destroy = function () {}; } return new S3; @@ -74,6 +68,18 @@ describe('S3 client', function () { AWS = new SnowflakeS3Util(s3, filesystem); }); + beforeEach(function () { + const stageInfo = { + location: mockLocation, + path: mockTable + '/' + mockPath + '/', + creds: {} + }; + meta = { + stageInfo, + SHA256_DIGEST: mockDigest, + client: AWS.createClient(stageInfo), + }; + }); it('extract bucket name and path', async function () { let result = AWS.extractBucketNameAndPath('sfc-eng-regression/test_sub_dir/'); @@ -117,6 +123,7 @@ describe('S3 client', function () { return new getObject; }; + this.destroy = function () {}; } return new S3; @@ -124,6 +131,7 @@ describe('S3 client', function () { }); s3 = require('s3'); const AWS = new SnowflakeS3Util(s3); + meta['client'] = AWS.createClient(meta['stageInfo']); await AWS.getFileHeader(meta, dataFile); assert.strictEqual(meta['resultStatus'], resultStatus.RENEW_TOKEN); @@ -144,6 +152,7 @@ describe('S3 client', function () { return new getObject; }; + this.destroy = function () {}; } return new S3; @@ -151,6 +160,7 @@ describe('S3 client', function () { }); s3 = require('s3'); const AWS = new SnowflakeS3Util(s3); + meta['client'] = AWS.createClient(meta['stageInfo']); await AWS.getFileHeader(meta, dataFile); assert.strictEqual(meta['resultStatus'], resultStatus.NOT_FOUND_FILE); @@ -171,6 +181,7 @@ describe('S3 client', function () { return new getObject; }; + this.destroy = function () {}; } return new S3; @@ -178,6 +189,7 @@ describe('S3 client', function () { }); s3 = require('s3'); const AWS = new SnowflakeS3Util(s3); + meta['client'] = AWS.createClient(meta['stageInfo']); await AWS.getFileHeader(meta, dataFile); assert.strictEqual(meta['resultStatus'], resultStatus.RENEW_TOKEN); @@ -198,6 +210,7 @@ describe('S3 client', function () { return new getObject; }; + this.destroy = function () {}; } return new S3; @@ -205,6 +218,7 @@ describe('S3 client', function () { }); s3 = require('s3'); const AWS = new SnowflakeS3Util(s3); + meta['client'] = AWS.createClient(meta['stageInfo']); await AWS.getFileHeader(meta, dataFile); assert.strictEqual(meta['resultStatus'], resultStatus.ERROR); @@ -230,6 +244,7 @@ describe('S3 client', function () { return new putObject; }; + this.destroy = function () {}; } return new S3; @@ -243,6 +258,7 @@ describe('S3 client', function () { s3 = require('s3'); filesystem = require('filesystem'); const AWS = new SnowflakeS3Util(s3, filesystem); + meta['client'] = AWS.createClient(meta['stageInfo']); await AWS.uploadFile(dataFile, meta, encryptionMetadata); assert.strictEqual(meta['resultStatus'], resultStatus.RENEW_TOKEN); @@ -263,6 +279,7 @@ describe('S3 client', function () { return new putObject; }; + this.destroy = function () {}; } return new S3; @@ -276,6 +293,7 @@ describe('S3 client', function () { s3 = require('s3'); filesystem = require('filesystem'); const AWS = new SnowflakeS3Util(s3, filesystem); + meta['client'] = AWS.createClient(meta['stageInfo']); await AWS.uploadFile(dataFile, meta, encryptionMetadata); assert.strictEqual(meta['resultStatus'], resultStatus.NEED_RETRY_WITH_LOWER_CONCURRENCY); @@ -296,6 +314,7 @@ describe('S3 client', function () { return new putObject; }; + this.destroy = function () {}; } return new S3; @@ -309,6 +328,7 @@ describe('S3 client', function () { s3 = require('s3'); filesystem = require('filesystem'); const AWS = new SnowflakeS3Util(s3, filesystem); + meta['client'] = AWS.createClient(meta['stageInfo']); await AWS.uploadFile(dataFile, meta, encryptionMetadata); assert.strictEqual(meta['resultStatus'], resultStatus.NEED_RETRY);