SNOW-1545648: Fix PUT command error if file path contains spaces and single quotes

Snowflake.Data.Tests/IntegrationTests/FileUploadDownloadLargeFilesIT.cs

@@ -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);
@@ -85,7 +85,7 @@ private void DownloadFile(string remoteFolderName, string downloadFolderName, st
                 command.ExecuteNonQuery();
             }
         }
-        
+
         private void RemoveFilesFromServer(string remoteFolderName)
         {
             using (var conn = new SnowflakeDbConnection())
@@ -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);

Snowflake.Data.Tests/IntegrationTests/SFPutGetTest.cs

@@ -1,4 +1,4 @@
-/*
+/*
  * Copyright (c) 2012-2023 Snowflake Computing Inc. All rights reserved.
  */
 
@@ -560,16 +560,36 @@
             }
         }
 
+        [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";
@@ -580,7 +600,16 @@
                 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);
 
@@ -610,16 +639,17 @@
         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();
@@ -661,7 +691,7 @@
         }
 
         // 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())
             {
@@ -671,7 +701,8 @@
                         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;
                 }
@@ -696,12 +727,14 @@
         }
 
         // 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;
@@ -764,7 +797,7 @@
         {
             var bytes = new byte[stream.Length];
             stream.Position = 0;
             stream.Read(bytes, 0, (int) stream.Length);
             return Encoding.UTF8.GetString(bytes).Split('\n');
         }
 

Snowflake.Data.Tests/UnitTests/SFFileTransferAgentTests.cs

@@ -1,4 +1,4 @@
-/*
+/*
  * Copyright (c) 2012-2023 Snowflake Computing Inc. All rights reserved.
  */
 
@@ -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()
         {
@@ -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);
+        }
     }
 }

Snowflake.Data/Core/FileTransfer/SFFileTransferAgent.cs

@@ -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;
         }