Skip to content

Commit

Permalink
SNOW-1545648: Fix PUT command error if file path contains spaces and …
Browse files Browse the repository at this point in the history 
…single quotes (#1066)

Co-authored-by: Krzysztof Nozderko <[email protected]>
  • Loading branch information
@sfc-gh-ext-simba-lf @sfc-gh-knozderko
sfc-gh-ext-simba-lf and sfc-gh-knozderko authored Jan 8, 2025
1 parent 580da21 commit 27b0ae4
Show file tree
Hide file tree
Showing 4 changed files with 97 additions and 24 deletions.
14 changes: 7 additions & 7 deletions Snowflake.Data.Tests/IntegrationTests/FileUploadDownloadLargeFilesIT.cs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -23,33 +23,33 @@ public class FileUploadDownloadLargeFilesIT : SFBaseTest
private static readonly string s_fullFileName = Path.Combine(s_localFolderName, FileName);
private static readonly string s_fullDownloadedFileName = Path.Combine(s_downloadFolderName, FileName);
private static readonly MD5 s_md5 = MD5.Create();

[OneTimeSetUp]
public static void GenerateLargeFileForTests()
{
CreateLocalDirectory(s_localFolderName);
GenerateLargeFile(s_fullFileName);
}

[OneTimeTearDown]
public static void DeleteGeneratedLargeFile()
{
RemoveLocalFile(s_fullFileName);
RemoveDirectory(s_localFolderName);
}

[Test]
public void TestThatUploadsAndDownloadsTheSameFile()
{
// act
UploadFile(s_fullFileName, s_remoteFolderName);
DownloadFile(s_remoteFolderName, s_downloadFolderName, FileName);

// assert
Assert.AreEqual(
CalcualteMD5(s_fullFileName),
CalcualteMD5(s_fullDownloadedFileName));

// cleanup
RemoveFilesFromServer(s_remoteFolderName);
RemoveLocalFile(s_fullDownloadedFileName);
Expand Down Expand Up @@ -85,7 +85,7 @@ private void DownloadFile(string remoteFolderName, string downloadFolderName, st
command.ExecuteNonQuery();
}
}

private void RemoveFilesFromServer(string remoteFolderName)
{
using (var conn = new SnowflakeDbConnection())
Expand All @@ -108,7 +108,7 @@ private static string CalcualteMD5(string fullFileName)
}

private static void RemoveLocalFile(string fullFileName) => File.Delete(fullFileName);

private static void CreateLocalDirectory(string path) => Directory.CreateDirectory(path);

private static void RemoveDirectory(string path) => Directory.Delete(path, true);
Expand Down
61 changes: 47 additions & 14 deletions Snowflake.Data.Tests/IntegrationTests/SFPutGetTest.cs
View file
Original file line number Diff line number Diff line change
@@ -1,4 +1,4 @@
/*
/*
* Copyright (c) 2012-2023 Snowflake Computing Inc. All rights reserved.
*/

Expand Down Expand Up @@ -560,16 +560,36 @@ public void TestPutGetGcsDownscopedCredential(
}
}

[Test]
public void TestPutGetFileWithSpaceAndSingleQuote(
[Values] StageType stageType,
[Values("/STAGE PATH WITH SPACE")] string stagePath)
{
PrepareTest(null, stageType, stagePath, false, true, true);
using (var conn = new SnowflakeDbConnection(ConnectionString))
{
conn.Open();
PutFile(conn, "", ResultStatus.UPLOADED, true);
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 @@ -580,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 @@ -610,16 +639,17 @@ private static bool IsCompressedByTheDriver()
string PutFile(
SnowflakeDbConnection conn,
String additionalAttribute = "",
ResultStatus expectedStatus = ResultStatus.UPLOADED)
ResultStatus expectedStatus = ResultStatus.UPLOADED,
bool encloseInSingleQuotes = false)
{
string queryId;
using (var command = conn.CreateCommand())
{
// Prepare PUT query
string 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 @@ -661,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 @@ -671,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 @@ -696,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
36 changes: 35 additions & 1 deletion Snowflake.Data.Tests/UnitTests/SFFileTransferAgentTests.cs
View file
Original file line number Diff line number Diff line change
@@ -1,4 +1,4 @@
/*
/*
* Copyright (c) 2012-2023 Snowflake Computing Inc. All rights reserved.
*/

Expand Down Expand Up @@ -69,6 +69,10 @@ class SFFileTransferAgentTest : SFBaseTest
// Mock file content
const string FileContent = "FTAFileContent";

// Mock file paths
static readonly string s_filePathWithoutSpaces = Path.Combine("C:\\Users\\Test\\", "folder_without_space", "*.*");
static readonly string s_filePathWithSpaces = Path.Combine("C:\\Users\\Test\\", "folder with space", "*.*");

[SetUp]
public void BeforeEachTest()
{
Expand Down Expand Up @@ -634,5 +638,35 @@ public void TestDownloadThrowsErrorDirectoryNotFound()
Assert.IsInstanceOf<DirectoryNotFoundException>(innerException);
Assert.That(innerException?.Message, Does.Match("Could not find a part of the path .*"));
}

[Test]
public void TestGetFilePathWithoutSpacesFromPutCommand()
{
TestGetFilePathFromPutCommand("PUT file://" + s_filePathWithoutSpaces + " @TestStage", s_filePathWithoutSpaces);
}

[Test]
public void TestGetFilePathWithSpacesFromPutCommand()
{
TestGetFilePathFromPutCommand("PUT file://" + s_filePathWithSpaces + " @TestStage", s_filePathWithSpaces);
}

[Test]
public void TestGetFilePathWithoutSpacesAndWithSingleQuotesFromPutCommand()
{
TestGetFilePathFromPutCommand("PUT 'file://" + s_filePathWithoutSpaces + "' @TestStage", s_filePathWithoutSpaces);
}

[Test]
public void TestGetFilePathWithSpacesAndWithSingleQuotesFromPutCommand()
{
TestGetFilePathFromPutCommand("PUT 'file://" + s_filePathWithSpaces + "' @TestStage", s_filePathWithSpaces);
}

public void TestGetFilePathFromPutCommand(string query, string expectedFilePath)
{
var actualFilePath = SFFileTransferAgent.getFilePathFromPutCommand(query);
Assert.AreEqual(expectedFilePath, actualFilePath);
}
}
}
10 changes: 8 additions & 2 deletions Snowflake.Data/Core/FileTransfer/SFFileTransferAgent.cs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -527,13 +527,19 @@ internal async Task updatePresignedUrlAsync(CancellationToken cancellationToken)
/// </summary>
/// <param name="query">The query containing the file path</param>
/// <returns>The file path contained by the query</returns>
private string getFilePathFromPutCommand(string query)
internal static string getFilePathFromPutCommand(string query)
{
// Extract file path from PUT command:
// E.g. "PUT file://C:<path-to-file> @DB.SCHEMA.%TABLE;"
int startIndex = query.IndexOf("file://") + "file://".Length;
int endIndex = query.Substring(startIndex).IndexOf('@') - 1;
string filePath = query.Substring(startIndex, endIndex);
string filePath = query.Substring(startIndex, endIndex).TrimEnd();

// Check if file path contains an enclosing (') char
if (filePath[filePath.Length - 1] == '\'')
{
filePath = filePath.Substring(0, filePath.Length - 1);
}
return filePath;
}

Expand Down

0 comments on commit 27b0ae4

Please sign in to comment.