SNOW-1959514: Pandas single quote character fix (#2307)

(cherry picked from commit ecd5d9f)

Author: sfc-gh-pbulawa

src/snowflake/connector/cursor.py

@@ -672,7 +672,10 @@ def _execute_helper(
         else:
             # or detect it.
             self._is_file_transfer = get_file_transfer_type(query) is not None
-        logger.debug("is_file_transfer: %s", self._is_file_transfer is not None)
+        logger.debug(
+            "is_file_transfer: %s",
+            self._is_file_transfer if self._is_file_transfer is not None else "None",
+        )
 
         real_timeout = (
             timeout if timeout and timeout > 0 else self._connection.network_timeout

src/snowflake/connector/pandas_tools.py

@@ -58,21 +58,26 @@ def build_location_helper(
     database: str | None, schema: str | None, name: str, quote_identifiers: bool
 ) -> str:
     """Helper to format table/stage/file format's location."""
-    if quote_identifiers:
-        location = (
-            (('"' + database + '".') if database else "")
-            + (('"' + schema + '".') if schema else "")
-            + ('"' + name + '"')
-        )
-    else:
-        location = (
-            (database + "." if database else "")
-            + (schema + "." if schema else "")
-            + name
-        )
+    location = (
+        (_escape_part_location(database, quote_identifiers) + "." if database else "")
+        + (_escape_part_location(schema, quote_identifiers) + "." if schema else "")
+        + _escape_part_location(name, quote_identifiers)
+    )
     return location
 
 
+def _escape_part_location(part: str, should_quote: bool) -> str:
+    if "'" in part:
+        should_quote = True
+    if should_quote:
+        if not part.startswith('"'):
+            part = '"' + part
+        if not part.endswith('"'):
+            part = part + '"'
+
+    return part
+
+
 def _do_create_temp_stage(
     cursor: SnowflakeCursor,
     stage_location: str,
@@ -473,6 +478,7 @@ def drop_object(name: str, object_type: str) -> None:
         drop_sql = f"DROP {object_type.upper()} IF EXISTS identifier(?) /* Python:snowflake.connector.pandas_tools.write_pandas() */"
         params = (name,)
         logger.debug(f"dropping {object_type} with '{drop_sql}'. params: %s", params)
+
         cursor.execute(
             drop_sql,
             _is_internal=True,
@@ -570,10 +576,11 @@ def drop_object(name: str, object_type: str) -> None:
                 num_statements=1,
             )
 
+        copy_stage_location = "@" + stage_location.replace("'", "\\'")
         copy_into_sql = (
             f"COPY INTO identifier(?) /* Python:snowflake.connector.pandas_tools.write_pandas() */ "
             f"({columns}) "
-            f"FROM (SELECT {parquet_columns} FROM @{stage_location}) "
+            f"FROM (SELECT {parquet_columns} FROM '{copy_stage_location}') "
             f"FILE_FORMAT=("
             f"TYPE=PARQUET "
             f"COMPRESSION={compression_map[compression]}"
@@ -582,7 +589,10 @@ def drop_object(name: str, object_type: str) -> None:
             f") "
             f"PURGE=TRUE ON_ERROR=?"
         )
-        params = (target_table_location, on_error)
+        params = (
+            target_table_location,
+            on_error,
+        )
         logger.debug(f"copying into with '{copy_into_sql}'. params: %s", params)
         copy_results = cursor.execute(
             copy_into_sql,

test/integ/pandas/test_pandas_tools.py

@@ -566,6 +566,8 @@ def mocked_execute(*args, **kwargs):
         (None, "schema", False, "schema"),
         (None, None, True, ""),
         (None, None, False, ""),
+        ("data'base", "schema", True, '"data\'base"."schema"'),
+        ("data'base", "schema", False, '"data\'base".schema'),
     ],
 )
 def test_stage_location_building(
@@ -1101,3 +1103,26 @@ def test_write_pandas_with_on_error(
             assert result["COUNT(*)"] == 1
         finally:
             cnx.execute_string(drop_sql)
+
+
+def test_pandas_with_single_quote(
+    conn_cnx: Callable[..., Generator[SnowflakeConnection]],
+):
+    random_table_name = random_string(5, "test'table")
+    table_name = f'"{random_table_name}"'
+    create_sql = f"CREATE OR REPLACE TABLE {table_name}(A INT)"
+    df_data = [[1]]
+    df = pandas.DataFrame(df_data, columns=["a"])
+    with conn_cnx() as cnx:  # type: SnowflakeConnection
+        try:
+            cnx.execute_string(create_sql)
+            write_pandas(
+                cnx,
+                df,
+                table_name,
+                quote_identifiers=False,
+                auto_create_table=False,
+                index=False,
+            )
+        finally:
+            cnx.execute_string(f"drop table if exists {table_name}")