Skip to content

Commit

Permalink
SNOW-1736920 Fix bindings uploading and error handling for GCP (#1041)
Browse files Browse the repository at this point in the history
  • Loading branch information
@sfc-gh-dstempniak
sfc-gh-dstempniak authored Oct 16, 2024
1 parent 1394f53 commit b146410
Show file tree
Hide file tree
Showing 5 changed files with 48 additions and 29 deletions.
1 change: 1 addition & 0 deletions Snowflake.Data.Tests/IntegrationTests/SFBindTestIT.cs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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());
Expand Down
22 changes: 15 additions & 7 deletions Snowflake.Data.Tests/Mock/MockGCSClient.cs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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<HttpWebResponse>();

Expand All @@ -46,26 +46,34 @@ 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<HttpWebResponse>();

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);
}

return response.Object;
}

// Create a mock response for DownloadFile
static internal HttpWebResponse CreateResponseForDownloadFile(HttpStatusCode httpStatusCode)
internal static HttpWebResponse CreateResponseForDownloadFile(HttpStatusCode? httpStatusCode)
{
var response = new Mock<HttpWebResponse>();

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,
Expand All @@ -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);
}

Expand Down
22 changes: 10 additions & 12 deletions Snowflake.Data.Tests/UnitTests/SFGCSClientTest.cs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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<WebRequest>();
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);
Expand All @@ -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<WebRequest>();
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);
Expand Down Expand Up @@ -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<WebRequest>();
Expand All @@ -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<WebRequest>();
Expand Down
16 changes: 12 additions & 4 deletions Snowflake.Data/Core/FileTransfer/StorageClient/SFGCSClient.cs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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 ||
Expand Down Expand Up @@ -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();
}
Expand Down Expand Up @@ -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();
}
Expand Down
16 changes: 10 additions & 6 deletions Snowflake.Data/Core/SFStatement.cs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -383,11 +383,14 @@ internal async Task<SFBaseResultSet> 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();
}
}

Expand Down Expand Up @@ -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();
}
}

Expand Down

0 comments on commit b146410

Please sign in to comment.