Skip to content

Commit

Permalink
SNOW-1672654 Support for empty encryptionMaterial (#1048)
Browse files Browse the repository at this point in the history
  • Loading branch information
@sfc-gh-dstempniak
sfc-gh-dstempniak authored Oct 31, 2024
1 parent 5e16693 commit d8e4b63
Show file tree
Hide file tree
Showing 9 changed files with 155 additions and 113 deletions.
110 changes: 68 additions & 42 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,7 +42,7 @@ 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,
Expand All @@ -63,7 +64,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 +74,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 +90,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 +115,7 @@ public void TearDown()
command.ExecuteNonQuery();
}
}

// Delete temp files if necessary
if (t_filesToDelete != null)
{
Expand All @@ -130,7 +136,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 +148,7 @@ public void TestPutFileAsteriskWildcard()
VerifyFilesAreUploaded(conn, files, t_internalStagePath);
}
}

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

[Test]
public void TestPutFileQuestionMarkWildcard()
{
Expand All @@ -180,7 +186,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 +198,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 +232,7 @@ public void TestPutGetOnClosedConnectionThrowsWithoutQueryId([Values("GET", "PUT
SnowflakeDbExceptionAssert.HasErrorCode(snowflakeDbException, SFError.EXECUTE_COMMAND_ON_CLOSED_CONNECTION);
}
}

[Test]
public void TestGetNonExistentFileReturnsFalseAndDoesNotThrow()
{
Expand All @@ -236,7 +242,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 +252,7 @@ public void TestGetNonExistentFileReturnsFalseAndDoesNotThrow()
}
}
}

[Test]
public void TestPutNonExistentFileThrowsWithQueryId()
{
Expand All @@ -256,14 +262,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 +291,7 @@ public void TestPutFileProvidesQueryIdOnFailure()
SnowflakeDbExceptionAssert.HasErrorCode(snowflakeDbException, SFError.IO_ERROR_ON_GETPUT_COMMAND);
}
}

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

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

// Assert
Assert.IsNotNull(queryId);
Assert.DoesNotThrow(()=>Guid.Parse(queryId));
Expand All @@ -337,11 +343,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 +357,7 @@ public void TestPutFileRelativePathWithDirectory()
VerifyFilesAreUploaded(conn, new List<string> { t_inputFilePath }, t_internalStagePath);
}
}

[Test]
public void TestPutFileRelativePathAsteriskWildcard()
{
Expand All @@ -374,7 +380,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 +390,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 +401,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 +421,7 @@ public void TestPutFileWithOverwriteFlagRunsSecondUpload()
PutFile(conn, overwriteAttribute, expectedStatus: ResultStatus.UPLOADED);
}
}

[Test]
public void TestPutDirectoryAsteriskWildcard()
{
Expand All @@ -431,7 +437,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 +465,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 +477,7 @@ public void TestPutDirectoryQuestionMarkWildcard()
VerifyFilesAreUploaded(conn, files, t_internalStagePath);
}
}

[Test]
public void TestPutDirectoryMixedWildcard()
{
Expand All @@ -487,7 +493,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 +505,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 +523,24 @@ public void TestPutGetCommand(
GetFile(conn);
}
}


[Test]
public void TestPutGetCommandForNamedStageWithoutClientSideEncryption(
[Values("none", "gzip")] string sourceFileCompressionType,
[Values("", "/DEEP/TEST_PATH")] string stagePath,
[Values] bool autoCompress)
{
PrepareTest(sourceFileCompressionType, StageType.NAMED, stagePath, autoCompress, false);

using (var conn = new SnowflakeDbConnection(ConnectionString))
{
conn.Open();
PutFile(conn);
CopyIntoTable(conn);
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 +559,15 @@ public void TestPutGetGcsDownscopedCredential(
GetFile(conn);
}
}

private void PrepareTest(string sourceFileCompressionType, StageType stageType, string stagePath, bool autoCompress)

private void PrepareTest(string sourceFileCompressionType, StageType stageType, string stagePath,
bool autoCompress, bool clientEncryption = true)
{
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 All @@ -570,7 +594,9 @@ private void PrepareTest(string sourceFileCompressionType, StageType stageType,
t_internalStagePath = $"@{t_schemaName}.%{t_tableName}{stagePath}";
break;
case StageType.NAMED:
t_internalStagePath = $"@{t_schemaName}.{t_stageName}{stagePath}";
t_internalStagePath = clientEncryption
? $"@{t_schemaName}.{t_stageName}{stagePath}"
: $"@{t_schemaName}.{t_stageNameSse}{stagePath}";
break;
}
}
Expand All @@ -579,11 +605,11 @@ 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 +730,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 +773,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
Loading

0 comments on commit d8e4b63

Please sign in to comment.