SNOW-1902019: Python CVEs january batch 2 (#2155)

Co-authored-by: Yijun Xie <[email protected]>

Author: sfc-gh-jszczerbinski

src/snowflake/connector/cursor.py

@@ -875,6 +875,7 @@ def execute(
         _skip_upload_on_content_match: bool = False,
         file_stream: IO[bytes] | None = None,
         num_statements: int | None = None,
+        _force_qmark_paramstyle: bool = False,
         _dataframe_ast: str | None = None,
     ) -> Self | dict[str, Any] | None:
         """Executes a command/query.
@@ -910,6 +911,7 @@ def execute(
             file_stream: File-like object to be uploaded with PUT
             num_statements: Query level parameter submitted in _statement_params constraining exact number of
             statements being submitted (or 0 if submitting an uncounted number) when using a multi-statement query.
+            _force_qmark_paramstyle: Force the use of qmark paramstyle regardless of the connection's paramstyle.
             _dataframe_ast: Base64-encoded dataframe request abstract syntax tree.
 
         Returns:
@@ -958,7 +960,7 @@ def execute(
             "dataframe_ast": _dataframe_ast,
         }
 
-        if self._connection.is_pyformat:
+        if self._connection.is_pyformat and not _force_qmark_paramstyle:
             query = self._preprocess_pyformat_query(command, params)
         else:
             # qmark and numeric paramstyle
@@ -1458,7 +1460,9 @@ def executemany(
         else:
             if re.search(";/s*$", command) is None:
                 command = command + "; "
-            if self._connection.is_pyformat:
+            if self._connection.is_pyformat and not kwargs.get(
+                "_force_qmark_paramstyle", False
+            ):
                 processed_queries = [
                     self._preprocess_pyformat_query(command, params)
                     for params in seqparams

src/snowflake/connector/pandas_tools.py

@@ -85,9 +85,16 @@ def _do_create_temp_stage(
     overwrite: bool,
     use_scoped_temp_object: bool,
 ) -> None:
-    create_stage_sql = f"CREATE {get_temp_type_for_object(use_scoped_temp_object)} STAGE /* Python:snowflake.connector.pandas_tools.write_pandas() */ {stage_location} FILE_FORMAT=(TYPE=PARQUET COMPRESSION={compression}{' BINARY_AS_TEXT=FALSE' if auto_create_table or overwrite else ''})"
-    logger.debug(f"creating stage with '{create_stage_sql}'")
-    cursor.execute(create_stage_sql, _is_internal=True).fetchall()
+    create_stage_sql = f"CREATE {get_temp_type_for_object(use_scoped_temp_object)} STAGE /* Python:snowflake.connector.pandas_tools.write_pandas() */ identifier(?) FILE_FORMAT=(TYPE=PARQUET COMPRESSION={compression}{' BINARY_AS_TEXT=FALSE' if auto_create_table or overwrite else ''})"
+    params = (stage_location,)
+    logger.debug(f"creating stage with '{create_stage_sql}'. params: %s", params)
+    cursor.execute(
+        create_stage_sql,
+        _is_internal=True,
+        _force_qmark_paramstyle=True,
+        params=params,
+        num_statements=1,
+    )
 
 
 def _create_temp_stage(
@@ -147,12 +154,19 @@ def _do_create_temp_file_format(
     use_scoped_temp_object: bool,
 ) -> None:
     file_format_sql = (
-        f"CREATE {get_temp_type_for_object(use_scoped_temp_object)} FILE FORMAT {file_format_location} "
+        f"CREATE {get_temp_type_for_object(use_scoped_temp_object)} FILE FORMAT identifier(?) "
         f"/* Python:snowflake.connector.pandas_tools.write_pandas() */ "
         f"TYPE=PARQUET COMPRESSION={compression}{sql_use_logical_type}"
     )
-    logger.debug(f"creating file format with '{file_format_sql}'")
-    cursor.execute(file_format_sql, _is_internal=True)
+    params = (file_format_location,)
+    logger.debug(f"creating file format with '{file_format_sql}'. params: %s", params)
+    cursor.execute(
+        file_format_sql,
+        _is_internal=True,
+        _force_qmark_paramstyle=True,
+        params=params,
+        num_statements=1,
+    )
 
 
 def _create_temp_file_format(
@@ -379,14 +393,20 @@ def write_pandas(
             # Upload parquet file
             upload_sql = (
                 "PUT /* Python:snowflake.connector.pandas_tools.write_pandas() */ "
-                "'file://{path}' @{stage_location} PARALLEL={parallel}"
+                "'file://{path}' ? PARALLEL={parallel}"
             ).format(
                 path=chunk_path.replace("\\", "\\\\").replace("'", "\\'"),
-                stage_location=stage_location,
                 parallel=parallel,
             )
-            logger.debug(f"uploading files with '{upload_sql}'")
-            cursor.execute(upload_sql, _is_internal=True)
+            params = ("@" + stage_location,)
+            logger.debug(f"uploading files with '{upload_sql}', params: %s", params)
+            cursor.execute(
+                upload_sql,
+                _is_internal=True,
+                _force_qmark_paramstyle=True,
+                params=params,
+                num_statements=1,
+            )
             # Remove chunk file
             os.remove(chunk_path)
 
@@ -403,9 +423,16 @@ def write_pandas(
     columns = quote + f"{quote},{quote}".join(snowflake_column_names) + quote
 
     def drop_object(name: str, object_type: str) -> None:
-        drop_sql = f"DROP {object_type.upper()} IF EXISTS {name} /* Python:snowflake.connector.pandas_tools.write_pandas() */"
-        logger.debug(f"dropping {object_type} with '{drop_sql}'")
-        cursor.execute(drop_sql, _is_internal=True)
+        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,
+            _force_qmark_paramstyle=True,
+            params=params,
+            num_statements=1,
+        )
 
     if auto_create_table or overwrite:
         file_format_location = _create_temp_file_format(
@@ -417,10 +444,17 @@ def drop_object(name: str, object_type: str) -> None:
             sql_use_logical_type,
             _use_scoped_temp_object,
         )
-        infer_schema_sql = f"SELECT COLUMN_NAME, TYPE FROM table(infer_schema(location=>'@{stage_location}', file_format=>'{file_format_location}'))"
-        logger.debug(f"inferring schema with '{infer_schema_sql}'")
+        infer_schema_sql = "SELECT COLUMN_NAME, TYPE FROM table(infer_schema(location=>?, file_format=>?))"
+        params = (f"@{stage_location}", file_format_location)
+        logger.debug(f"inferring schema with '{infer_schema_sql}'. params: %s", params)
         column_type_mapping = dict(
-            cursor.execute(infer_schema_sql, _is_internal=True).fetchall()
+            cursor.execute(
+                infer_schema_sql,
+                _is_internal=True,
+                _force_qmark_paramstyle=True,
+                params=params,
+                num_statements=1,
+            ).fetchall()
         )
         # Infer schema can return the columns out of order depending on the chunking we do when uploading
         # so we have to iterate through the dataframe columns to make sure we create the table with its
@@ -440,12 +474,21 @@ def drop_object(name: str, object_type: str) -> None:
         )
 
         create_table_sql = (
-            f"CREATE {table_type.upper()} TABLE IF NOT EXISTS {target_table_location} "
+            f"CREATE {table_type.upper()} TABLE IF NOT EXISTS identifier(?) "
             f"({create_table_columns})"
             f" /* Python:snowflake.connector.pandas_tools.write_pandas() */ "
         )
-        logger.debug(f"auto creating table with '{create_table_sql}'")
-        cursor.execute(create_table_sql, _is_internal=True)
+        params = (target_table_location,)
+        logger.debug(
+            f"auto creating table with '{create_table_sql}'. params: %s", params
+        )
+        cursor.execute(
+            create_table_sql,
+            _is_internal=True,
+            _force_qmark_paramstyle=True,
+            params=params,
+            num_statements=1,
+        )
         # need explicit casting when the underlying table schema is inferred
         parquet_columns = "$1:" + ",$1:".join(
             f"{quote}{snowflake_col}{quote}::{column_type_mapping[col]}"
@@ -464,12 +507,19 @@ def drop_object(name: str, object_type: str) -> None:
 
     try:
         if overwrite and (not auto_create_table):
-            truncate_sql = f"TRUNCATE TABLE {target_table_location} /* Python:snowflake.connector.pandas_tools.write_pandas() */"
-            logger.debug(f"truncating table with '{truncate_sql}'")
-            cursor.execute(truncate_sql, _is_internal=True)
+            truncate_sql = "TRUNCATE TABLE identifier(?) /* Python:snowflake.connector.pandas_tools.write_pandas() */"
+            params = (target_table_location,)
+            logger.debug(f"truncating table with '{truncate_sql}'. params: %s", params)
+            cursor.execute(
+                truncate_sql,
+                _is_internal=True,
+                _force_qmark_paramstyle=True,
+                params=params,
+                num_statements=1,
+            )
 
         copy_into_sql = (
-            f"COPY INTO {target_table_location} /* Python:snowflake.connector.pandas_tools.write_pandas() */ "
+            f"COPY INTO identifier(?) /* Python:snowflake.connector.pandas_tools.write_pandas() */ "
             f"({columns}) "
             f"FROM (SELECT {parquet_columns} FROM @{stage_location}) "
             f"FILE_FORMAT=("
@@ -478,10 +528,17 @@ def drop_object(name: str, object_type: str) -> None:
             f"{' BINARY_AS_TEXT=FALSE' if auto_create_table or overwrite else ''}"
             f"{sql_use_logical_type}"
             f") "
-            f"PURGE=TRUE ON_ERROR={on_error}"
+            f"PURGE=TRUE ON_ERROR=?"
         )
-        logger.debug(f"copying into with '{copy_into_sql}'")
-        copy_results = cursor.execute(copy_into_sql, _is_internal=True).fetchall()
+        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,
+            _is_internal=True,
+            _force_qmark_paramstyle=True,
+            params=params,
+            num_statements=1,
+        ).fetchall()
 
         if overwrite and auto_create_table:
             original_table_location = build_location_helper(
@@ -491,9 +548,16 @@ def drop_object(name: str, object_type: str) -> None:
                 quote_identifiers=quote_identifiers,
             )
             drop_object(original_table_location, "table")
-            rename_table_sql = f"ALTER TABLE {target_table_location} RENAME TO {original_table_location} /* Python:snowflake.connector.pandas_tools.write_pandas() */"
-            logger.debug(f"rename table with '{rename_table_sql}'")
-            cursor.execute(rename_table_sql, _is_internal=True)
+            rename_table_sql = "ALTER TABLE identifier(?) RENAME TO identifier(?) /* Python:snowflake.connector.pandas_tools.write_pandas() */"
+            params = (target_table_location, original_table_location)
+            logger.debug(f"rename table with '{rename_table_sql}'. params: %s", params)
+            cursor.execute(
+                rename_table_sql,
+                _is_internal=True,
+                _force_qmark_paramstyle=True,
+                params=params,
+                num_statements=1,
+            )
     except ProgrammingError:
         if overwrite and auto_create_table:
             # drop table only if we created a new one with a random name

test/integ/pandas/test_pandas_tools.py

@@ -64,7 +64,7 @@ def assert_result_equals(
 
 
 def test_fix_snow_746341(
-    conn_cnx: Callable[..., Generator[SnowflakeConnection, None, None]]
+    conn_cnx: Callable[..., Generator[SnowflakeConnection, None, None]],
 ):
     cat = '"cat"'
     df = pandas.DataFrame([[1], [2]], columns=[f"col_'{cat}'"])
@@ -534,8 +534,7 @@ def test_table_location_building(
 
         def mocked_execute(*args, **kwargs):
             if len(args) >= 1 and args[0].startswith("COPY INTO"):
-                location = args[0].split(" ")[2]
-                assert location == expected_location
+                assert kwargs["params"][0] == expected_location
             cur = SnowflakeCursor(cnx)
             cur._result = iter([])
             return cur
@@ -906,7 +905,7 @@ def test_auto_create_table_similar_column_names(
 
 
 def test_all_pandas_types(
-    conn_cnx: Callable[..., Generator[SnowflakeConnection, None, None]]
+    conn_cnx: Callable[..., Generator[SnowflakeConnection, None, None]],
 ):
     table_name = random_string(5, "all_types_")
     datetime_with_tz = datetime(1997, 6, 3, 14, 21, 32, 00, tzinfo=timezone.utc)
@@ -997,7 +996,7 @@ def test_no_create_internal_object_privilege_in_target_schema(
             def mock_execute(*args, **kwargs):
                 if (
                     f"CREATE TEMP {object_type}" in args[0]
-                    and "target_schema_no_create_" in args[0]
+                    and "target_schema_no_create_" in kwargs["params"][0]
                 ):
                     raise ProgrammingError("Cannot create temp object in target schema")
                 cursor = cnx.cursor()
@@ -1027,3 +1026,48 @@ def mock_execute(*args, **kwargs):
         finally:
             cnx.execute_string(f"drop schema if exists {source_schema}")
             cnx.execute_string(f"drop schema if exists {target_schema}")
+
+
+def test_write_pandas_with_on_error(
+    conn_cnx: Callable[..., Generator[SnowflakeConnection, None, None]],
+):
+    """Tests whether overwriting table using a Pandas DataFrame works as expected."""
+    random_table_name = random_string(5, "userspoints_")
+    df_data = [("Dash", 50)]
+    df = pandas.DataFrame(df_data, columns=["name", "points"])
+
+    table_name = random_table_name
+    col_id = "id"
+    col_name = "name"
+    col_points = "points"
+
+    create_sql = (
+        f"CREATE OR REPLACE TABLE {table_name}"
+        f"({col_name} STRING, {col_points} INT, {col_id} INT AUTOINCREMENT)"
+    )
+
+    select_count_sql = f"SELECT count(*) FROM {table_name}"
+    drop_sql = f"DROP TABLE IF EXISTS {table_name}"
+    with conn_cnx() as cnx:  # type: SnowflakeConnection
+        cnx.execute_string(create_sql)
+        try:
+            # Write dataframe with 1 row
+            success, nchunks, nrows, _ = write_pandas(
+                cnx,
+                df,
+                random_table_name,
+                quote_identifiers=False,
+                auto_create_table=False,
+                overwrite=True,
+                index=True,
+                on_error="continue",
+            )
+            # Check write_pandas output
+            assert success
+            assert nchunks == 1
+            assert nrows == 1
+            result = cnx.cursor(DictCursor).execute(select_count_sql).fetchone()
+            # Check number of rows
+            assert result["COUNT(*)"] == 1
+        finally:
+            cnx.execute_string(drop_sql)