From 4a84f76103b50e69ea18a546335a9ccd2a21bae8 Mon Sep 17 00:00:00 2001 From: sfc-gh-ext-simba-lf Date: Fri, 8 Sep 2023 14:20:19 -0700 Subject: [PATCH 1/4] SNOW-723810: Add coverage for SFRemoteStorageUtil --- .../Mock/MockRemoteStorageClient.cs | 147 +++++ .../UnitTests/SFRemoteStorageClientTest.cs | 572 ++++++++++++++++++ .../StorageClient/SFRemoteStorageUtil.cs | 33 +- 3 files changed, 739 insertions(+), 13 deletions(-) create mode 100644 Snowflake.Data.Tests/Mock/MockRemoteStorageClient.cs create mode 100644 Snowflake.Data.Tests/UnitTests/SFRemoteStorageClientTest.cs diff --git a/Snowflake.Data.Tests/Mock/MockRemoteStorageClient.cs b/Snowflake.Data.Tests/Mock/MockRemoteStorageClient.cs new file mode 100644 index 000000000..3992283c4 --- /dev/null +++ b/Snowflake.Data.Tests/Mock/MockRemoteStorageClient.cs @@ -0,0 +1,147 @@ +/* + * Copyright (c) 2012-2023 Snowflake Computing Inc. All rights reserved. + */ + +using Moq; +using Snowflake.Data.Core.FileTransfer.StorageClient; +using System.IO; +using System.Net; +using System.Text; + +namespace Snowflake.Data.Tests.Mock +{ + + class MockRemoteStorageClient + { + // Mock data for downloaded file + internal const string FileContent = "RemoteStorageClientTest"; + + // Mock content length + internal const int ContentLength = 9999; + + // Mock error message + internal const string ErrorMessage = "Mock GCS Remote Storage Error"; + + // Variables for the encryption data + static Stream s_encryptedStream = null; + static string s_encryptedStreamIV = null; + static string s_encryptedStreamKey = null; + + static internal void SetEncryptionData(Stream stream, string iv, string key) + { + s_encryptedStream = stream; + s_encryptedStreamIV = iv; + s_encryptedStreamKey = key; + } + + // Sets up the mock sequence when Remote Storage uploads a file + // 1. Get the file header + // 2. Upload the file + static internal HttpWebResponse SequenceResponseForUploadFile(ref bool firstRequest, HttpStatusCode statusCode, HttpStatusCode statusCodeAfterRetry) + { + try + { + // For the first call to GetResponse(), return the file header response + // For the second call to GetResponse(), return the upload response + var response = firstRequest ? + CreateResponseForFileHeader(statusCode) : + CreateResponseForUploadFile(statusCodeAfterRetry); + + return response; + } + finally + { + // If an exception is thrown, the value will still be changed + firstRequest = false; + } + } + + // Create a mock response for GetFileHeader + static internal HttpWebResponse CreateResponseForFileHeader(HttpStatusCode httpStatusCode) + { + var response = new Mock(); + + if (httpStatusCode == HttpStatusCode.OK) + { + response.Setup(c => c.Headers).Returns(new WebHeaderCollection()); + response.Object.Headers.Add("content-length", MockGCSClient.ContentLength.ToString()); + response.Object.Headers.Add(SFGCSClient.GCS_METADATA_SFC_DIGEST, MockGCSClient.SFCDigest); + } + else + { + response.SetupGet(c => c.StatusCode) + .Returns(httpStatusCode); + throw new WebException(ErrorMessage, null, 0, response.Object); + } + + return response.Object; + } + + // Create a mock response for UploadFile + static internal HttpWebResponse CreateResponseForUploadFile(HttpStatusCode httpStatusCode) + { + var response = new Mock(); + + if (httpStatusCode != HttpStatusCode.OK) + { + response.SetupGet(c => c.StatusCode) + .Returns(httpStatusCode); + throw new WebException(ErrorMessage, null, 0, response.Object); + } + + return response.Object; + } + + // Create a mock response for DownloadFile + static internal HttpWebResponse CreateResponseForDownloadFile(HttpStatusCode httpStatusCode) + { + var response = new Mock(); + + if (httpStatusCode == HttpStatusCode.OK) + { + response.Setup(c => c.Headers).Returns(new WebHeaderCollection()); + + // For downloads with encryption material + if (s_encryptedStream != null) + { + // Set the position to 0 and return the encrypted stream + s_encryptedStream.Position = 0; + response.Setup(c => c.GetResponseStream()).Returns(s_encryptedStream); + + // Set the iv and key to the ones used for encrypting the stream + response.Object.Headers.Add(SFGCSClient.GCS_METADATA_ENCRYPTIONDATAPROP, + "{" + + $"\"ContentEncryptionIV\": \"{s_encryptedStreamIV}\", " + + $"\"WrappedContentKey\": {{\"EncryptedKey\":\"{s_encryptedStreamKey}\"}}" + + "}"); + + s_encryptedStreamIV = null; + s_encryptedStreamKey = null; + s_encryptedStream = null; + } + else // For unencrypted downloads + { + response.Setup(c => c.GetResponseStream()).Returns(new MemoryStream(Encoding.ASCII.GetBytes(FileContent))); + response.Object.Headers.Add(SFGCSClient.GCS_METADATA_ENCRYPTIONDATAPROP, + "{" + + $"\"ContentEncryptionIV\": \"{MockGCSClient.GcsIV}\", " + + $"\"WrappedContentKey\": {{\"EncryptedKey\":\"{MockGCSClient.GcsKey}\"}}" + + "}"); + } + + response.Object.Headers.Add(SFGCSClient.GCS_METADATA_MATDESC_KEY, MockGCSClient.GcsMatdesc); + response.Object.Headers.Add(SFGCSClient.GCS_METADATA_SFC_DIGEST, MockGCSClient.SFCDigest); + response.Object.Headers.Add(SFGCSClient.GCS_FILE_HEADER_CONTENT_LENGTH, MockGCSClient.ContentLength.ToString()); + } + else + { + response.SetupGet(c => c.StatusCode) + .Returns(httpStatusCode); + throw new WebException(ErrorMessage, null, 0, response.Object); + } + + return response.Object; + } + } +} + diff --git a/Snowflake.Data.Tests/UnitTests/SFRemoteStorageClientTest.cs b/Snowflake.Data.Tests/UnitTests/SFRemoteStorageClientTest.cs new file mode 100644 index 000000000..4be002666 --- /dev/null +++ b/Snowflake.Data.Tests/UnitTests/SFRemoteStorageClientTest.cs @@ -0,0 +1,572 @@ +/* + * Copyright (c) 2012-2023 Snowflake Computing Inc. All rights reserved. + */ + +namespace Snowflake.Data.Tests.UnitTests +{ + using NUnit.Framework; + using Snowflake.Data.Core; + using Snowflake.Data.Core.FileTransfer.StorageClient; + using Snowflake.Data.Core.FileTransfer; + using System.Collections.Generic; + using System; + using Snowflake.Data.Tests.Mock; + using System.Threading.Tasks; + using System.Threading; + using System.IO; + using System.Text; + using System.Net; + using Moq; + + [TestFixture] + class SFRemoteStorageClientTest : SFBaseTest + { + // Mock data for file metadata + const string EndPoint = "mockEndPoint.com"; + + const string LocationStage = "mock-customer-stage"; + const string LocationId = "mock-id"; + const string LocationTables = "tables"; + const string LocationKey = "mock-key"; + const string LocationPath = LocationTables + "/" + LocationKey + "/"; + const string Location = LocationStage + "/" + LocationId + "/" + LocationPath; + + const string Region = "us-west-2"; + + const string StorageAccount = "mockStorageAccount"; + + const string RealSourceFilePath = "realSrcFilePath.txt"; + + PutGetEncryptionMaterial EncryptionMaterial = new PutGetEncryptionMaterial() + { + queryId = "MOCK/QUERY/ID/==", + queryStageMasterKey = "MOCKQUERYSTAGEMASTERKE==", + smkId = 9999 + }; + + // Mock unsupported stage type + const string UnsupportedStageType = "UNSUPPORTED"; + + // Settings for mock client + const int Parallel = 0; + + // File name for mock test files + const string LocalLocation = "./"; + const string DownloadFileName = "mockDownloadFileName.txt"; + const string DownloadFilePath = LocalLocation + DownloadFileName; + + // Mock upload file size + const int UploadFileSize = 9999; + const int DestFileSizeWhenFileAlreadyExists = 0; + + // Mock client and metadata + SFGCSClient _client; + PutGetResponseData _responseData; + SFFileMetadata _fileMetadata; + + // Token for async tests + CancellationToken _cancellationToken; + + // Flags for non-async and async mock methods + const bool NotAsync = false; + const bool IsAsync = true; + + [SetUp] + public void BeforeTest() + { + _fileMetadata = new SFFileMetadata() + { + destFileName = DownloadFileName, + localLocation = LocalLocation, + MaxBytesInMemory = 1024, + memoryStream = new MemoryStream(Encoding.UTF8.GetBytes(MockRemoteStorageClient.FileContent)), + parallel = Parallel, + realSrcFilePath = RealSourceFilePath, + stageInfo = new PutGetStageInfo() + { + endPoint = EndPoint, + isClientSideEncrypted = true, + location = Location, + locationType = SFRemoteStorageUtil.GCS_FS, + path = LocationPath, + presignedUrl = null, + region = Region, + stageCredentials = new Dictionary() + { + {"AWS_KEY_ID", "MOCK_AWS_KEY_ID"}, + {"AWS_SECRET_KEY", "MOCK_AWS_SECRET_KEY"}, + {"AWS_TOKEN", "MOCK_AWS_TOKEN"}, + {"AWS_ID", "MOCK_AWS_ID"}, + {"AWS_KEY", "MOCK_AWS_KEY"}, + {"AZURE_SAS_TOKEN", "MOCK_AZURE_SAS_TOKEN"}, + {"GCS_ACCESS_TOKEN", "MOCK_GCS_ACCESS_TOKEN"} + }, + storageAccount = StorageAccount + }, + uploadSize = UploadFileSize + }; + + _responseData = new PutGetResponseData() + { + autoCompress = true, + stageInfo = _fileMetadata.stageInfo + }; + + // Set the mock GCS client to use + _client = new SFGCSClient(_fileMetadata.stageInfo); + _fileMetadata.client = _client; + + _cancellationToken = new CancellationToken(); + } + + [TearDown] + public void AfterTest() + { + // Delete temporary files from upload + if (File.Exists(_fileMetadata.realSrcFilePath)) + { + File.Delete(_fileMetadata.realSrcFilePath); + } + + // Delete temporary files from download + if (File.Exists(DownloadFilePath)) + { + File.Delete(DownloadFilePath); + } + } + + [Test] + [Ignore("RemoteStorageClientTest")] + public void RemoteStorageClientTestDone() + { + // Do nothing; + } + + [Test] + [TestCase(SFRemoteStorageUtil.LOCAL_FS)] + [TestCase(SFRemoteStorageUtil.S3_FS)] + [TestCase(SFRemoteStorageUtil.AZURE_FS)] + [TestCase(SFRemoteStorageUtil.GCS_FS)] + [TestCase(UnsupportedStageType)] // Any other stage type should return null + public void TestGetRemoteStorageClient(string stageType) + { + _responseData.stageInfo.locationType = stageType; + + if (stageType == SFRemoteStorageUtil.LOCAL_FS) + { + Assert.Throws(() => SFRemoteStorageUtil.GetRemoteStorage(_responseData)); + } + else + { + ISFRemoteStorageClient client = SFRemoteStorageUtil.GetRemoteStorage(_responseData); + + if (stageType == SFRemoteStorageUtil.S3_FS) + { + Assert.IsInstanceOf(client); + } + else if (stageType == SFRemoteStorageUtil.AZURE_FS) + { + Assert.IsInstanceOf(client); + } + else if (stageType == SFRemoteStorageUtil.GCS_FS) + { + Assert.IsInstanceOf(client); + } + else + { + Assert.IsNull(client); + } + } + } + + [Test] + [TestCase(false, false)] + [TestCase(false, true)] + [TestCase(true, false)] + [TestCase(true, true)] + public void TestGetFileBytes(bool containsEncryptionMaterial, bool useMemoryStream) + { + // Arrange + var mockWebRequest = new Mock(); + mockWebRequest.Setup(client => client.GetResponse()) + .Returns(() => + { + return MockRemoteStorageClient.CreateResponseForFileHeader(HttpStatusCode.OK); + }); + _client.SetCustomWebRequest(mockWebRequest.Object); + + // Add encryption material to the file metadata + _fileMetadata.encryptionMaterial = containsEncryptionMaterial ? EncryptionMaterial : null; + + // Use file for upload + if (!useMemoryStream) + { + File.WriteAllText(_fileMetadata.realSrcFilePath, MockRemoteStorageClient.FileContent); + _fileMetadata.memoryStream = null; + } + + // Act + SFRemoteStorageUtil.UploadOneFile(_fileMetadata); + + // Assert + Assert.AreEqual(DestFileSizeWhenFileAlreadyExists, _fileMetadata.destFileSize); + Assert.AreEqual(ResultStatus.SKIPPED.ToString(), _fileMetadata.resultStatus); + } + + private Mock CreateBaseMockClient() + { + // Setup the mock GCS client for remote storage tests + var mockWebRequest = new Mock(); + mockWebRequest.Setup(c => c.Headers).Returns(new WebHeaderCollection()); + + return mockWebRequest; + } + + private void SetUpMockClientForUpload(HttpStatusCode statusCode, HttpStatusCode statusCodeAfterRetry, bool isAsync) + { + // Setup the mock GCS client for remote storage tests + var mockWebRequest = CreateBaseMockClient(); + bool firstRequest = true; + + if (isAsync) + { + mockWebRequest.Setup(client => client.GetRequestStreamAsync()) + .Returns(() => Task.FromResult(new MemoryStream())); + mockWebRequest.Setup(client => client.GetResponseAsync()) + .Returns(() => + { + return Task.FromResult(MockRemoteStorageClient.SequenceResponseForUploadFile(ref firstRequest, statusCode, statusCodeAfterRetry)); + }); + } + else + { + mockWebRequest.Setup(client => client.GetRequestStream()) + .Returns(() => new MemoryStream()); + mockWebRequest.Setup(client => client.GetResponse()) + .Returns(() => + { + return MockRemoteStorageClient.SequenceResponseForUploadFile(ref firstRequest, statusCode, statusCodeAfterRetry); + }); + } + + _client.SetCustomWebRequest(mockWebRequest.Object); + } + + [Test] + [TestCase(HttpStatusCode.NotFound, HttpStatusCode.OK, ResultStatus.UPLOADED)] + public void TestUploadOneFileWithRetry(HttpStatusCode httpStatusCode, HttpStatusCode httpStatusCodeAfterRetry, ResultStatus expectedResultStatus) + { + // Arrange + SetUpMockClientForUpload(httpStatusCode, httpStatusCodeAfterRetry, NotAsync); + + // Act + SFRemoteStorageUtil.UploadOneFileWithRetry(_fileMetadata); + + // Assert + Assert.AreEqual(expectedResultStatus.ToString(), _fileMetadata.resultStatus); + } + + [Test] + [TestCase(HttpStatusCode.NotFound, HttpStatusCode.OK, ResultStatus.UPLOADED)] + public async Task TestUploadOneFileAsyncWithRetry(HttpStatusCode httpStatusCode, HttpStatusCode httpStatusCodeAfterRetry, ResultStatus expectedResultStatus) + { + // Arrange + SetUpMockClientForUpload(httpStatusCode, httpStatusCodeAfterRetry, IsAsync); + + // Act + await SFRemoteStorageUtil.UploadOneFileWithRetryAsync(_fileMetadata, _cancellationToken).ConfigureAwait(false); + + // Assert + Assert.AreEqual(expectedResultStatus.ToString(), _fileMetadata.resultStatus); + } + + [Test] + [TestCase(HttpStatusCode.OK, null, ResultStatus.SKIPPED)] + [TestCase(HttpStatusCode.NotFound, HttpStatusCode.OK, ResultStatus.UPLOADED)] + [TestCase(HttpStatusCode.NotFound, HttpStatusCode.BadRequest, ResultStatus.RENEW_PRESIGNED_URL)] + [TestCase(HttpStatusCode.NotFound, HttpStatusCode.Unauthorized, ResultStatus.RENEW_TOKEN)] + public void TestUploadOneFile(HttpStatusCode httpStatusCode, HttpStatusCode httpStatusCodeAfterRetry, ResultStatus expectedResultStatus) + { + // Arrange + SetUpMockClientForUpload(httpStatusCode, httpStatusCodeAfterRetry, NotAsync); + + // Act + SFRemoteStorageUtil.UploadOneFile(_fileMetadata); + + // Assert + if (expectedResultStatus == ResultStatus.SKIPPED) + { + Assert.AreEqual(DestFileSizeWhenFileAlreadyExists, _fileMetadata.destFileSize); + } + Assert.AreEqual(expectedResultStatus.ToString(), _fileMetadata.resultStatus); + } + + [Test] + [TestCase(HttpStatusCode.OK, null, ResultStatus.SKIPPED)] + [TestCase(HttpStatusCode.NotFound, HttpStatusCode.OK, ResultStatus.UPLOADED)] + [TestCase(HttpStatusCode.NotFound, HttpStatusCode.BadRequest, ResultStatus.RENEW_PRESIGNED_URL)] + [TestCase(HttpStatusCode.NotFound, HttpStatusCode.Unauthorized, ResultStatus.RENEW_TOKEN)] + public async Task TestUploadOneFileAsync(HttpStatusCode httpStatusCode, HttpStatusCode httpStatusCodeAfterRetry, ResultStatus expectedResultStatus) + { + // Arrange + SetUpMockClientForUpload(httpStatusCode, httpStatusCodeAfterRetry, IsAsync); + + // Act + await SFRemoteStorageUtil.UploadOneFileAsync(_fileMetadata, _cancellationToken).ConfigureAwait(false); + + // Assert + if (expectedResultStatus == ResultStatus.SKIPPED) + { + Assert.AreEqual(DestFileSizeWhenFileAlreadyExists, _fileMetadata.destFileSize); + } + Assert.AreEqual(expectedResultStatus.ToString(), _fileMetadata.resultStatus); + } + + [Test] + [TestCase(HttpStatusCode.NotFound, HttpStatusCode.Forbidden, ResultStatus.NEED_RETRY)] + [TestCase(HttpStatusCode.NotFound, HttpStatusCode.InternalServerError, ResultStatus.NEED_RETRY)] + [TestCase(HttpStatusCode.NotFound, HttpStatusCode.ServiceUnavailable, ResultStatus.NEED_RETRY)] + public void TestUploadOneFileThrowsForRetryErrors(HttpStatusCode httpStatusCode, HttpStatusCode httpStatusCodeAfterRetry, ResultStatus expectedResultStatus) + { + // Arrange + SetUpMockClientForUpload(httpStatusCode, httpStatusCodeAfterRetry, NotAsync); + + // Act + Exception ex = Assert.Throws(() => SFRemoteStorageUtil.UploadOneFile(_fileMetadata)); + + // Assert + Assert.That(ex.Message, Does.Match(MockRemoteStorageClient.ErrorMessage)); + Assert.AreEqual(expectedResultStatus.ToString(), _fileMetadata.resultStatus); + } + + [Test] + [TestCase(HttpStatusCode.NotFound, HttpStatusCode.Forbidden, ResultStatus.NEED_RETRY)] + [TestCase(HttpStatusCode.NotFound, HttpStatusCode.InternalServerError, ResultStatus.NEED_RETRY)] + [TestCase(HttpStatusCode.NotFound, HttpStatusCode.ServiceUnavailable, ResultStatus.NEED_RETRY)] + public void TestUploadOneFileAsyncThrowsForRetryErrors(HttpStatusCode httpStatusCode, HttpStatusCode httpStatusCodeAfterRetry, ResultStatus expectedResultStatus) + { + // Arrange + SetUpMockClientForUpload(httpStatusCode, httpStatusCodeAfterRetry, IsAsync); + + // Act + Exception ex = Assert.ThrowsAsync(async () => await SFRemoteStorageUtil.UploadOneFileAsync(_fileMetadata, _cancellationToken).ConfigureAwait(false)); + + // Assert + Assert.That(ex.Message, Does.Match(MockRemoteStorageClient.ErrorMessage)); + Assert.AreEqual(expectedResultStatus.ToString(), _fileMetadata.resultStatus); + } + + + [Test] + [TestCase(HttpStatusCode.NotFound, null, ResultStatus.ERROR)] + public void TestUploadOneFileThrowsForUnknownErrors(HttpStatusCode httpStatusCode, HttpStatusCode httpStatusCodeAfterRetry, ResultStatus expectedResultStatus) + { + // Arrange + SetUpMockClientForUpload(httpStatusCode, httpStatusCodeAfterRetry, NotAsync); + + // Act + Exception ex = Assert.Throws(() => SFRemoteStorageUtil.UploadOneFile(_fileMetadata)); + + // Assert + Assert.That(ex.Message, Does.Match($"Unknown Error in uploading a file: .*")); + Assert.AreEqual(expectedResultStatus.ToString(), _fileMetadata.resultStatus); + } + + [Test] + [TestCase(HttpStatusCode.NotFound, null, ResultStatus.ERROR)] + public void TestUploadOneFileAsyncThrowsForUnknownErrors(HttpStatusCode httpStatusCode, HttpStatusCode httpStatusCodeAfterRetry, ResultStatus expectedResultStatus) + { + // Arrange + SetUpMockClientForUpload(httpStatusCode, httpStatusCodeAfterRetry, IsAsync); + + // Act + Exception ex = Assert.ThrowsAsync(async () => await SFRemoteStorageUtil.UploadOneFileAsync(_fileMetadata, _cancellationToken).ConfigureAwait(false)); + + // Assert + Assert.That(ex.Message, Does.Match($"Unknown Error in uploading a file: .*")); + Assert.AreEqual(expectedResultStatus.ToString(), _fileMetadata.resultStatus); + } + + private void SetUpMockClientForDownload(HttpStatusCode statusCode, bool isAsync) + { + // Setup the mock GCS client for remote storage tests + var mockWebRequest = CreateBaseMockClient(); + if (isAsync) + { + mockWebRequest.Setup(client => client.GetResponseAsync()) + .Returns(() => + { + return Task.FromResult(MockRemoteStorageClient.CreateResponseForDownloadFile(statusCode)); + }); + + } + else + { + mockWebRequest.Setup(client => client.GetResponse()) + .Returns(() => + { + return MockRemoteStorageClient.CreateResponseForDownloadFile(statusCode); + }); + } + + _client.SetCustomWebRequest(mockWebRequest.Object); + } + + [Test] + [TestCase(HttpStatusCode.OK, ResultStatus.DOWNLOADED)] + [TestCase(HttpStatusCode.Unauthorized, ResultStatus.RENEW_TOKEN)] + public void TestDownloadOneFile(HttpStatusCode httpStatusCode, ResultStatus expectedResultStatus) + { + // Arrange + SetUpMockClientForDownload(httpStatusCode, NotAsync); + + // Act + SFRemoteStorageUtil.DownloadOneFile(_fileMetadata); + + // Assert + if (expectedResultStatus == ResultStatus.DOWNLOADED) + { + string text = File.ReadAllText(_fileMetadata.localLocation + _fileMetadata.destFileName); + Assert.AreEqual(MockRemoteStorageClient.FileContent, text); + } + Assert.AreEqual(expectedResultStatus.ToString(), _fileMetadata.resultStatus); + } + + [Test] + [TestCase(HttpStatusCode.OK, ResultStatus.DOWNLOADED)] + [TestCase(HttpStatusCode.Unauthorized, ResultStatus.RENEW_TOKEN)] + public async Task TestDownloadOneFileAsync(HttpStatusCode httpStatusCode, ResultStatus expectedResultStatus) + { + // Arrange + SetUpMockClientForDownload(httpStatusCode, IsAsync); + + // Act + await SFRemoteStorageUtil.DownloadOneFileAsync(_fileMetadata, _cancellationToken).ConfigureAwait(false); + + // Assert + if (expectedResultStatus == ResultStatus.DOWNLOADED) + { + string text = File.ReadAllText(_fileMetadata.localLocation + _fileMetadata.destFileName); + Assert.AreEqual(MockRemoteStorageClient.FileContent, text); + } + Assert.AreEqual(expectedResultStatus.ToString(), _fileMetadata.resultStatus); + } + + [Test] + [TestCase(HttpStatusCode.Forbidden, ResultStatus.NEED_RETRY)] + [TestCase(HttpStatusCode.InternalServerError, ResultStatus.NEED_RETRY)] + [TestCase(HttpStatusCode.ServiceUnavailable, ResultStatus.NEED_RETRY)] + public void TestDownloadOneFileThrowsForRetryErrors(HttpStatusCode httpStatusCode, ResultStatus expectedResultStatus) + { + // Arrange + SetUpMockClientForDownload(httpStatusCode, NotAsync); + + // Act + Exception ex = Assert.Throws(() => SFRemoteStorageUtil.DownloadOneFile(_fileMetadata)); + + // Assert + Assert.That(ex.Message, Does.Match(MockRemoteStorageClient.ErrorMessage)); + Assert.AreEqual(expectedResultStatus.ToString(), _fileMetadata.resultStatus); + } + + [Test] + [TestCase(HttpStatusCode.Forbidden, ResultStatus.NEED_RETRY)] + [TestCase(HttpStatusCode.InternalServerError, ResultStatus.NEED_RETRY)] + [TestCase(HttpStatusCode.ServiceUnavailable, ResultStatus.NEED_RETRY)] + public void TestDownloadOneFileAsyncThrowsForRetryErrors(HttpStatusCode httpStatusCode, ResultStatus expectedResultStatus) + { + // Arrange + SetUpMockClientForDownload(httpStatusCode, IsAsync); + + // Act + Exception ex = Assert.ThrowsAsync(async () => await SFRemoteStorageUtil.DownloadOneFileAsync(_fileMetadata, _cancellationToken).ConfigureAwait(false)); + + // Assert + Assert.That(ex.Message, Does.Match(MockRemoteStorageClient.ErrorMessage)); + Assert.AreEqual(expectedResultStatus.ToString(), _fileMetadata.resultStatus); + } + + [Test] + [TestCase(HttpStatusCode.NotFound)] + public void TestDownloadOneFileThrowsForUnknownErrors(HttpStatusCode httpStatusCode) + { + // Arrange + SetUpMockClientForDownload(httpStatusCode, NotAsync); + + // Act + Exception ex = Assert.Throws(() => SFRemoteStorageUtil.DownloadOneFile(_fileMetadata)); + + // Assert + Assert.That(ex.Message, Does.Match($"Unknown Error in downloading a file: .*")); + Assert.IsNull(_fileMetadata.resultStatus); + } + + [Test] + [TestCase(HttpStatusCode.NotFound)] + public void TestDownloadOneFileAsyncThrowsForUnknownErrors(HttpStatusCode httpStatusCode) + { + // Arrange + SetUpMockClientForDownload(httpStatusCode, IsAsync); + + // Act + Exception ex = Assert.ThrowsAsync(async () => await SFRemoteStorageUtil.DownloadOneFileAsync(_fileMetadata, _cancellationToken).ConfigureAwait(false)); + + // Assert + Assert.That(ex.Message, Does.Match($"Unknown Error in downloading a file: .*")); + Assert.IsNull(_fileMetadata.resultStatus); + } + + private void SetUpMockEncryptedFileForDownload() + { + _fileMetadata.encryptionMaterial = EncryptionMaterial; + + // Write file to encrypt + File.WriteAllText(DownloadFilePath, MockRemoteStorageClient.FileContent); + + // Get encrypted stream from file + SFEncryptionMetadata encryptionMetadata = new SFEncryptionMetadata(); + Stream stream = EncryptionProvider.EncryptFile( + DownloadFilePath, + _fileMetadata.encryptionMaterial, + encryptionMetadata, + FileTransferConfiguration.FromFileMetadata(_fileMetadata)); + + // Set up the stream and metadata for decryption + MockRemoteStorageClient.SetEncryptionData(stream, encryptionMetadata.iv, encryptionMetadata.key); + } + + [Test] + [TestCase(HttpStatusCode.OK, ResultStatus.DOWNLOADED)] + public void TestDownloadOneFileWithEncryptionMaterial(HttpStatusCode httpStatusCode, ResultStatus expectedResultStatus) + { + // Arrange + SetUpMockEncryptedFileForDownload(); + SetUpMockClientForDownload(httpStatusCode, NotAsync); + + // Act + SFRemoteStorageUtil.DownloadOneFile(_fileMetadata); + + // Assert + string text = File.ReadAllText(DownloadFilePath); + Assert.AreEqual(MockRemoteStorageClient.FileContent, text); + Assert.AreEqual(expectedResultStatus.ToString(), _fileMetadata.resultStatus); + } + + [Test] + [TestCase(HttpStatusCode.OK, ResultStatus.DOWNLOADED)] + public async Task TestDownloadOneFileAsyncWithEncryptionMaterial(HttpStatusCode httpStatusCode, ResultStatus expectedResultStatus) + { + // Arrange + SetUpMockEncryptedFileForDownload(); + SetUpMockClientForDownload(httpStatusCode, IsAsync); + + // Act + await SFRemoteStorageUtil.DownloadOneFileAsync(_fileMetadata, _cancellationToken).ConfigureAwait(false); + + // Assert + string text = File.ReadAllText(DownloadFilePath); + Assert.AreEqual(MockRemoteStorageClient.FileContent, text); + Assert.AreEqual(expectedResultStatus.ToString(), _fileMetadata.resultStatus); + } + } +} diff --git a/Snowflake.Data/Core/FileTransfer/StorageClient/SFRemoteStorageUtil.cs b/Snowflake.Data/Core/FileTransfer/StorageClient/SFRemoteStorageUtil.cs index 45489954f..c5b358641 100644 --- a/Snowflake.Data/Core/FileTransfer/StorageClient/SFRemoteStorageUtil.cs +++ b/Snowflake.Data/Core/FileTransfer/StorageClient/SFRemoteStorageUtil.cs @@ -377,6 +377,11 @@ internal static void DownloadOneFile(SFFileMetadata fileMetadata) fileMetadata.destFileSize = fileInfo.Length; return; } + else if (fileMetadata.resultStatus == ResultStatus.RENEW_TOKEN.ToString() || + fileMetadata.resultStatus == ResultStatus.RENEW_PRESIGNED_URL.ToString()) + { + return; + } else { HandleDownloadFileErr(ref fileMetadata, ref maxConcurrency, ref lastErr, retry, maxRetry); @@ -410,12 +415,6 @@ internal static async Task DownloadOneFileAsync(SFFileMetadata fileMetadata, Can } ISFRemoteStorageClient client = fileMetadata.client; - FileHeader fileHeader = await client.GetFileHeaderAsync(fileMetadata, cancellationToken).ConfigureAwait(false); - - if (fileHeader != null) - { - fileMetadata.srcFileSize = fileHeader.contentLength; - } int maxConcurrency = fileMetadata.parallel; Exception lastErr = null; @@ -440,15 +439,23 @@ internal static async Task DownloadOneFileAsync(SFFileMetadata fileMetadata, Can * One example of this is the utils that use presigned url * for upload / download and not the storage client library. **/ + FileHeader fileHeader = null; if (fileMetadata.presignedUrl != null) { fileHeader = await client.GetFileHeaderAsync(fileMetadata, cancellationToken).ConfigureAwait(false); + + if (fileHeader != null) + { + fileMetadata.srcFileSize = fileHeader.contentLength; + } } + SFEncryptionMetadata encryptionMetadata = fileHeader != null ? fileHeader.encryptionMetadata : fileMetadata.encryptionMetadata; + string tmpDstName = EncryptionProvider.DecryptFile( fullDstPath, fileMetadata.encryptionMaterial, - fileHeader.encryptionMetadata, + encryptionMetadata, FileTransferConfiguration.FromFileMetadata(fileMetadata)); File.Delete(fullDstPath); @@ -464,6 +471,11 @@ internal static async Task DownloadOneFileAsync(SFFileMetadata fileMetadata, Can fileMetadata.destFileSize = fileInfo.Length; return; } + else if (fileMetadata.resultStatus == ResultStatus.RENEW_TOKEN.ToString() || + fileMetadata.resultStatus == ResultStatus.RENEW_PRESIGNED_URL.ToString()) + { + return; + } else { HandleDownloadFileErr(ref fileMetadata, ref maxConcurrency, ref lastErr, retry, maxRetry); @@ -490,12 +502,7 @@ internal static async Task DownloadOneFileAsync(SFFileMetadata fileMetadata, Can /// The max retry private static void HandleDownloadFileErr(ref SFFileMetadata fileMetadata, ref int maxConcurrency, ref Exception lastErr, int retry, int maxRetry) { - if (fileMetadata.resultStatus == ResultStatus.RENEW_TOKEN.ToString() || - fileMetadata.resultStatus == ResultStatus.RENEW_PRESIGNED_URL.ToString()) - { - return; - } - else if (fileMetadata.resultStatus == ResultStatus.NEED_RETRY_WITH_LOWER_CONCURRENCY.ToString()) + if (fileMetadata.resultStatus == ResultStatus.NEED_RETRY_WITH_LOWER_CONCURRENCY.ToString()) { lastErr = fileMetadata.lastError; // Failed to download file, retrying with max concurrency From 89d16f43f82cc999b3cef121133d4478c24f39aa Mon Sep 17 00:00:00 2001 From: sfc-gh-ext-simba-lf Date: Fri, 8 Sep 2023 14:25:20 -0700 Subject: [PATCH 2/4] SNOW-723810: Add coverage for SFRemoteStorageUtil --- .../Core/FileTransfer/StorageClient/SFRemoteStorageUtil.cs | 5 ----- 1 file changed, 5 deletions(-) diff --git a/Snowflake.Data/Core/FileTransfer/StorageClient/SFRemoteStorageUtil.cs b/Snowflake.Data/Core/FileTransfer/StorageClient/SFRemoteStorageUtil.cs index c5b358641..92e8dd467 100644 --- a/Snowflake.Data/Core/FileTransfer/StorageClient/SFRemoteStorageUtil.cs +++ b/Snowflake.Data/Core/FileTransfer/StorageClient/SFRemoteStorageUtil.cs @@ -443,11 +443,6 @@ internal static async Task DownloadOneFileAsync(SFFileMetadata fileMetadata, Can if (fileMetadata.presignedUrl != null) { fileHeader = await client.GetFileHeaderAsync(fileMetadata, cancellationToken).ConfigureAwait(false); - - if (fileHeader != null) - { - fileMetadata.srcFileSize = fileHeader.contentLength; - } } SFEncryptionMetadata encryptionMetadata = fileHeader != null ? fileHeader.encryptionMetadata : fileMetadata.encryptionMetadata; From e9b3f01eb6acff034a5af83352da3586074f9141 Mon Sep 17 00:00:00 2001 From: sfc-gh-ext-simba-lf Date: Fri, 15 Sep 2023 15:09:53 -0700 Subject: [PATCH 3/4] SNOW-723810: Refactor to work with parallelization --- .../Mock/MockRemoteStorageClient.cs | 29 ++++++++++--------- .../UnitTests/SFRemoteStorageClientTest.cs | 28 +++++++++--------- 2 files changed, 30 insertions(+), 27 deletions(-) diff --git a/Snowflake.Data.Tests/Mock/MockRemoteStorageClient.cs b/Snowflake.Data.Tests/Mock/MockRemoteStorageClient.cs index 3992283c4..c93173d58 100644 --- a/Snowflake.Data.Tests/Mock/MockRemoteStorageClient.cs +++ b/Snowflake.Data.Tests/Mock/MockRemoteStorageClient.cs @@ -4,6 +4,7 @@ using Moq; using Snowflake.Data.Core.FileTransfer.StorageClient; +using System; using System.IO; using System.Net; using System.Text; @@ -23,15 +24,15 @@ class MockRemoteStorageClient internal const string ErrorMessage = "Mock GCS Remote Storage Error"; // Variables for the encryption data - static Stream s_encryptedStream = null; - static string s_encryptedStreamIV = null; - static string s_encryptedStreamKey = null; + [ThreadStatic] static Stream t_encryptedStream = null; + [ThreadStatic] static string t_encryptedStreamIV = null; + [ThreadStatic] static string t_encryptedStreamKey = null; static internal void SetEncryptionData(Stream stream, string iv, string key) { - s_encryptedStream = stream; - s_encryptedStreamIV = iv; - s_encryptedStreamKey = key; + t_encryptedStream = stream; + t_encryptedStreamIV = iv; + t_encryptedStreamKey = key; } // Sets up the mock sequence when Remote Storage uploads a file @@ -102,22 +103,22 @@ static internal HttpWebResponse CreateResponseForDownloadFile(HttpStatusCode htt response.Setup(c => c.Headers).Returns(new WebHeaderCollection()); // For downloads with encryption material - if (s_encryptedStream != null) + if (t_encryptedStream != null) { // Set the position to 0 and return the encrypted stream - s_encryptedStream.Position = 0; - response.Setup(c => c.GetResponseStream()).Returns(s_encryptedStream); + t_encryptedStream.Position = 0; + response.Setup(c => c.GetResponseStream()).Returns(t_encryptedStream); // Set the iv and key to the ones used for encrypting the stream response.Object.Headers.Add(SFGCSClient.GCS_METADATA_ENCRYPTIONDATAPROP, "{" + - $"\"ContentEncryptionIV\": \"{s_encryptedStreamIV}\", " + - $"\"WrappedContentKey\": {{\"EncryptedKey\":\"{s_encryptedStreamKey}\"}}" + + $"\"ContentEncryptionIV\": \"{t_encryptedStreamIV}\", " + + $"\"WrappedContentKey\": {{\"EncryptedKey\":\"{t_encryptedStreamKey}\"}}" + "}"); - s_encryptedStreamIV = null; - s_encryptedStreamKey = null; - s_encryptedStream = null; + t_encryptedStreamIV = null; + t_encryptedStreamKey = null; + t_encryptedStream = null; } else // For unencrypted downloads { diff --git a/Snowflake.Data.Tests/UnitTests/SFRemoteStorageClientTest.cs b/Snowflake.Data.Tests/UnitTests/SFRemoteStorageClientTest.cs index 4be002666..f9ea9f6ea 100644 --- a/Snowflake.Data.Tests/UnitTests/SFRemoteStorageClientTest.cs +++ b/Snowflake.Data.Tests/UnitTests/SFRemoteStorageClientTest.cs @@ -35,7 +35,7 @@ class SFRemoteStorageClientTest : SFBaseTest const string StorageAccount = "mockStorageAccount"; - const string RealSourceFilePath = "realSrcFilePath.txt"; + [ThreadStatic] private static string t_realSourceFilePath; PutGetEncryptionMaterial EncryptionMaterial = new PutGetEncryptionMaterial() { @@ -51,9 +51,8 @@ class SFRemoteStorageClientTest : SFBaseTest const int Parallel = 0; // File name for mock test files + [ThreadStatic] private static string t_downloadFileName; const string LocalLocation = "./"; - const string DownloadFileName = "mockDownloadFileName.txt"; - const string DownloadFilePath = LocalLocation + DownloadFileName; // Mock upload file size const int UploadFileSize = 9999; @@ -74,14 +73,17 @@ class SFRemoteStorageClientTest : SFBaseTest [SetUp] public void BeforeTest() { + t_realSourceFilePath = TestNameWithWorker + "_realSrcFilePath.txt"; + t_downloadFileName = TestNameWithWorker + "_mockFileName.txt"; + _fileMetadata = new SFFileMetadata() { - destFileName = DownloadFileName, + destFileName = t_downloadFileName, localLocation = LocalLocation, MaxBytesInMemory = 1024, memoryStream = new MemoryStream(Encoding.UTF8.GetBytes(MockRemoteStorageClient.FileContent)), parallel = Parallel, - realSrcFilePath = RealSourceFilePath, + realSrcFilePath = t_realSourceFilePath, stageInfo = new PutGetStageInfo() { endPoint = EndPoint, @@ -129,9 +131,9 @@ public void AfterTest() } // Delete temporary files from download - if (File.Exists(DownloadFilePath)) + if (File.Exists(t_downloadFileName)) { - File.Delete(DownloadFilePath); + File.Delete(t_downloadFileName); } } @@ -426,7 +428,7 @@ public void TestDownloadOneFile(HttpStatusCode httpStatusCode, ResultStatus expe // Assert if (expectedResultStatus == ResultStatus.DOWNLOADED) { - string text = File.ReadAllText(_fileMetadata.localLocation + _fileMetadata.destFileName); + string text = File.ReadAllText(t_downloadFileName); Assert.AreEqual(MockRemoteStorageClient.FileContent, text); } Assert.AreEqual(expectedResultStatus.ToString(), _fileMetadata.resultStatus); @@ -446,7 +448,7 @@ public async Task TestDownloadOneFileAsync(HttpStatusCode httpStatusCode, Result // Assert if (expectedResultStatus == ResultStatus.DOWNLOADED) { - string text = File.ReadAllText(_fileMetadata.localLocation + _fileMetadata.destFileName); + string text = File.ReadAllText(t_downloadFileName); Assert.AreEqual(MockRemoteStorageClient.FileContent, text); } Assert.AreEqual(expectedResultStatus.ToString(), _fileMetadata.resultStatus); @@ -521,12 +523,12 @@ private void SetUpMockEncryptedFileForDownload() _fileMetadata.encryptionMaterial = EncryptionMaterial; // Write file to encrypt - File.WriteAllText(DownloadFilePath, MockRemoteStorageClient.FileContent); + File.WriteAllText(t_downloadFileName, MockRemoteStorageClient.FileContent); // Get encrypted stream from file SFEncryptionMetadata encryptionMetadata = new SFEncryptionMetadata(); Stream stream = EncryptionProvider.EncryptFile( - DownloadFilePath, + t_downloadFileName, _fileMetadata.encryptionMaterial, encryptionMetadata, FileTransferConfiguration.FromFileMetadata(_fileMetadata)); @@ -547,7 +549,7 @@ public void TestDownloadOneFileWithEncryptionMaterial(HttpStatusCode httpStatusC SFRemoteStorageUtil.DownloadOneFile(_fileMetadata); // Assert - string text = File.ReadAllText(DownloadFilePath); + string text = File.ReadAllText(t_downloadFileName); Assert.AreEqual(MockRemoteStorageClient.FileContent, text); Assert.AreEqual(expectedResultStatus.ToString(), _fileMetadata.resultStatus); } @@ -564,7 +566,7 @@ public async Task TestDownloadOneFileAsyncWithEncryptionMaterial(HttpStatusCode await SFRemoteStorageUtil.DownloadOneFileAsync(_fileMetadata, _cancellationToken).ConfigureAwait(false); // Assert - string text = File.ReadAllText(DownloadFilePath); + string text = File.ReadAllText(t_downloadFileName); Assert.AreEqual(MockRemoteStorageClient.FileContent, text); Assert.AreEqual(expectedResultStatus.ToString(), _fileMetadata.resultStatus); } From 3a9b9cbb8d6e3fb6baa2c4efecd6ddc8aeb5c49d Mon Sep 17 00:00:00 2001 From: sfc-gh-ext-simba-lf Date: Tue, 19 Sep 2023 14:13:23 -0700 Subject: [PATCH 4/4] SNOW-723810: Rename test to be more descriptive --- Snowflake.Data.Tests/UnitTests/SFRemoteStorageClientTest.cs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/Snowflake.Data.Tests/UnitTests/SFRemoteStorageClientTest.cs b/Snowflake.Data.Tests/UnitTests/SFRemoteStorageClientTest.cs index f9ea9f6ea..fa7df6250 100644 --- a/Snowflake.Data.Tests/UnitTests/SFRemoteStorageClientTest.cs +++ b/Snowflake.Data.Tests/UnitTests/SFRemoteStorageClientTest.cs @@ -186,7 +186,7 @@ public void TestGetRemoteStorageClient(string stageType) [TestCase(false, true)] [TestCase(true, false)] [TestCase(true, true)] - public void TestGetFileBytes(bool containsEncryptionMaterial, bool useMemoryStream) + public void TestUploadFileOrStreamWithAndWithoutEncryption(bool containsEncryptionMaterial, bool useMemoryStream) { // Arrange var mockWebRequest = new Mock();