SNOW-896805: Replace urllib with axios (#623)

* SNOW-896805: Replace urllib with axios * SNOW-896805: Remove unnecessary code changes * SNOW-896805: Upgrade axios to the last version - 1.5.0 * SNOW-896805: Unzip GCS response without axios * SNOW-896805: Cleanup large result set error handling --------- Co-authored-by: Przemyslaw Motacki <[email protected]>
snowflakedb · Sep 11, 2023 · 7abaf84 · 7abaf84
1 parent e80fe5d
commit 7abaf84
Show file tree

Hide file tree

Showing 5 changed files with 61 additions and 28 deletions.
diff --git a/lib/agent/socket_util.js b/lib/agent/socket_util.js
@@ -74,6 +74,8 @@ exports.secureSocket = function (socket, host, agent, mock)
     // stop listening for the secure event
     socket.removeListener('secure', validate);
 
+    Logger.getInstance().trace('socket reused = %s', socket.isSessionReused());
+
     // if the server has resumed our existing session, unblock all
     // writes without performing any additional validation
     if (socket.isSessionReused())
@@ -106,8 +108,6 @@ exports.secureSocket = function (socket, host, agent, mock)
         socket.uncork();
       });
     }
-
-    Logger.getInstance().trace('socket reused = %s', socket.isSessionReused());
   };
 
   // when the socket is secure, perform additional validation

diff --git a/lib/http/base.js b/lib/http/base.js
@@ -6,6 +6,7 @@ const zlib = require('zlib');
 const Util = require('../util');
 const Errors = require('../errors');
 const Logger = require('../logger');
+const axios = require('axios');
 
 const DEFAULT_REQUEST_TIMEOUT = 360000;
 
@@ -76,7 +77,12 @@ HttpClient.prototype.request = function (options)
         body: body,
         timeout: timeout,
         requestOCSP: true,
-        rejectUnauthorized: true
+        rejectUnauthorized: true,
+        // axios does not know how to decode response with content-encoding GZIP (it should be gzip)
+        // that we receive from GCS, so let's get response as arraybuffer and unzip it outside axios
+        // issue in axios about case insensitive content-encoding is marked as won't fix: https://github.com/axios/axios/issues/4280
+        // for all other responses we manually parse jsons or other structures from the server so they need to be text
+        responseType: options.url.includes('storage.googleapis.com') ? 'arraybuffer' : 'text',
       };
 
     let mock;
@@ -97,21 +103,39 @@ HttpClient.prototype.request = function (options)
     requestOptions.httpsAgent.keepAlive = agentAndProxyOptions.agentOptions.keepAlive;
     requestOptions.retryDelay = this.constructExponentialBackoffStrategy();
 
-    request = require('urllib').request(requestOptions.url, requestOptions, async function (err, data, response)
-    {
-      // if a callback has been specified, normalize the
-      // response before passing it to the callback
-      if (Util.isFunction(options.callback))
-      {
-        if (err)
-        {
-          await options.callback(err, normalizeResponse(null), null);
+    request = axios.request(requestOptions).then(response => {
+      if (Util.isFunction(options.callback)) {
+        if (options.url.includes('storage.googleapis.com')) {
+          // we request that GCS returns body as arraybuffer, not text
+          // when it is GZIPped then we have to unzip it
+          // otherwise we should convert arraybuffer to string
+          try {
+            if (response.headers['content-encoding'] === 'GZIP') {
+              const unzippedData = zlib.gunzipSync(response.data).toString('utf-8');
+              return options.callback(null, normalizeResponse(response), unzippedData);
+            } else {
+              return options.callback(null, normalizeResponse(response), new TextDecoder('utf-8').decode(response.data));
+            }
+          } catch (e) {
+            return options.callback(e, null, null);
+          }
+        } else {
+          return options.callback(null, normalizeResponse(response), response.data);
         }
-        else
-        {
-          response.body = response.data.toString();
-          await options.callback(null, normalizeResponse(response), data.toString());
+      } else {
+        Logger.getInstance().trace(`Callback function was not provided for the call to ${options.url}`);
+        return null;
+      }
+    }).catch(err => {
+      if (Util.isFunction(options.callback)) {
+        if (err.response) { // axios returns error for not 2xx responses - let's unwrap it
+          options.callback(null, normalizeResponse(err.response), err.response.data);
+        } else {
+          options.callback(err, normalizeResponse(null), null);
         }
+        return null;
+      } else {
+        throw err;
       }
     });
   };
@@ -263,5 +287,10 @@ function normalizeResponse(response)
     };
   }
 
+  if (response) {
+    response.body = response.data; // converting axios response body to old expected body attribute
+    response.statusCode = response.status; // converting axios status to old expected statusCode
+  }
+
   return response;
 }
diff --git a/lib/services/large_result_set.js b/lib/services/large_result_set.js
@@ -26,12 +26,12 @@ function LargeResultSetService(connectionConfig, httpClient)
    * @param err Client error or null/undefined
    * @return {boolean}
    */
-  function isRetryableClientError(err)
-  {
+  function isRetryableClientError(err) {
     return err && (
-      err.code === 'ECONNRESET' || err.code === 'ETIMEDOUT' ||
-      // errors named with TimoutError suffix are thrown from urllib
-      err.name === 'ResponseTimeoutError' || err.name === 'ConnectionTimeoutError'
+        err.code === 'ECONNRESET' ||
+        err.code === 'ETIMEDOUT' ||
+        // error code ECONNABORTED is thrown from axios on timeout
+        (err.name === 'AxiosError' && err.code === 'ECONNABORTED')
     );
   }
 
@@ -44,7 +44,7 @@ function LargeResultSetService(connectionConfig, httpClient)
 
   function isUnsuccessfulResponse(response) {
     // even for 200 OK S3 can return xml error (large files are normally binary)
-    return response && (response.statusCode !== 200 || response.getResponseHeader('Content-Type') === 'application/xml')
+    return response && (response.statusCode !== 200 || response.getResponseHeader('Content-Type') === 'application/xml');
   }
 
   /**
@@ -101,7 +101,9 @@ function LargeResultSetService(connectionConfig, httpClient)
           }
         }
       }
-
+      if (response) {
+        Logger.getInstance().trace(`Response headers are: ${JSON.stringify(response.headers)}`);
+      }
       // if we have an error, clear the body
       if (err)
       {
@@ -111,7 +113,11 @@ function LargeResultSetService(connectionConfig, httpClient)
       // if a callback was specified, invoke it
       if (Util.isFunction(options.callback))
       {
-        options.callback(err, body);
+        try {
+          options.callback(err, body);
+        } catch (e) {
+          Logger.getInstance().error(`Callback failed with ${e}`);
+        }
       }
     };
 

diff --git a/lib/services/sf.js b/lib/services/sf.js
@@ -675,7 +675,6 @@ function StateAbstract(options)
     }
 
     return httpClient.request(realRequestOptions);
-    // comment //
   }
 
   this.snowflakeService = options.snowflakeService;

diff --git a/package.json b/package.json
@@ -10,7 +10,7 @@
     "asn1.js-rfc2560": "^5.0.0",
     "asn1.js-rfc5280": "^3.0.0",
     "aws-sdk": "^2.878.0",
-    "axios": "^0.27.2",
+    "axios": "^1.5.0",
     "big-integer": "^1.6.43",
     "bignumber.js": "^2.4.0",
     "binascii": "0.0.2",
@@ -33,7 +33,6 @@
     "simple-lru-cache": "^0.0.2",
     "string-similarity": "^4.0.4",
     "tmp": "^0.2.1",
-    "urllib": "^2.41.0",
     "uuid": "^3.3.2",
     "winston": "^3.1.0"
   },
@@ -77,4 +76,4 @@
     "url": "https://www.snowflake.com/"
   },
   "license": "Apache-2.0"
-}
+}
-Original file line number
+Diff line change
@@ Expand Up / @@ -675,7 +675,6 @@ function StateAbstract(options) @@
         }
         return httpClient.request(realRequestOptions);
-        // comment //
       }
       this.snowflakeService = options.snowflakeService;
@@ Expand Down @@