From 1a7d139f71443060234f5ab45e4eb4830549cbe1 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Micha=C5=82=20Hofman?= <michal.hofman@snowflake.com>
Date: Mon, 19 Feb 2024 13:05:58 +0100
Subject: [PATCH] SNOW-950923 revert repackaging exception when non-put/get
 command are executed

---
 Snowflake.Data/Core/SFStatement.cs | 141 +++++++++++++++++------------
 1 file changed, 83 insertions(+), 58 deletions(-)

diff --git a/Snowflake.Data/Core/SFStatement.cs b/Snowflake.Data/Core/SFStatement.cs
index 974bd7918..ef55cca9e 100644
--- a/Snowflake.Data/Core/SFStatement.cs
+++ b/Snowflake.Data/Core/SFStatement.cs
@@ -343,71 +343,47 @@ internal SFBaseResultSet Execute(int timeout, string sql, Dictionary<string, Bin
         {
             // Trim the sql query and check if this is a PUT/GET command
             string trimmedSql = TrimSql(sql);
-
             try
             {
                 if (IsPutOrGetCommand(trimmedSql))
                 {
-                    isPutGetQuery = true;
-                    PutGetExecResponse response =
-                        ExecuteHelper<PutGetExecResponse, PutGetResponseData>(
-                             timeout,
-                             sql,
-                             bindings,
-                             describeOnly);
+                    return ExecuteSqlWithPutGet(timeout, trimmedSql, bindings, describeOnly);
+                }
 
-                    logger.Debug("PUT/GET queryId: " + (response.data != null ? response.data.queryId : "Unknown"));
+                return ExecuteSqlOtherThanPutGet(timeout, trimmedSql, bindings, describeOnly);
+            }
+            finally
+            {
+                CleanUpCancellationTokenSources();
+                ClearQueryRequestId();
+            }
+        }
 
-                    SFFileTransferAgent fileTransferAgent =
-                        new SFFileTransferAgent(trimmedSql, SfSession, response.data, CancellationToken.None);
+        internal SFBaseResultSet ExecuteSqlWithPutGet(int timeout, string sql, Dictionary<string, BindingDTO> bindings, bool describeOnly)
+        {
+            try
+            {
+                isPutGetQuery = true;
+                PutGetExecResponse response =
+                    ExecuteHelper<PutGetExecResponse, PutGetResponseData>(
+                        timeout,
+                        sql,
+                        bindings,
+                        describeOnly);
 
-                    // Start the file transfer
-                    fileTransferAgent.execute();
+                logger.Debug("PUT/GET queryId: " + (response.data != null ? response.data.queryId : "Unknown"));
 
-                    if (response.data != null)
-                        _lastQueryId = response.data.queryId;
+                SFFileTransferAgent fileTransferAgent =
+                    new SFFileTransferAgent(sql, SfSession, response.data, CancellationToken.None);
 
-                    // Get the results of the upload/download
-                    return fileTransferAgent.result();
-                }
-                else
-                {
-                    int arrayBindingThreshold = 0;
-                    if (SfSession.ParameterMap.ContainsKey(SFSessionParameter.CLIENT_STAGE_ARRAY_BINDING_THRESHOLD))
-                    {
-                        String val = (String)SfSession.ParameterMap[SFSessionParameter.CLIENT_STAGE_ARRAY_BINDING_THRESHOLD];
-                        arrayBindingThreshold = Int32.Parse(val);
-                    }
-
-                    int numBinding = GetBindingCount(bindings);
-
-                    if (0 < arrayBindingThreshold
-                        && arrayBindingThreshold <= numBinding
-                        && !describeOnly)
-                    {
-                        try
-                        {
-                            AssignQueryRequestId();
-                            SFBindUploader uploader = new SFBindUploader(SfSession, _requestId);
-                            uploader.Upload(bindings);
-                            _bindStage = uploader.getStagePath();
-                            ClearQueryRequestId();
-                        }
-                        catch (Exception e)
-                        {
-                            logger.Warn("Exception encountered trying to upload binds to stage. Attaching binds in payload instead. {0}", e);
-                        }
-                    }
+                // Start the file transfer
+                fileTransferAgent.execute();
 
-                    QueryExecResponse response =
-                        ExecuteHelper<QueryExecResponse, QueryExecResponseData>(
-                             timeout,
-                             sql,
-                             bindings,
-                             describeOnly);
+                if (response.data != null)
+                    _lastQueryId = response.data.queryId;
 
-                    return BuildResultSet(response, CancellationToken.None);
-                }
+                // Get the results of the upload/download
+                return fileTransferAgent.result();
             }
             catch (SnowflakeDbException ex)
             {
@@ -420,13 +396,62 @@ internal SFBaseResultSet Execute(int timeout, string sql, Dictionary<string, Bin
                 logger.Error("Query execution failed.", ex);
                 throw new SnowflakeDbException(ex, SFError.INTERNAL_ERROR);
             }
-            finally
+        }
+        
+        internal SFBaseResultSet ExecuteSqlOtherThanPutGet(int timeout, string sql, Dictionary<string, BindingDTO> bindings, bool describeOnly)
+        {
+            try
             {
-                CleanUpCancellationTokenSources();
-                ClearQueryRequestId();
+                int arrayBindingThreshold = 0;
+                if (SfSession.ParameterMap.ContainsKey(SFSessionParameter.CLIENT_STAGE_ARRAY_BINDING_THRESHOLD))
+                {
+                    String val =
+                        (String)SfSession.ParameterMap[SFSessionParameter.CLIENT_STAGE_ARRAY_BINDING_THRESHOLD];
+                    arrayBindingThreshold = Int32.Parse(val);
+                }
+
+                int numBinding = GetBindingCount(bindings);
+
+                if (0 < arrayBindingThreshold
+                    && arrayBindingThreshold <= numBinding
+                    && !describeOnly)
+                {
+                    try
+                    {
+                        AssignQueryRequestId();
+                        SFBindUploader uploader = new SFBindUploader(SfSession, _requestId);
+                        uploader.Upload(bindings);
+                        _bindStage = uploader.getStagePath();
+                        ClearQueryRequestId();
+                    }
+                    catch (Exception e)
+                    {
+                        logger.Warn(
+                            "Exception encountered trying to upload binds to stage. Attaching binds in payload instead. {0}",
+                            e);
+                    }
+                }
+
+                QueryExecResponse response =
+                    ExecuteHelper<QueryExecResponse, QueryExecResponseData>(
+                        timeout,
+                        sql,
+                        bindings,
+                        describeOnly);
+
+                return BuildResultSet(response, CancellationToken.None);
+            }
+            catch (Exception ex)
+            {
+                logger.Error("Query execution failed.", ex);
+                if (ex is SnowflakeDbException snowflakeDbException)
+                {
+                    _lastQueryId = snowflakeDbException.QueryId ?? _lastQueryId;
+                }
+                throw;
             }
         }
-
+        
         internal async Task<SFBaseResultSet> GetResultWithIdAsync(string resultId, CancellationToken cancellationToken)
         {
             var req = BuildResultRequestWithId(resultId);