SNOW-871416: [HTAP] Query Context Caching for Node (#631)

* Query Context Cache implementation

lib/connection/connection_config.js

@@ -48,7 +48,8 @@ const DEFAULT_PARAMS =
   'validateDefaultParameters',
   'arrayBindingThreshold',
   'gcsUseDownscopedCredential',
-  'forceStageBindError'
+  'forceStageBindError',
+  'disableQueryContextCache',
 ];
 
 function consolidateHostAndAccount(options)
@@ -468,6 +469,14 @@ function ConnectionConfig(options, validateCredentials, qaMode, clientInfo)
     forceStageBindError = options.forceStageBindError;
   }
 
+  let disableQueryContextCache = false;
+  if (Util.exists(options.disableQueryContextCache)) {
+    Errors.checkArgumentValid(Util.isBoolean(options.disableQueryContextCache),
+    ErrorCodes.ERR_CONN_CREATE_INVALID_DISABLED_QUERY_CONTEXT_CACHE);
+
+    disableQueryContextCache = options.disableQueryContextCache;
+  }
+
   if (validateDefaultParameters)
   {
     for (const [key] of Object.entries(options))
@@ -747,6 +756,15 @@ function ConnectionConfig(options, validateCredentials, qaMode, clientInfo)
     return forceStageBindError;
   };
 
+  /**
+   * Returns whether the Query Context Cache is enabled or not by the configuration
+   *
+   * @returns {Boolean}
+   */
+  this.getDisableQueryContextCache = function () {
+    return disableQueryContextCache;
+  }
+
   // save config options
   this.username = options.username;
   this.password = options.password;

lib/connection/result/result.js

@@ -53,6 +53,7 @@ function Result(options) {
   this._returnedRows = data.returned;
   this._totalRows = data.total;
   this._statementTypeId = data.statementTypeId;
+  this._queryContext = data.queryContext;
 
   // if no chunk headers were specified, but a query-result-master-key (qrmk)
   // was specified, build the chunk headers from the qrmk
@@ -73,7 +74,7 @@ function Result(options) {
 
   // convert the parameters array to a map
   parametersMap = {};
-  parametersArray = data.parameters;
+  parametersArray = data.parameters || [];
   for (index = 0, length = parametersArray.length; index < length; index++) {
     parameter = parametersArray[index];
     parametersMap[parameter.name] = parameter.value;
@@ -125,6 +126,10 @@ function Result(options) {
     this._statement,
     this._services);
 
+  this.getQueryContext = function () {
+    return this._queryContext;
+  }
+
   /* Disable the ChunkCache until the implementation is complete.
    * 
    *   // create a chunk cache and save a reference to it in case we need to

lib/connection/statement.js

@@ -576,6 +576,15 @@ function BaseStatement(
     sendCancelStatement(context, statement, callback);
   };
 
+  //Integration Testing purpose.
+  this.getQueryContextCacheSize = function () {
+    return services.sf.getQueryContextCacheSize();
+  }
+
+  this.getQueryContextDTOSize = function () {
+    return services.sf.getQueryContextDTO().entries.length;
+  }
+
   /**
    * Issues a request to get the statement result again.
    *
@@ -806,8 +815,9 @@ function createOnStatementRequestSuccRow(statement, context)
             connectionConfig: context.connectionConfig,
             rowMode: context.rowMode
           });
-        // save the query id
+
         context.queryId = context.result.getQueryId();
+        this.services.sf.deserializeQueryContext(context.result.getQueryContext());
       }
     }
     else
@@ -1334,6 +1344,10 @@ function sendRequestPreExec(statementContext, onResultAvailable)
     json.isInternal = statementContext.internal;
   }
 
+  if(!statementContext.disableQueryContextCache){
+    json.queryContextDTO = statementContext.services.sf.getQueryContextDTO();
+  }
+
   // use the snowflake service to issue the request
   sendSfRequest(statementContext,
     {
@@ -1389,6 +1403,10 @@ this.sendRequest = function (statementContext, onResultAvailable)
     json.isInternal = statementContext.internal;
   }
 
+  if(!statementContext.disableQueryContextCache){
+    json.queryContextDTO = statementContext.services.sf.getQueryContextDTO();
+  }
+
   var options =
   {
     method: 'POST',

lib/constants/error_messages.js

@@ -64,6 +64,7 @@ exports[404037] = 'Invalid arrayBindingThreshold. The specified value must be a
 exports[404038] = 'Invalid gcsUseDownscopedCredential. The specified value must be a boolean.';
 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.';
 
 // 405001
 exports[405001] = 'Invalid callback. The specified value must be a function.';

lib/errors.js

@@ -69,6 +69,7 @@ codes.ERR_CONN_CREATE_INVALID_ARRAY_BINDING_THRESHOLD = 404037;
 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
 
 // 405001
 codes.ERR_CONN_CONNECT_INVALID_CALLBACK = 405001;

lib/parameters.js

@@ -60,6 +60,7 @@ names.CLIENT_SESSION_KEEP_ALIVE_HEARTBEAT_FREQUENCY = 'CLIENT_SESSION_KEEP_ALIVE
 names.JS_TREAT_INTEGER_AS_BIGINT = 'JS_TREAT_INTEGER_AS_BIGINT';
 names.CLIENT_STAGE_ARRAY_BINDING_THRESHOLD = 'CLIENT_STAGE_ARRAY_BINDING_THRESHOLD';
 names.MULTI_STATEMENT_COUNT = 'MULTI_STATEMENT_COUNT';
+names.QUERY_CONTEXT_CACHE_SIZE = 'QUERY_CONTEXT_CACHE_SIZE';
 
 var parameters =
   [
@@ -106,6 +107,12 @@ var parameters =
         value: 1,
         desc: 'When 1, multi statement is disable, when 0, multi statement is unlimited'
       }),
+    new Parameter(
+      {
+        name: names.QUERY_CONTEXT_CACHE_SIZE,
+        value: 5,
+        desc: 'Query Context Cache Size'
+      }),
   ];
 
 // put all the parameters in a map so they're easy to retrieve and update