From 48683592c6857494af55ad186e0ce1faed386618 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Micha=C5=82=20Hofman?= <michal.hofman@snowflake.com>
Date: Wed, 6 Dec 2023 16:58:29 +0100
Subject: [PATCH] SNOW-950923 QueryID for PUT/GET without changes to the
 related exceptions

---
 .../IntegrationTests/SFPutGetTest.cs          |  5 ++--
 .../UnitTests/SFFileTransferAgentTests.cs     | 24 ++++++-------------
 Snowflake.Data/Client/SnowflakeDbException.cs |  9 +------
 Snowflake.Data/Core/ErrorMessages.resx        |  5 +---
 .../Core/FileTransfer/SFFileTransferAgent.cs  | 10 +-------
 Snowflake.Data/Core/SFStatement.cs            |  8 +++----
 6 files changed, 16 insertions(+), 45 deletions(-)

diff --git a/Snowflake.Data.Tests/IntegrationTests/SFPutGetTest.cs b/Snowflake.Data.Tests/IntegrationTests/SFPutGetTest.cs
index 49e835e7c..0f6126cab 100644
--- a/Snowflake.Data.Tests/IntegrationTests/SFPutGetTest.cs
+++ b/Snowflake.Data.Tests/IntegrationTests/SFPutGetTest.cs
@@ -545,10 +545,9 @@ string PutFile(
                     }
                     Assert.IsNull(reader.GetString((int)SFResultSet.PutGetResponseRowTypeInfo.ErrorDetails));
                 }
-                catch (SnowflakeDbException e)
+                catch (Exception e)
                 {
-                    queryId = e.QueryId;
-                    Assert.AreEqual(queryId, ((SnowflakeDbCommand)command).GetQueryId());
+                    queryId = ((SnowflakeDbCommand)command).GetQueryId();
                 }
             }
             return queryId;
diff --git a/Snowflake.Data.Tests/UnitTests/SFFileTransferAgentTests.cs b/Snowflake.Data.Tests/UnitTests/SFFileTransferAgentTests.cs
index d9210c65b..0cf147ff4 100644
--- a/Snowflake.Data.Tests/UnitTests/SFFileTransferAgentTests.cs
+++ b/Snowflake.Data.Tests/UnitTests/SFFileTransferAgentTests.cs
@@ -490,17 +490,15 @@ public void TestUploadThrowsArgumentExceptionForMissingRootDirectoryWithWildcard
 
             // Set command to upload
             _responseData.command = CommandTypes.UPLOAD.ToString();
-            _responseData.queryId = Guid.NewGuid().ToString();
             _fileTransferAgent = new SFFileTransferAgent(_putQuery,
                 _session,
                 _responseData,
                 _cancellationToken);
 
             // Act
-            SnowflakeDbException ex = Assert.Throws<SnowflakeDbException>(() => _fileTransferAgent.execute());
+            Exception ex = Assert.Throws<ArgumentException>(() => _fileTransferAgent.execute());
 
             // Assert
-            Assert.AreEqual(_responseData.queryId, ex.QueryId);
             Assert.That(ex.Message, Does.Match($"No file found for: {tempUploadRootDirectory}\\*/{tempUploadSecondDirectory}\\*/{mockFileName}"));
 
             for (int i = 0; i < numberOfDirectories; i++)
@@ -584,21 +582,17 @@ public void TestDownloadThrowsErrorFileNotFound()
 
             // Set command to download
             _responseData.command = CommandTypes.DOWNLOAD.ToString();
-            _responseData.queryId = Guid.NewGuid().ToString();
             _fileTransferAgent = new SFFileTransferAgent(GetQuery,
                 _session,
                 _responseData,
                 _cancellationToken);
 
             // Act
-            SnowflakeDbException ex = Assert.Throws<SnowflakeDbException>(() => _fileTransferAgent.execute());
+            Exception ex = Assert.Throws<AggregateException>(() => _fileTransferAgent.execute());
 
             // Assert
-            Assert.AreEqual(_responseData.queryId, ex.QueryId);
-            Assert.IsInstanceOf<AggregateException>(ex.InnerException);
-            var innerException = ((AggregateException)ex.InnerException)?.InnerExceptions[0];
-            Assert.IsInstanceOf<FileNotFoundException>(innerException);
-            Assert.That(innerException?.Message, Does.Match("Could not find file .*"));
+            Assert.IsInstanceOf<FileNotFoundException>(ex.InnerException);
+            Assert.That(ex.InnerException.Message, Does.Match("Could not find file .*"));
         }
 
         [Test]
@@ -615,21 +609,17 @@ public void TestDownloadThrowsErrorDirectoryNotFound()
 
             // Set command to download
             _responseData.command = CommandTypes.DOWNLOAD.ToString();
-            _responseData.queryId = Guid.NewGuid().ToString();
             _fileTransferAgent = new SFFileTransferAgent(GetQuery,
                 _session,
                 _responseData,
                 _cancellationToken);
 
             // Act
-            SnowflakeDbException ex = Assert.Throws<SnowflakeDbException>(() => _fileTransferAgent.execute());
+            Exception ex = Assert.Throws<AggregateException>(() => _fileTransferAgent.execute());
 
             // Assert
-            Assert.AreEqual(_responseData.queryId, ex.QueryId);
-            Assert.IsInstanceOf<AggregateException>(ex.InnerException);
-            var innerException = ((AggregateException)ex.InnerException)?.InnerExceptions[0];
-            Assert.IsInstanceOf<DirectoryNotFoundException>(innerException);
-            Assert.That(innerException?.Message, Does.Match("Could not find a part of the path .*"));
+            Assert.IsInstanceOf<DirectoryNotFoundException>(ex.InnerException);
+            Assert.That(ex.InnerException.Message, Does.Match("Could not find a part of the path .*"));
         }
     }
 }
diff --git a/Snowflake.Data/Client/SnowflakeDbException.cs b/Snowflake.Data/Client/SnowflakeDbException.cs
index f61349762..47ab44138 100755
--- a/Snowflake.Data/Client/SnowflakeDbException.cs
+++ b/Snowflake.Data/Client/SnowflakeDbException.cs
@@ -27,7 +27,7 @@ public sealed class SnowflakeDbException : DbException
 
         private int VendorCode;
 
-        public string QueryId { get; set; }
+        public string QueryId { get; }
 
         public override int ErrorCode
         {
@@ -45,13 +45,6 @@ public SnowflakeDbException(string sqlState, int vendorCode, string errorMessage
             QueryId = queryId;
         }
 
-        public SnowflakeDbException(SFError error, string queryId, Exception innerException)
-            : base(FormatExceptionMessage(error, new object[] {innerException.Message}, string.Empty, queryId), innerException)
-        {
-            VendorCode = error.GetAttribute<SFErrorAttr>().errorCode;
-            QueryId = queryId;
-        }
-
         public SnowflakeDbException(SFError error, params object[] args)
             : base(FormatExceptionMessage(error, args, string.Empty, string.Empty))
         {
diff --git a/Snowflake.Data/Core/ErrorMessages.resx b/Snowflake.Data/Core/ErrorMessages.resx
index 1bc4537ef..83235cf27 100755
--- a/Snowflake.Data/Core/ErrorMessages.resx
+++ b/Snowflake.Data/Core/ErrorMessages.resx
@@ -183,7 +183,4 @@
   <data name="BROWSER_RESPONSE_TIMEOUT" xml:space="preserve">
     <value>Browser response timed out after {0} seconds.</value>
   </data>
-  <data name="IO_ERROR_ON_GETPUT_COMMAND" xml:space="preserve">
-    <value>IO operation failed. Error: {0}</value>
-  </data>
-</root>
\ No newline at end of file
+</root>
diff --git a/Snowflake.Data/Core/FileTransfer/SFFileTransferAgent.cs b/Snowflake.Data/Core/FileTransfer/SFFileTransferAgent.cs
index 45d55f137..b0cbf1e65 100644
--- a/Snowflake.Data/Core/FileTransfer/SFFileTransferAgent.cs
+++ b/Snowflake.Data/Core/FileTransfer/SFFileTransferAgent.cs
@@ -231,15 +231,7 @@ public void execute()
             catch (Exception e)
             {
                 Logger.Error("Error while transferring file(s): " + e.Message);
-                if (e is SnowflakeDbException snowflakeException)
-                {
-                    if (snowflakeException.QueryId == null)
-                    {
-                        snowflakeException.QueryId = TransferMetadata.queryId;
-                    }
-                    throw snowflakeException;
-                }
-                throw new SnowflakeDbException(SFError.IO_ERROR_ON_GETPUT_COMMAND, TransferMetadata.queryId, e);
+                throw;
             }
             
         }
diff --git a/Snowflake.Data/Core/SFStatement.cs b/Snowflake.Data/Core/SFStatement.cs
index 670d48960..d23882778 100644
--- a/Snowflake.Data/Core/SFStatement.cs
+++ b/Snowflake.Data/Core/SFStatement.cs
@@ -357,16 +357,16 @@ internal SFBaseResultSet Execute(int timeout, string sql, Dictionary<string, Bin
                              describeOnly);
 
                     logger.Debug("PUT/GET queryId: " + (response.data != null ? response.data.queryId : "Unknown"));
+                    
+                    if (response.data != null)
+                        _lastQueryId = response.data.queryId;
 
                     SFFileTransferAgent fileTransferAgent =
                         new SFFileTransferAgent(trimmedSql, SfSession, response.data, CancellationToken.None);
 
                     // Start the file transfer
                     fileTransferAgent.execute();
-
-                    if (response.data != null)
-                        _lastQueryId = response.data.queryId;
-
+                    
                     // Get the results of the upload/download
                     return fileTransferAgent.result();
                 }