From b1464100aa5e80fdb968555e3eeca8540c1c9f20 Mon Sep 17 00:00:00 2001 From: Dariusz Stempniak Date: Wed, 16 Oct 2024 11:59:46 +0200 Subject: [PATCH] SNOW-1736920 Fix bindings uploading and error handling for GCP (#1041) --- .../IntegrationTests/SFBindTestIT.cs | 1 + Snowflake.Data.Tests/Mock/MockGCSClient.cs | 22 +++++++++++++------ .../UnitTests/SFGCSClientTest.cs | 22 +++++++++---------- .../FileTransfer/StorageClient/SFGCSClient.cs | 16 ++++++++++---- Snowflake.Data/Core/SFStatement.cs | 16 +++++++++----- 5 files changed, 48 insertions(+), 29 deletions(-) diff --git a/Snowflake.Data.Tests/IntegrationTests/SFBindTestIT.cs b/Snowflake.Data.Tests/IntegrationTests/SFBindTestIT.cs index 1683700cb..956362fe8 100755 --- a/Snowflake.Data.Tests/IntegrationTests/SFBindTestIT.cs +++ b/Snowflake.Data.Tests/IntegrationTests/SFBindTestIT.cs @@ -649,6 +649,7 @@ public void TestPutArrayBind() var count = cmd.ExecuteNonQuery(); Assert.AreEqual(total * 3, count); + cmd.Parameters.Clear(); cmd.CommandText = $"SELECT * FROM {TableName}"; IDataReader reader = cmd.ExecuteReader(); Assert.IsTrue(reader.Read()); diff --git a/Snowflake.Data.Tests/Mock/MockGCSClient.cs b/Snowflake.Data.Tests/Mock/MockGCSClient.cs index cb36918ae..a25d4279a 100644 --- a/Snowflake.Data.Tests/Mock/MockGCSClient.cs +++ b/Snowflake.Data.Tests/Mock/MockGCSClient.cs @@ -25,7 +25,7 @@ class MockGCSClient internal const string GcsFileContent = "GCSClientTest"; // Create a mock response for GetFileHeader - static internal HttpWebResponse CreateResponseForFileHeader(HttpStatusCode httpStatusCode) + internal static HttpWebResponse CreateResponseForFileHeader(HttpStatusCode httpStatusCode) { var response = new Mock(); @@ -46,14 +46,18 @@ static internal HttpWebResponse CreateResponseForFileHeader(HttpStatusCode httpS } // Create a mock response for UploadFile - static internal HttpWebResponse CreateResponseForUploadFile(HttpStatusCode httpStatusCode) + internal static HttpWebResponse CreateResponseForUploadFile(HttpStatusCode? httpStatusCode) { var response = new Mock(); - if (httpStatusCode != HttpStatusCode.OK) + if (httpStatusCode is null) + { + throw new WebException("Mock GCS Error - no response", null, 0, null); + } + else if (httpStatusCode != HttpStatusCode.OK) { response.SetupGet(c => c.StatusCode) - .Returns(httpStatusCode); + .Returns(httpStatusCode.Value); throw new WebException("Mock GCS Error", null, 0, response.Object); } @@ -61,11 +65,15 @@ static internal HttpWebResponse CreateResponseForUploadFile(HttpStatusCode httpS } // Create a mock response for DownloadFile - static internal HttpWebResponse CreateResponseForDownloadFile(HttpStatusCode httpStatusCode) + internal static HttpWebResponse CreateResponseForDownloadFile(HttpStatusCode? httpStatusCode) { var response = new Mock(); - if (httpStatusCode == HttpStatusCode.OK) + if (httpStatusCode is null) + { + throw new WebException("Mock GCS Error - no response", null, 0, null); + } + else if (httpStatusCode == HttpStatusCode.OK) { response.Setup(c => c.Headers).Returns(new WebHeaderCollection()); response.Object.Headers.Add(SFGCSClient.GCS_METADATA_ENCRYPTIONDATAPROP, @@ -82,7 +90,7 @@ static internal HttpWebResponse CreateResponseForDownloadFile(HttpStatusCode htt else { response.SetupGet(c => c.StatusCode) - .Returns(httpStatusCode); + .Returns(httpStatusCode.Value); throw new WebException("Mock GCS Error", null, 0, response.Object); } diff --git a/Snowflake.Data.Tests/UnitTests/SFGCSClientTest.cs b/Snowflake.Data.Tests/UnitTests/SFGCSClientTest.cs index 925ce4c98..0fad57542 100644 --- a/Snowflake.Data.Tests/UnitTests/SFGCSClientTest.cs +++ b/Snowflake.Data.Tests/UnitTests/SFGCSClientTest.cs @@ -223,16 +223,14 @@ private void AssertForGetFileHeaderTests(ResultStatus expectedResultStatus, File [TestCase(HttpStatusCode.Forbidden, ResultStatus.NEED_RETRY)] [TestCase(HttpStatusCode.InternalServerError, ResultStatus.NEED_RETRY)] [TestCase(HttpStatusCode.ServiceUnavailable, ResultStatus.NEED_RETRY)] - public void TestUploadFile(HttpStatusCode httpStatusCode, ResultStatus expectedResultStatus) + [TestCase(null, ResultStatus.ERROR)] + public void TestUploadFile(HttpStatusCode? httpStatusCode, ResultStatus expectedResultStatus) { // Arrange var mockWebRequest = new Mock(); mockWebRequest.Setup(c => c.Headers).Returns(new WebHeaderCollection()); mockWebRequest.Setup(client => client.GetResponse()) - .Returns(() => - { - return MockGCSClient.CreateResponseForUploadFile(httpStatusCode); - }); + .Returns(() => MockGCSClient.CreateResponseForUploadFile(httpStatusCode)); mockWebRequest.Setup(client => client.GetRequestStream()) .Returns(() => new MemoryStream()); _client.SetCustomWebRequest(mockWebRequest.Object); @@ -257,16 +255,14 @@ public void TestUploadFile(HttpStatusCode httpStatusCode, ResultStatus expectedR [TestCase(HttpStatusCode.Forbidden, ResultStatus.NEED_RETRY)] [TestCase(HttpStatusCode.InternalServerError, ResultStatus.NEED_RETRY)] [TestCase(HttpStatusCode.ServiceUnavailable, ResultStatus.NEED_RETRY)] - public async Task TestUploadFileAsync(HttpStatusCode httpStatusCode, ResultStatus expectedResultStatus) + [TestCase(null, ResultStatus.ERROR)] + public async Task TestUploadFileAsync(HttpStatusCode? httpStatusCode, ResultStatus expectedResultStatus) { // Arrange var mockWebRequest = new Mock(); mockWebRequest.Setup(c => c.Headers).Returns(new WebHeaderCollection()); mockWebRequest.Setup(client => client.GetResponseAsync()) - .Returns(() => - { - return Task.FromResult((WebResponse)MockGCSClient.CreateResponseForUploadFile(httpStatusCode)); - }); + .Returns(() => Task.FromResult((WebResponse)MockGCSClient.CreateResponseForUploadFile(httpStatusCode))); mockWebRequest.Setup(client => client.GetRequestStreamAsync()) .Returns(() => Task.FromResult((Stream) new MemoryStream())); _client.SetCustomWebRequest(mockWebRequest.Object); @@ -301,7 +297,8 @@ private void AssertForUploadFileTests(ResultStatus expectedResultStatus) [TestCase(HttpStatusCode.Forbidden, ResultStatus.NEED_RETRY)] [TestCase(HttpStatusCode.InternalServerError, ResultStatus.NEED_RETRY)] [TestCase(HttpStatusCode.ServiceUnavailable, ResultStatus.NEED_RETRY)] - public void TestDownloadFile(HttpStatusCode httpStatusCode, ResultStatus expectedResultStatus) + [TestCase(null, ResultStatus.ERROR)] + public void TestDownloadFile(HttpStatusCode? httpStatusCode, ResultStatus expectedResultStatus) { // Arrange var mockWebRequest = new Mock(); @@ -325,7 +322,8 @@ public void TestDownloadFile(HttpStatusCode httpStatusCode, ResultStatus expecte [TestCase(HttpStatusCode.Forbidden, ResultStatus.NEED_RETRY)] [TestCase(HttpStatusCode.InternalServerError, ResultStatus.NEED_RETRY)] [TestCase(HttpStatusCode.ServiceUnavailable, ResultStatus.NEED_RETRY)] - public async Task TestDownloadFileAsync(HttpStatusCode httpStatusCode, ResultStatus expectedResultStatus) + [TestCase(null, ResultStatus.ERROR)] + public async Task TestDownloadFileAsync(HttpStatusCode? httpStatusCode, ResultStatus expectedResultStatus) { // Arrange var mockWebRequest = new Mock(); diff --git a/Snowflake.Data/Core/FileTransfer/StorageClient/SFGCSClient.cs b/Snowflake.Data/Core/FileTransfer/StorageClient/SFGCSClient.cs index 14abaad4a..9e588e921 100644 --- a/Snowflake.Data/Core/FileTransfer/StorageClient/SFGCSClient.cs +++ b/Snowflake.Data/Core/FileTransfer/StorageClient/SFGCSClient.cs @@ -349,7 +349,7 @@ public void DownloadFile(SFFileMetadata fileMetadata, string fullDstPath, int ma try { // Issue the GET request - WebRequest request = _customWebRequest == null ? FormBaseRequest(fileMetadata, "GET") : _customWebRequest; + WebRequest request = _customWebRequest == null ? FormBaseRequest(fileMetadata, "GET") : _customWebRequest; using (HttpWebResponse response = (HttpWebResponse)request.GetResponse()) { // Write to file @@ -444,7 +444,7 @@ private void HandleDownloadResponse(HttpWebResponse response, SFFileMetadata fil private SFFileMetadata HandleFileHeaderErrForPresignedUrls(WebException ex, SFFileMetadata fileMetadata) { Logger.Error("Failed to get file header for presigned url: " + ex.Message); - + HttpWebResponse response = (HttpWebResponse)ex.Response; if (response.StatusCode == HttpStatusCode.Unauthorized || response.StatusCode == HttpStatusCode.Forbidden || @@ -509,7 +509,11 @@ private SFFileMetadata HandleUploadFileErr(WebException ex, SFFileMetadata fileM fileMetadata.lastError = ex; HttpWebResponse response = (HttpWebResponse)ex.Response; - if (response.StatusCode == HttpStatusCode.BadRequest && GCS_ACCESS_TOKEN != null) + if (response is null) + { + fileMetadata.resultStatus = ResultStatus.ERROR.ToString(); + } + else if (response.StatusCode == HttpStatusCode.BadRequest && GCS_ACCESS_TOKEN != null) { fileMetadata.resultStatus = ResultStatus.RENEW_PRESIGNED_URL.ToString(); } @@ -539,7 +543,11 @@ private SFFileMetadata HandleDownloadFileErr(WebException ex, SFFileMetadata fil fileMetadata.lastError = ex; HttpWebResponse response = (HttpWebResponse)ex.Response; - if (response.StatusCode == HttpStatusCode.Unauthorized) + if (response is null) + { + fileMetadata.resultStatus = ResultStatus.ERROR.ToString(); + } + else if (response.StatusCode == HttpStatusCode.Unauthorized) { fileMetadata.resultStatus = ResultStatus.RENEW_TOKEN.ToString(); } diff --git a/Snowflake.Data/Core/SFStatement.cs b/Snowflake.Data/Core/SFStatement.cs index e84690d54..146a10130 100644 --- a/Snowflake.Data/Core/SFStatement.cs +++ b/Snowflake.Data/Core/SFStatement.cs @@ -383,11 +383,14 @@ internal async Task ExecuteAsync(int timeout, string sql, Dicti SFBindUploader uploader = new SFBindUploader(SfSession, _requestId); await uploader.UploadAsync(bindings, cancellationToken).ConfigureAwait(false); _bindStage = uploader.getStagePath(); - ClearQueryRequestId(); } catch (Exception e) { - logger.Warn("Exception encountered trying to upload binds to stage. Attaching binds in payload instead. {0}", e); + logger.Warn("Exception encountered trying to upload binds to stage. Attaching binds in payload instead. Exception: " + e.Message); + } + finally + { + ClearQueryRequestId(); } } @@ -532,13 +535,14 @@ private SFBaseResultSet ExecuteSqlOtherThanPutGet(int timeout, string sql, Dicti 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); + logger.Warn("Exception encountered trying to upload binds to stage. Attaching binds in payload instead. Exception: " + e.Message); + } + finally + { + ClearQueryRequestId(); } }