From 3feb537eb66b1638447bd93a0b313189c92d8b38 Mon Sep 17 00:00:00 2001
From: Louie Farol <louie.farol@insightsoftware.com>
Date: Thu, 28 Nov 2024 09:57:12 -0800
Subject: [PATCH] SNOW-1545648: Fix PUT command error if file path contains
 spaces and single quotes

---
 .../UnitTests/SFFileTransferAgentTests.cs       | 17 ++++++++++++++++-
 .../Core/FileTransfer/SFFileTransferAgent.cs    | 14 ++++++++++----
 2 files changed, 26 insertions(+), 5 deletions(-)

diff --git a/Snowflake.Data.Tests/UnitTests/SFFileTransferAgentTests.cs b/Snowflake.Data.Tests/UnitTests/SFFileTransferAgentTests.cs
index 4e7c2041e..93f29e96a 100644
--- a/Snowflake.Data.Tests/UnitTests/SFFileTransferAgentTests.cs
+++ b/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
+        const string FilePathWithoutSpaces = "C:/Users/Test/folder_without_space/*.*";
+        const string FilePathWithSpaces = "C:/Users/Test/folder with space/*.*";
+
         [SetUp]
         public void BeforeEachTest()
         {
@@ -634,5 +638,16 @@ public void TestDownloadThrowsErrorDirectoryNotFound()
             Assert.IsInstanceOf<DirectoryNotFoundException>(innerException);
             Assert.That(innerException?.Message, Does.Match("Could not find a part of the path .*"));
         }
+
+        [Test]
+        [TestCase("PUT file://" + FilePathWithoutSpaces + " @TestStage", FilePathWithoutSpaces)]
+        [TestCase("PUT file://" + FilePathWithSpaces + " @TestStage", FilePathWithSpaces)]
+        [TestCase("PUT 'file://" + FilePathWithoutSpaces + "' @TestStage", FilePathWithoutSpaces)]
+        [TestCase("PUT 'file://" + FilePathWithSpaces + "' @TestStage", FilePathWithSpaces)]
+        public void TestGetFilePathFromPutCommand(string query, string expectedFilePath)
+        {
+            var actualFilePath = SFFileTransferAgent.GetFilePathFromPutCommand(query);
+            Assert.AreEqual(expectedFilePath, actualFilePath);
+        }
     }
 }
diff --git a/Snowflake.Data/Core/FileTransfer/SFFileTransferAgent.cs b/Snowflake.Data/Core/FileTransfer/SFFileTransferAgent.cs
index 1ee8557b6..bbe16af0b 100644
--- a/Snowflake.Data/Core/FileTransfer/SFFileTransferAgent.cs
+++ b/Snowflake.Data/Core/FileTransfer/SFFileTransferAgent.cs
@@ -1,4 +1,4 @@
-/*
+/*
  * Copyright (c) 2021 Snowflake Computing Inc. All rights reserved.
  */
 
@@ -451,7 +451,7 @@ private void updatePresignedUrl()
                 {
                     foreach (SFFileMetadata fileMeta in FilesMetas)
                     {
-                        string filePathToReplace = getFilePathFromPutCommand(Query);
+                        string filePathToReplace = GetFilePathFromPutCommand(Query);
                         string fileNameToReplaceWith = fileMeta.destFileName;
                         string queryWithSingleFile = Query;
                         queryWithSingleFile = queryWithSingleFile.Replace(filePathToReplace, fileNameToReplaceWith);
@@ -492,7 +492,7 @@ internal async Task updatePresignedUrlAsync(CancellationToken cancellationToken)
                 {
                     foreach (SFFileMetadata fileMeta in FilesMetas)
                     {
-                        string filePathToReplace = getFilePathFromPutCommand(Query);
+                        string filePathToReplace = GetFilePathFromPutCommand(Query);
                         string fileNameToReplaceWith = fileMeta.destFileName;
                         string queryWithSingleFile = Query;
                         queryWithSingleFile = queryWithSingleFile.Replace(filePathToReplace, fileNameToReplaceWith);
@@ -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);
+
+            // Check if file path contains an enclosing (') char
+            if (filePath[filePath.Length - 1] == '\'')
+            {
+                filePath = filePath.Substring(0, filePath.Length - 1);
+            }
             return filePath;
         }