From 2f6a4abfe6089f78aa8ede769c3159af0ff2a635 Mon Sep 17 00:00:00 2001
From: sfc-gh-ext-simba-lf <lfarol@magnitude.com>
Date: Thu, 21 Sep 2023 14:48:38 -0700
Subject: [PATCH] SNOW-728803: Add test for special case when resubmitting
 requets

---
 lib/connection/statement.js        |   6 ++
 test/unit/mock/mock_http_client.js | 105 +++++++++++++++++++++++++++++
 test/unit/snowflake_test.js        |  65 ++++++++++++++++++
 3 files changed, 176 insertions(+)

diff --git a/lib/connection/statement.js b/lib/connection/statement.js
index baff69682..d8d8e46c1 100644
--- a/lib/connection/statement.js
+++ b/lib/connection/statement.js
@@ -420,6 +420,12 @@ function createContextPreExec(
     {
       statementContext.requestId = statementOptions.requestId;
     }
+
+    // SNOW-728803: In QA mode, set resubmitRequest to true to test that sqlText
+    // is not overwritten if a sqlText is already specified by the user
+    if (statementOptions.requestId === 'SNOW-728803-requestId') {
+      statementContext.resubmitRequest = true;
+    }
   }
 
   return statementContext;
diff --git a/test/unit/mock/mock_http_client.js b/test/unit/mock/mock_http_client.js
index cfefe231a..483c337ff 100644
--- a/test/unit/mock/mock_http_client.js
+++ b/test/unit/mock/mock_http_client.js
@@ -697,6 +697,111 @@ function buildRequestOutputMappings(clientInfo)
             }
         }
     },
+    {
+      request:
+        {
+          method: 'POST',
+          url: 'http://fakeaccount.snowflakecomputing.com/queries/v1/query-request?requestId=',
+          headers:
+            {
+              'Accept': 'application/snowflake',
+              'Authorization': 'Snowflake Token="SESSION_TOKEN"',
+              'Content-Type': 'application/json'
+            },
+          json:
+            {
+              disableOfflineChunks: false,
+              sqlText: 'select 1;'
+            }
+        },
+      output:
+        {
+          err: null,
+          response:
+            {
+              statusCode: 401,
+              statusMessage: 'Unauthorized',
+              body:
+                {
+                  'data' : null,
+                  'code': '390103',
+                  'message': 'Session token not found in the request data.',
+                  'success': false,
+                }
+            }
+        }
+    },
+    {
+      request:
+        {
+          method: 'POST',
+          url: 'http://fakeaccount.snowflakecomputing.com/queries/v1/query-request?requestId=SNOW-728803-requestId',
+          headers:
+            {
+              'Accept': 'application/snowflake',
+              'Authorization': 'Snowflake Token="SESSION_TOKEN"',
+              'Content-Type': 'application/json'
+            },
+          json:
+            {
+              disableOfflineChunks: false,
+              sqlText: 'SELECT \'Error retrieving query results for request id: SNOW-728803-requestId, please use RESULT_SCAN instead\' AS ErrorMessage;'
+            }
+        },
+      output:
+        {
+          err: null,
+          response:
+            {
+              body:
+                {
+                  'message': 'The specified sqlText should not be overwritten when resubmitting the request',
+                  'success': false
+                }
+            }
+        }
+    },
+    {
+      request:
+        {
+          method: 'POST',
+          url: 'http://fakeaccount.snowflakecomputing.com/queries/v1/query-request?requestId=SNOW-728803-requestId',
+          headers:
+            {
+              'Accept': 'application/snowflake',
+              'Authorization': 'Snowflake Token="SESSION_TOKEN"',
+              'Content-Type': 'application/json'
+            },
+          json:
+            {
+              disableOfflineChunks: false,
+              sqlText: 'select 1;'
+            }
+        },
+      output:
+        {
+          err: null,
+          response:
+            {
+              statusCode: 200,
+              statusMessage: 'OK',
+              body:
+                {
+                  'data':
+                    {
+                      'parameters': [],
+                      'rowtype': [],
+                      'rowset': [['1']],
+                      'total': 1,
+                      'returned': 1
+                    },
+                  'message': null,
+                  'code': null,
+                  'success': true
+                }
+            }
+        }
+    },
     {
       request:
         {
diff --git a/test/unit/snowflake_test.js b/test/unit/snowflake_test.js
index a50c7de05..8bb6d7afd 100644
--- a/test/unit/snowflake_test.js
+++ b/test/unit/snowflake_test.js
@@ -936,6 +936,71 @@ describe('connection.execute() statement failure', function ()
   });
 });
 
+describe('connection.execute() with requestId', function () {
+  const connection = snowflake.createConnection(connectionOptions);
+  const sqlText = 'select 1;';
+  const requestId = 'SNOW-728803-requestId';
+
+  it('keep original sqlText when resubmitting requests', function (done) {
+    let statement;
+
+    async.series(
+      [
+        function (callback) {
+          connection.connect(function (err, conn) {
+            assert.ok(!err, 'there should be no error');
+            assert.strictEqual(conn, connection,
+              'the connect() callback should be invoked with the statement');
+
+            callback();
+          });
+        },
+        function (callback) {
+          // 1st request fails with an error
+          statement = connection.execute(
+            {
+              sqlText: sqlText,
+              complete: function (err, stmt) {
+                assert.ok(err, 'there should be an error');
+                assert.strictEqual(stmt, statement,
+                  'the execute() callback should be invoked with the statement');
+
+                callback();
+              }
+            });
+        },
+        function (callback) {
+          // 2nd request with sqlText and requestId specified
+          statement = connection.execute(
+            {
+              sqlText: sqlText,
+              requestId: requestId,
+              complete: function (err, stmt) {
+                // if there's an error, fail the test with the error
+                if (err) {
+                  done(err);
+                }
+
+                assert.ok(!err, 'there should be no error');
+                assert.strictEqual(stmt, statement,
+                  'the execute() callback should be invoked with the statement');
+
+                callback();
+              }
+            });
+
+          // the sql text and request id should be the same as what was passed
+          // in
+          assert.strictEqual(statement.getSqlText(), sqlText);
+          assert.strictEqual(statement.getRequestId(), requestId);
+        }
+      ],
+      function () {
+        done();
+      });
+  });
+});
+
 describe('too many concurrent requests', function ()
 {
   it('too many concurrent requests per user', function (done)