SNOW-736355: [HTAP] Add Retry Context to retried request

lib/connection/connection_config.js

@@ -49,6 +49,7 @@ const DEFAULT_PARAMS =
   'arrayBindingThreshold',
   'gcsUseDownscopedCredential',
   'forceStageBindError',
+  'includeRetryReason',
   'disableQueryContextCache',
 ];
 
@@ -493,6 +494,14 @@ function ConnectionConfig(options, validateCredentials, qaMode, clientInfo)
     }
   }
 
+  let includeRetryReason = true;
+  if (Util.exists(options.includeRetryReason)) {
+    Errors.checkArgumentValid(Util.isBoolean(options.includeRetryReason),
+    ErrorCodes.ERR_CONN_CREATE_INVALID_INCLUDE_RETRY_REASON);
+
+    includeRetryReason = options.includeRetryReason;
+  }
+
   /**
    * Returns an object that contains information about the proxy hostname, port,
    * etc. for when http requests are made.
@@ -757,6 +766,15 @@ function ConnectionConfig(options, validateCredentials, qaMode, clientInfo)
   };
 
   /**
+   * Returns whether the Retry reason is included or not in the retry url
+   *    
+   * @returns {Boolean}
+   */
+  this.getIncludeRetryReason = function () {
+    return includeRetryReason;
+  }
+
+  /** 
    * Returns whether the Query Context Cache is enabled or not by the configuration
    *
    * @returns {Boolean}

lib/connection/statement.js

@@ -1580,6 +1580,7 @@ function sendSfRequest(statementContext, options, appendQueryParamOnRetry)
   var numRetries = 0;
   var maxNumRetries = connectionConfig.getRetrySfMaxNumRetries();
   var sleep = connectionConfig.getRetrySfStartingSleepTime();
+  let lastStatusCodeForRetry;
 
   // create a function to send the request
   var sendRequest = function ()
@@ -1588,7 +1589,14 @@ function sendSfRequest(statementContext, options, appendQueryParamOnRetry)
     // retry, update the url
     if ((numRetries > 0) && appendQueryParamOnRetry)
     {
-      options.url = Util.url.appendParam(urlOrig, 'retry', true);
+      const retryOption = {
+        url: urlOrig,
+        retryCount: numRetries,
+        retryReason: lastStatusCodeForRetry,
+        includeRetryReason: connectionConfig.getIncludeRetryReason(),
+      }
+
+      options.url = Util.url.appendRetryParam(retryOption);
     }
 
     sf.request(options);
@@ -1606,6 +1614,7 @@ function sendSfRequest(statementContext, options, appendQueryParamOnRetry)
     {
       // increment the retry count
       numRetries++;
+      lastStatusCodeForRetry = err.response ? err.response.statusCode : 0
 
       // use exponential backoff with decorrelated jitter to compute the
       // next sleep time.

lib/constants/error_messages.js

@@ -65,6 +65,7 @@ exports[404038] = 'Invalid gcsUseDownscopedCredential. The specified value must
 exports[404039] = 'Invalid forceStageBindError. The specified value must be a number.';
 exports[404040] = 'Invalid browser timeout value. The specified value must be a positive number.';
 exports[404041] = 'Invalid disablQueryContextCache. The specified value must be a boolean.';
+exports[404042] = 'Invalid includeRetryReason. The specified value must be a boolean.'
 
 // 405001
 exports[405001] = 'Invalid callback. The specified value must be a function.';

lib/errors.js

@@ -70,6 +70,7 @@ codes.ERR_CONN_CREATE_INVALID_GCS_USE_DOWNSCOPED_CREDENTIAL = 404038;
 codes.ERR_CONN_CREATE_INVALID_FORCE_STAGE_BIND_ERROR = 404039;
 codes.ERR_CONN_CREATE_INVALID_BROWSER_TIMEOUT = 404040;
 codes.ERR_CONN_CREATE_INVALID_DISABLED_QUERY_CONTEXT_CACHE = 404041
+codes.ERR_CONN_CREATE_INVALID_INCLUDE_RETRY_REASON =404042
 
 // 405001
 codes.ERR_CONN_CONNECT_INVALID_CALLBACK = 405001;

lib/util.js

@@ -344,6 +344,15 @@ exports.url =
       }
 
       return url;
+    },
+
+    appendRetryParam: function (option) {
+      let retryUrl = this.appendParam(option.url, 'retryCount', option.retryCount);
+      if (option.includeRetryReason) {
+        retryUrl = this.appendParam(retryUrl, 'retryReason', option.retryReason);
+      }
+
+      return retryUrl;
     }
   };
 

test/unit/connection/connection_config_test.js

@@ -523,7 +523,18 @@ describe('ConnectionConfig: basic', function ()
           account: 'account',
           disableQueryContextCache: 1234
         },
-        errorCode: ErrorCodes.ERR_CONN_CREATE_INVALID_DISABLED_QUERY_CONTEXT_CACHE
+        errorCode: ErrorCodes.ERR_CONN_CREATE_INVALID_DISABLED_QUERY_CONTEXT_CACHE, 
+      },
+      {
+        name: 'invalid includeRetryReason',
+        options:
+        {
+          username: 'username',
+          password: 'password',
+          account: 'account',
+          includeRetryReason: 'invalid'
+        },
+        errorCode: ErrorCodes.ERR_CONN_CREATE_INVALID_INCLUDE_RETRY_REASON,
       },
     ];
 

test/unit/util_test.js

@@ -488,6 +488,40 @@ describe('Util', function ()
     }
   });
 
+  describe('Append retry parameters', function () {
+    const testCases =
+      [
+        {
+          testName: "test appending retry params with retry reason",
+          option: {
+            url: 'http://www.something.snowflakecomputing.com',
+            retryCount: 3,
+            retryReason: 429,
+            includeRetryReason: true,
+          },
+          result: 'http://www.something.snowflakecomputing.com?retryCount=3&retryReason=429'
+        },
+        {
+          testName: "test appending retry params without retry reason",
+          option: {
+            url: 'http://www.something.snowflakecomputing.com',
+            retryCount: 3,
+            retryReason: 429,
+            includeRetryReason: false,
+          },
+          result: 'http://www.something.snowflakecomputing.com?retryCount=3'
+        }
+      ];
+
+      for (let i = 0; i < testCases.length; i++) {
+        const testCase = testCases[i];
+        it(testCase.testName, function () {
+          const url = Util.url.appendRetryParam(testCase.option);
+          assert.strictEqual(url, testCase.result);
+        })
+      }
+  })
+
   it('Util.apply()', function ()
   {
     assert.strictEqual(Util.apply(null, null), null);