add endPoint implementation when the driver makes a request with Axios

snowflakedb · Dec 17, 2024 · 8a13cfc · 8a13cfc
1 parent 0aaa740
commit 8a13cfc
Show file tree

Hide file tree

Showing 2 changed files with 39 additions and 26 deletions.
diff --git a/lib/file_transfer_agent/gcs_util.js b/lib/file_transfer_agent/gcs_util.js
@@ -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,26 +442,39 @@ 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;
   };
 
+  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 endPoint;
+  };
+
   /**
-* Left strip the specified character from a string.
-*
-* @param {String} str
-* @param {Character} remove
-*
-* @returns {String}
-*/
+  * 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);

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);
       } );