Skip to content

Commit

Permalink
SNOW-1672654 fix for empty encryptionMaterial when client side encryp…
Browse files Browse the repository at this point in the history 
…tion is disabled
  • Loading branch information
@sfc-gh-dstempniak
sfc-gh-dstempniak committed Oct 28, 2024
1 parent 8199563 commit 9fe83a1
Show file tree
Hide file tree
Showing 7 changed files with 72 additions and 61 deletions.
92 changes: 51 additions & 41 deletions Snowflake.Data.Tests/IntegrationTests/SFPutGetTest.cs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -31,6 +31,7 @@ class SFPutGetTest : SFBaseTest
[ThreadStatic] private static string t_schemaName;
[ThreadStatic] private static string t_tableName;
[ThreadStatic] private static string t_stageName;
[ThreadStatic] private static string t_stageNameSse; // server side encryption without client side encryption
[ThreadStatic] private static string t_fileName;
[ThreadStatic] private static string t_outputFileName;
[ThreadStatic] private static string t_inputFilePath;
Expand All @@ -41,12 +42,13 @@ class SFPutGetTest : SFBaseTest
[ThreadStatic] private static string t_destCompressionType;
[ThreadStatic] private static bool t_autoCompress;
[ThreadStatic] private static List<string> t_filesToDelete;

public enum StageType
{
USER,
TABLE,
NAMED
NAMED,
NAMED_SSE // without client side encryption
}

[OneTimeSetUp]
Expand All @@ -63,7 +65,7 @@ public static void OneTimeTearDown()
// Delete temp output directory and downloaded files
Directory.Delete(s_outputDirectory, true);
}

[SetUp]
public void SetUp()
{
Expand All @@ -73,6 +75,7 @@ public void SetUp()
t_schemaName = testConfig.schema;
t_tableName = $"TABLE_{threadSuffix}";
t_stageName = $"STAGE_{threadSuffix}";
t_stageNameSse = $"STAGE_{threadSuffix}_SSE";
t_filesToDelete = new List<string>();

using (var conn = new SnowflakeDbConnection(ConnectionString))
Expand All @@ -88,6 +91,10 @@ public void SetUp()
// Create temp stage
command.CommandText = $"CREATE OR REPLACE STAGE {t_schemaName}.{t_stageName}";
command.ExecuteNonQuery();

// Create temp stage without client side encryption
command.CommandText = $"CREATE OR REPLACE STAGE {t_schemaName}.{t_stageNameSse} ENCRYPTION = (TYPE = 'SNOWFLAKE_SSE')";
command.ExecuteNonQuery();
}
}
}
Expand All @@ -109,7 +116,7 @@ public void TearDown()
command.ExecuteNonQuery();
}
}

// Delete temp files if necessary
if (t_filesToDelete != null)
{
Expand All @@ -130,7 +137,7 @@ public void TestPutFileAsteriskWildcard()
$"{absolutePathPrefix}_three.csv"
};
PrepareFileData(files);

// Set the PUT query variables
t_inputFilePath = $"{absolutePathPrefix}*";
t_internalStagePath = $"@{t_schemaName}.{t_stageName}";
Expand All @@ -142,7 +149,7 @@ public void TestPutFileAsteriskWildcard()
VerifyFilesAreUploaded(conn, files, t_internalStagePath);
}
}

[Test]
public void TestPutFileAsteriskWildcardWithExtension()
{
Expand All @@ -167,7 +174,7 @@ public void TestPutFileAsteriskWildcardWithExtension()
VerifyFilesAreUploaded(conn, files, t_internalStagePath);
}
}

[Test]
public void TestPutFileQuestionMarkWildcard()
{
Expand All @@ -180,7 +187,7 @@ public void TestPutFileQuestionMarkWildcard()
PrepareFileData(files);
// Create file which should be omitted during the transfer
PrepareFileData($"{absolutePathPrefix}_four.csv");

// Set the PUT query variables
t_inputFilePath = $"{absolutePathPrefix}_?.csv";
t_internalStagePath = $"@{t_schemaName}.{t_stageName}";
Expand All @@ -192,14 +199,14 @@ public void TestPutFileQuestionMarkWildcard()
VerifyFilesAreUploaded(conn, files, t_internalStagePath);
}
}

[Test]
public void TestPutFileRelativePathWithoutDirectory()
{
// Set the PUT query variables
t_inputFilePath = $"{Guid.NewGuid()}_1.csv";
t_internalStagePath = $"@{t_schemaName}.{t_stageName}";

PrepareFileData(t_inputFilePath);

using (var conn = new SnowflakeDbConnection(ConnectionString))
Expand All @@ -226,7 +233,7 @@ public void TestPutGetOnClosedConnectionThrowsWithoutQueryId([Values("GET", "PUT
SnowflakeDbExceptionAssert.HasErrorCode(snowflakeDbException, SFError.EXECUTE_COMMAND_ON_CLOSED_CONNECTION);
}
}

[Test]
public void TestGetNonExistentFileReturnsFalseAndDoesNotThrow()
{
Expand All @@ -236,7 +243,7 @@ public void TestGetNonExistentFileReturnsFalseAndDoesNotThrow()
// Act
using (var conn = new SnowflakeDbConnection(ConnectionString))
{
conn.Open();
conn.Open();
var sql = $"GET {t_internalStagePath}/{t_fileName} file://{s_outputDirectory}";
using (var command = conn.CreateCommand())
{
Expand All @@ -246,7 +253,7 @@ public void TestGetNonExistentFileReturnsFalseAndDoesNotThrow()
}
}
}

[Test]
public void TestPutNonExistentFileThrowsWithQueryId()
{
Expand All @@ -256,14 +263,14 @@ public void TestPutNonExistentFileThrowsWithQueryId()
// Act
using (var conn = new SnowflakeDbConnection(ConnectionString))
{
conn.Open();
conn.Open();
var snowflakeDbException = Assert.Throws<SnowflakeDbException>(() => PutFile(conn));
Assert.IsNotNull(snowflakeDbException);
Assert.IsNotNull(snowflakeDbException.QueryId);
SnowflakeDbExceptionAssert.HasErrorCode(snowflakeDbException, SFError.IO_ERROR_ON_GETPUT_COMMAND);
}
}

[Test]
public void TestPutFileProvidesQueryIdOnFailure()
{
Expand All @@ -285,7 +292,7 @@ public void TestPutFileProvidesQueryIdOnFailure()
SnowflakeDbExceptionAssert.HasErrorCode(snowflakeDbException, SFError.IO_ERROR_ON_GETPUT_COMMAND);
}
}

[Test]
public void TestPutFileWithSyntaxErrorProvidesQueryIdOnFailure()
{
Expand All @@ -308,7 +315,7 @@ public void TestPutFileWithSyntaxErrorProvidesQueryIdOnFailure()
Assert.That(snowflakeDbException.InnerException, Is.Null);
}
}

[Test]
public void TestPutFileProvidesQueryIdOnSuccess()
{
Expand All @@ -323,7 +330,7 @@ public void TestPutFileProvidesQueryIdOnSuccess()
{
conn.Open();
var queryId = PutFile(conn);

// Assert
Assert.IsNotNull(queryId);
Assert.DoesNotThrow(()=>Guid.Parse(queryId));
Expand All @@ -337,11 +344,11 @@ public void TestPutFileRelativePathWithDirectory()
var guid = Guid.NewGuid();
var relativePath = $"{guid}";
Directory.CreateDirectory(relativePath);

// Set the PUT query variables
t_inputFilePath = $"{relativePath}{Path.DirectorySeparatorChar}{guid}_1.csv";
t_internalStagePath = $"@{t_schemaName}.{t_stageName}";

PrepareFileData(t_inputFilePath);

using (var conn = new SnowflakeDbConnection(ConnectionString))
Expand All @@ -351,7 +358,7 @@ public void TestPutFileRelativePathWithDirectory()
VerifyFilesAreUploaded(conn, new List<string> { t_inputFilePath }, t_internalStagePath);
}
}

[Test]
public void TestPutFileRelativePathAsteriskWildcard()
{
Expand All @@ -374,7 +381,7 @@ public void TestPutFileRelativePathAsteriskWildcard()
VerifyFilesAreUploaded(conn, files, t_internalStagePath);
}
}

[Test]
// presigned url is enabled on CI so we need to disable the test
// it should be enabled when downscoped credential is the default option
Expand All @@ -384,7 +391,7 @@ public void TestPutFileWithoutOverwriteFlagSkipsSecondUpload()
// Set the PUT query variables
t_inputFilePath = $"{Guid.NewGuid()}.csv";
t_internalStagePath = $"@{t_schemaName}.{t_stageName}";

PrepareFileData(t_inputFilePath);

using (var conn = new SnowflakeDbConnection(ConnectionString))
Expand All @@ -395,18 +402,18 @@ public void TestPutFileWithoutOverwriteFlagSkipsSecondUpload()
PutFile(conn, expectedStatus: ResultStatus.SKIPPED);
}
}

[Test]
public void TestPutFileWithOverwriteFlagRunsSecondUpload()
{
var overwriteAttribute = "OVERWRITE=TRUE";

// Set the PUT query variables
t_inputFilePath = $"{Guid.NewGuid()}.csv";
t_internalStagePath = $"@{t_schemaName}.{t_stageName}";

PrepareFileData(t_inputFilePath);

using (var conn = new SnowflakeDbConnection(ConnectionString))
{
conn.Open();
Expand All @@ -415,7 +422,7 @@ public void TestPutFileWithOverwriteFlagRunsSecondUpload()
PutFile(conn, overwriteAttribute, expectedStatus: ResultStatus.UPLOADED);
}
}

[Test]
public void TestPutDirectoryAsteriskWildcard()
{
Expand All @@ -431,7 +438,7 @@ public void TestPutDirectoryAsteriskWildcard()
PrepareFileData(fullPath);
files.Add(fullPath);
}

// Set the PUT query variables
t_inputFilePath = $"{path}*{Path.DirectorySeparatorChar}*";
t_internalStagePath = $"@{t_schemaName}.{t_stageName}";
Expand Down Expand Up @@ -459,7 +466,7 @@ public void TestPutDirectoryQuestionMarkWildcard()
PrepareFileData(fullPath);
files.Add(fullPath);
}

// Set the PUT query variables
t_inputFilePath = $"{path}_?{Path.DirectorySeparatorChar}{guid}_?_file.csv";
t_internalStagePath = $"@{t_schemaName}.{t_stageName}";
Expand All @@ -471,7 +478,7 @@ public void TestPutDirectoryQuestionMarkWildcard()
VerifyFilesAreUploaded(conn, files, t_internalStagePath);
}
}

[Test]
public void TestPutDirectoryMixedWildcard()
{
Expand All @@ -487,7 +494,7 @@ public void TestPutDirectoryMixedWildcard()
PrepareFileData(fullPath);
files.Add(fullPath);
}

// Set the PUT query variables
t_inputFilePath = $"{path}_*{Path.DirectorySeparatorChar}{guid}_?_file.csv";
t_internalStagePath = $"@{t_schemaName}.{t_stageName}";
Expand All @@ -499,7 +506,7 @@ public void TestPutDirectoryMixedWildcard()
VerifyFilesAreUploaded(conn, files, t_internalStagePath);
}
}

[Test]
public void TestPutGetCommand(
[Values("none", "gzip", "bzip2", "brotli", "deflate", "raw_deflate", "zstd")] string sourceFileCompressionType,
Expand All @@ -517,7 +524,7 @@ public void TestPutGetCommand(
GetFile(conn);
}
}

// Test small file upload/download with GCS_USE_DOWNSCOPED_CREDENTIAL set to true
[Test]
[IgnoreOnEnvIs("snowflake_cloud_env", new [] { "AWS", "AZURE" })]
Expand All @@ -536,14 +543,14 @@ public void TestPutGetGcsDownscopedCredential(
GetFile(conn);
}
}

private void PrepareTest(string sourceFileCompressionType, StageType stageType, string stagePath, bool autoCompress)
{
t_stageType = stageType;
t_sourceCompressionType = sourceFileCompressionType;
t_autoCompress = autoCompress;
// Prepare temp file name with specified file extension
t_fileName = Guid.NewGuid() + ".csv" +
t_fileName = Guid.NewGuid() + ".csv" +
(t_autoCompress? SFFileCompressionTypes.LookUpByName(t_sourceCompressionType).FileExtension: "");
t_inputFilePath = Path.GetTempPath() + t_fileName;
if (IsCompressedByTheDriver())
Expand Down Expand Up @@ -572,18 +579,21 @@ private void PrepareTest(string sourceFileCompressionType, StageType stageType,
case StageType.NAMED:
t_internalStagePath = $"@{t_schemaName}.{t_stageName}{stagePath}";
break;
case StageType.NAMED_SSE:
t_internalStagePath = $"@{t_schemaName}.{t_stageNameSse}{stagePath}";
break;
}
}

private static bool IsCompressedByTheDriver()
{
return t_sourceCompressionType == "none" && t_autoCompress;
}

// PUT - upload file from local directory to the stage
string PutFile(
SnowflakeDbConnection conn,
String additionalAttribute = "",
SnowflakeDbConnection conn,
String additionalAttribute = "",
ResultStatus expectedStatus = ResultStatus.UPLOADED)
{
string queryId;
Expand Down Expand Up @@ -704,7 +714,7 @@ private void ProcessFile(String command, SnowflakeDbConnection connection)
{
switch (command)
{
case "GET":
case "GET":
GetFile(connection);
break;
case "PUT":
Expand Down Expand Up @@ -747,7 +757,7 @@ private static void PrepareFileData(string file)
// Prepare csv raw data and write to temp files
var rawDataRow = string.Join(",", s_colData) + "\n";
var rawData = string.Concat(Enumerable.Repeat(rawDataRow, NumberOfRows));

File.WriteAllText(file, rawData);
t_filesToDelete.Add(file);
}
Expand Down
1 change: 0 additions & 1 deletion Snowflake.Data.Tests/UnitTests/SFAzureClientTest.cs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -69,7 +69,6 @@ class SFAzureClientTest : SFBaseTest
stageInfo = new PutGetStageInfo()
{
endPoint = EndPoint,
isClientSideEncrypted = true,
location = Location,
locationType = SFRemoteStorageUtil.AZURE_FS,
path = LocationPath,
Expand Down
1 change: 0 additions & 1 deletion Snowflake.Data.Tests/UnitTests/SFGCSClientTest.cs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -62,7 +62,6 @@ class SFGCSClientTest : SFBaseTest
stageInfo = new PutGetStageInfo()
{
endPoint = null,
isClientSideEncrypted = true,
location = Location,
locationType = SFRemoteStorageUtil.GCS_FS,
path = LocationPath,
Expand Down
1 change: 0 additions & 1 deletion Snowflake.Data.Tests/UnitTests/SFRemoteStorageClientTest.cs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -87,7 +87,6 @@ class SFRemoteStorageClientTest : SFBaseTest
stageInfo = new PutGetStageInfo()
{
endPoint = EndPoint,
isClientSideEncrypted = true,
location = Location,
locationType = SFRemoteStorageUtil.GCS_FS,
path = LocationPath,
Expand Down
Loading

0 comments on commit 9fe83a1

Please sign in to comment.