Skip to content

Commit

Permalink
Modify test to include get
Browse files Browse the repository at this point in the history
  • Loading branch information
@sfc-gh-ext-simba-lf
sfc-gh-ext-simba-lf committed Dec 6, 2024
1 parent f74a4c8 commit e7302a0
Showing 1 changed file with 35 additions and 39 deletions.
74 changes: 35 additions & 39 deletions Snowflake.Data.Tests/IntegrationTests/SFPutGetTest.cs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -561,38 +561,35 @@ public void TestPutGetGcsDownscopedCredential(
}

[Test]
public void TestPutFileWithSpaceAndSingleQuote()
public void TestPutGetFileWithSpaceAndSingleQuote(
[Values] StageType stageType,
[Values("/STAGE PATH WITH SPACE")] string stagePath)
{
var absolutePathPrefix = Path.Combine(Path.GetTempPath(), $"{Guid.NewGuid()} file path with space");
var files = new List<string> {
$"{absolutePathPrefix}_one.csv",
$"{absolutePathPrefix}_two.csv",
$"{absolutePathPrefix}_three.csv"
};
PrepareFileData(files);

// Set the PUT query variables
t_inputFilePath = $"{absolutePathPrefix}*";
t_internalStagePath = $"@{t_schemaName}.{t_stageName}";

PrepareTest(null, stageType, stagePath, false, true, true);
using (var conn = new SnowflakeDbConnection(ConnectionString))
{
conn.Open();
PutFile(conn, "", ResultStatus.UPLOADED, true);
VerifyFilesAreUploaded(conn, files, t_internalStagePath);
CopyIntoTable(conn, true);
GetFile(conn, true);
}
}

private void PrepareTest(string sourceFileCompressionType, StageType stageType, string stagePath,
bool autoCompress, bool clientEncryption = true)
bool autoCompress, bool clientEncryption = true, bool makeFilePathWithSpace = false)
{
t_stageType = stageType;
t_sourceCompressionType = sourceFileCompressionType;
t_autoCompress = autoCompress;
// Prepare temp file name with specified file extension
t_fileName = Guid.NewGuid() + ".csv" +
(t_autoCompress? SFFileCompressionTypes.LookUpByName(t_sourceCompressionType).FileExtension: "");
t_inputFilePath = Path.GetTempPath() + t_fileName;
(t_autoCompress ? SFFileCompressionTypes.LookUpByName(t_sourceCompressionType).FileExtension : "");
var sourceFolderWithSpace = $"{Guid.NewGuid()} source file path with space";
var inputPathBase = makeFilePathWithSpace ?
Path.Combine(s_outputDirectory, sourceFolderWithSpace) :
Path.GetTempPath();
t_inputFilePath = Path.Combine(inputPathBase, t_fileName);

if (IsCompressedByTheDriver())
{
t_destCompressionType = "gzip";
Expand All @@ -603,7 +600,16 @@ private void PrepareTest(string sourceFileCompressionType, StageType stageType,
t_destCompressionType = t_sourceCompressionType;
t_outputFileName = t_fileName;
}
t_outputFilePath = $@"{s_outputDirectory}/{t_outputFileName}";
var destinationFolderWithSpace = $"{Guid.NewGuid()} destination file path with space";
var outputPathBase = makeFilePathWithSpace ?
Path.Combine(s_outputDirectory, destinationFolderWithSpace) :
s_outputDirectory;
t_outputFilePath = Path.Combine(outputPathBase, t_outputFileName);
if (makeFilePathWithSpace)
{
Directory.CreateDirectory(inputPathBase);
Directory.CreateDirectory(outputPathBase);
}
t_filesToDelete.Add(t_outputFilePath);
PrepareFileData(t_inputFilePath);

Expand Down Expand Up @@ -640,23 +646,10 @@ string PutFile(
using (var command = conn.CreateCommand())
{
// Prepare PUT query
string putQuery;
if (encloseInSingleQuotes)
{
// Enclosing the file path in single quotes require forward slash
t_inputFilePath = t_inputFilePath.Replace("\\", "/");
putQuery =
$"PUT 'file://{t_inputFilePath}' {t_internalStagePath}" +
$" AUTO_COMPRESS={(t_autoCompress ? "TRUE" : "FALSE")}";
}
else
{
putQuery =
$"PUT file://{t_inputFilePath} {t_internalStagePath}" +
$" AUTO_COMPRESS={(t_autoCompress ? "TRUE" : "FALSE")}" +
$" {additionalAttribute}";
}

var putQuery = encloseInSingleQuotes ?
$"PUT 'file://{t_inputFilePath.Replace("\\", "/")}' '{t_internalStagePath}'" :
$"PUT file://{t_inputFilePath} {t_internalStagePath}";
putQuery += $" AUTO_COMPRESS={(t_autoCompress ? "TRUE" : "FALSE")}" + $" {additionalAttribute}";
// Upload file
command.CommandText = putQuery;
var reader = command.ExecuteReader();
Expand Down Expand Up @@ -698,7 +691,7 @@ string PutFile(
}

// COPY INTO - Copy data from the stage into temp table
private void CopyIntoTable(SnowflakeDbConnection conn)
private void CopyIntoTable(SnowflakeDbConnection conn, bool encloseInSingleQuotes = false)
{
using (var command = conn.CreateCommand())
{
Expand All @@ -708,7 +701,8 @@ private void CopyIntoTable(SnowflakeDbConnection conn)
command.CommandText = $"COPY INTO {t_schemaName}.{t_tableName}";
break;
default:
command.CommandText =
command.CommandText = encloseInSingleQuotes ?
$"COPY INTO {t_schemaName}.{t_tableName} FROM '{t_internalStagePath}/{t_fileName}'" :
$"COPY INTO {t_schemaName}.{t_tableName} FROM {t_internalStagePath}/{t_fileName}";
break;
}
Expand All @@ -733,12 +727,14 @@ private void CopyIntoTable(SnowflakeDbConnection conn)
}

// GET - Download from the stage into local directory
private void GetFile(DbConnection conn)
private void GetFile(DbConnection conn, bool encloseInSingleQuotes = false)
{
using (var command = conn.CreateCommand())
{
// Prepare GET query
var getQuery = $"GET {t_internalStagePath}/{t_fileName} file://{s_outputDirectory}";
var getQuery = encloseInSingleQuotes ?
$"GET '{t_internalStagePath}/{t_fileName}' 'file://{Path.GetDirectoryName(t_outputFilePath).Replace("\\", "/")}'" :
$"GET {t_internalStagePath}/{t_fileName} file://{s_outputDirectory}";

// Download file
command.CommandText = getQuery;
Expand Down

0 comments on commit e7302a0

Please sign in to comment.