Merge branch 'master' into SNOW-1789666-azure-core-http-deprecated

snowflakedb · Jan 9, 2025 · b8d6b47 · b8d6b47
2 parents 44a2ca0 + cd3994a
commit b8d6b47
Show file tree

Hide file tree

Showing 5 changed files with 77 additions and 35 deletions.
diff --git a/index.d.ts b/index.d.ts
@@ -874,8 +874,8 @@ declare module 'snowflake-sdk' {
     }
 
     export interface FileAndStageBindStatement extends RowStatement {
-        hasNext(): () => boolean;
-        NextResult(): () => void;
+        hasNext: () => boolean;
+        NextResult: () => void;
     }
 
     export interface SnowflakeErrorExternal extends Error {

diff --git a/lib/file_transfer_agent/gcs_util.js b/lib/file_transfer_agent/gcs_util.js
@@ -4,7 +4,7 @@
 
 const EncryptionMetadata = require('./encrypt_util').EncryptionMetadata;
 const FileHeader = require('./file_util').FileHeader;
-const { shouldPerformGCPBucket } = require('../util');
+const { shouldPerformGCPBucket, lstrip } = require('../util');
 
 const GCS_METADATA_PREFIX = 'x-goog-meta-';
 const SFC_DIGEST = 'sfc-digest';
@@ -26,7 +26,11 @@ const EXPIRED_TOKEN = 'ExpiredToken';
 
 const ERRORNO_WSAECONNABORTED = 10053;  // network connection was aborted
 
-// GCS Location
+/** 
+ * @typedef {object} GCSLocation
+ * @property {string} bucketName
+ * @property {string} path 
+ */
 function GCSLocation(bucketName, path) {
   return {
     'bucketName': bucketName,
@@ -69,14 +73,7 @@ function GCSUtil(httpclient, filestream) {
         }
       });
 
-      //TODO: SNOW-1789759 hardcoded region will be replaced in the future
-      const isRegionalUrlEnabled = (stageInfo.region).toLowerCase() === 'me-central2' || stageInfo.useRegionalUrl;
-      let endPoint = null;
-      if (stageInfo['endPoint']) {
-        endPoint = stageInfo['endPoint'];
-      } else if (isRegionalUrlEnabled) {
-        endPoint = `storage.${stageInfo.region.toLowerCase()}.rep.googleapis.com`;
-      }
+      const endPoint = this.getGCSCustomEndPoint(stageInfo);
       const storage = endPoint ? new Storage({ interceptors_: interceptors, apiEndpoint: endPoint }) : new Storage({ interceptors_: interceptors });
       client = { gcsToken: gcsToken, gcsClient: storage };
     } else {
@@ -91,7 +88,7 @@ function GCSUtil(httpclient, filestream) {
   *
   * @param {String} stageLocation
   *
-  * @returns {Object}
+  * @returns {GCSLocation}
   */
   this.extractBucketNameAndPath = function (stageLocation) {
     let containerName = stageLocation;
@@ -135,7 +132,7 @@ function GCSUtil(httpclient, filestream) {
             }
           });
       } else {
-        const url = this.generateFileURL(meta['stageInfo']['location'], lstrip(filename, '/'));
+        const url = this.generateFileURL(meta.stageInfo, lstrip(filename, '/'));
         const accessToken = meta['client'].gcsToken;
         const gcsHeaders = { 'Authorization': `Bearer ${accessToken}` };
         let encryptionMetadata;
@@ -239,7 +236,7 @@ function GCSUtil(httpclient, filestream) {
     if (!uploadUrl) {
       const tempFilename = meta['dstFileName'].substring(meta['dstFileName'].indexOf('/') + 1, meta['dstFileName'].length);
 
-      uploadUrl = this.generateFileURL(meta['stageInfo']['location'], tempFilename);
+      uploadUrl = this.generateFileURL(meta.stageInfo, tempFilename);
       accessToken = meta['client'].gcsToken;
     }
     let contentEncoding = '';
@@ -346,7 +343,7 @@ function GCSUtil(httpclient, filestream) {
 
     if (!downloadUrl) {
       downloadUrl = this.generateFileURL(
-        meta.stageInfo['location'], lstrip(meta['srcFileName'], '/')
+        meta.stageInfo, lstrip(meta['srcFileName'], '/')
       );
       accessToken = meta['client'].gcsToken;
       gcsHeaders = { 'Authorization': `Bearer ${accessToken}` };
@@ -445,32 +442,30 @@ function GCSUtil(httpclient, filestream) {
   /**
   * Generate file URL based on bucket.
   *
-  * @param {String} stageLocation
+  * @param {Object} stageInfo
   * @param {String} filename
   *
   * @returns {String}
   */
-  this.generateFileURL = function (stageLocation, filename) {
-    const gcsLocation = this.extractBucketNameAndPath(stageLocation);
+  this.generateFileURL = function (stageInfo, filename) {
+    const gcsLocation = this.extractBucketNameAndPath(stageInfo.location);
     const fullFilePath = `${gcsLocation.path}${filename}`;
-    const link = 'https://storage.googleapis.com/' + gcsLocation.bucketName + '/' + fullFilePath;
+    const endPoint = this.getGCSCustomEndPoint(stageInfo);
+    const link = `${endPoint != null ? endPoint : 'https://storage.googleapis.com'}/${gcsLocation.bucketName}/${fullFilePath}`;
     return link;
   };
 
-  /**
-* Left strip the specified character from a string.
-*
-* @param {String} str
-* @param {Character} remove
-*
-* @returns {String}
-*/
-  function lstrip(str, remove) {
-    while (str.length > 0 && remove.indexOf(str.charAt(0)) !== -1) {
-      str = str.substr(1);
+  this.getGCSCustomEndPoint = function (stageInfo) {
+    //TODO: SNOW-1789759 hardcoded region will be replaced in the future
+    const isRegionalUrlEnabled = (stageInfo.region).toLowerCase() === 'me-central2' || stageInfo.useRegionalUrl;
+    let endPoint = null;
+    if (stageInfo['endPoint']) {
+      endPoint = stageInfo['endPoint'];
+    } else if (isRegionalUrlEnabled) {
+      endPoint = `storage.${stageInfo.region.toLowerCase()}.rep.googleapis.com`;
     }
-    return str;
-  }
+    return endPoint;
+  };
 }
 
 module.exports = GCSUtil;
diff --git a/lib/util.js b/lib/util.js
@@ -753,4 +753,19 @@ exports.isNotEmptyAsString = function (variable) {
 
 exports.isWindows = function () {
   return os.platform() === 'win32';
+};
+
+/**
+* Left strip the specified character from a string.
+*
+* @param {String} str
+* @param {Character} remove
+*
+* @returns {String}
+*/
+exports.lstrip = function (str, remove) {
+  while (str.length > 0 && remove.indexOf(str.charAt(0)) !== -1) {
+    str = str.substr(1);
+  }
+  return str;
 };
diff --git a/test/unit/file_transfer_agent/gcs_test.js b/test/unit/file_transfer_agent/gcs_test.js
@@ -34,7 +34,10 @@ describe('GCS client', function () {
     meta = {
       stageInfo: {
         location: mockLocation,
-        path: mockTable + '/' + mockPath + '/'
+        path: mockTable + '/' + mockPath + '/',
+        endPoint: null,
+        useRegionalUrl: false,
+        region: 'mockLocation',
       },
       presignedUrl: mockPresignedUrl,
       dstFileName: mockPresignedUrl,
@@ -133,7 +136,7 @@ describe('GCS client', function () {
 
     testCases.forEach(({ name, stageInfo, result }) => {
       it(name, () => {
-        const client = GCS.createClient({ ...stageInfo, ...meta.stageInfo, creds: { GCS_ACCESS_TOKEN: 'mockToken' } });
+        const client = GCS.createClient({ ...meta.stageInfo, ...stageInfo,  creds: { GCS_ACCESS_TOKEN: 'mockToken' } });
         assert.strictEqual(client.gcsClient.apiEndpoint, result);
       } );
 

diff --git a/test/unit/util_test.js b/test/unit/util_test.js
@@ -1215,4 +1215,33 @@ describe('Util', function () {
     }
   });
 
+  describe('lstrip function Test', function () {
+    const testCases = [
+      {
+        name: 'remove consecutive characters /',
+        str: '///////////helloworld',
+        remove: '/',
+        result: 'helloworld'
+      },
+      {
+        name: 'when the first character is not matched with the remove character',
+        str: '/\\/\\helloworld',
+        remove: '\\',
+        result: '/\\/\\helloworld'
+      },
+      {
+        name: 'when the first and the third characters are matched',
+        str: '@1@12345helloworld',
+        remove: '@',
+        result: '1@12345helloworld'
+      },
+    ];
+
+    for (const { name, str, remove,  result } of testCases) {
+      it(name, function () {
+        assert.strictEqual(Util.lstrip(str, remove), result);
+      });
+    }
+  });
+
 });