diff --git a/src/snowflake/connector/pandas_tools.py b/src/snowflake/connector/pandas_tools.py index f197cfafa..c0d21b595 100644 --- a/src/snowflake/connector/pandas_tools.py +++ b/src/snowflake/connector/pandas_tools.py @@ -166,7 +166,7 @@ def _create_temp_file_format( def write_pandas( - conn: SnowflakeConnection, + conn: SnowflakeConnection | SnowflakeCursor, df: pandas.DataFrame, table_name: str, database: str | None = None, @@ -216,7 +216,7 @@ def write_pandas( success, nchunks, nrows, _ = write_pandas(cnx, df, 'customers') Args: - conn: Connection to be used to communicate with Snowflake. + conn: Connection or Cursor to be used to communicate with Snowflake. df: Dataframe we'd like to write back. table_name: Table name where we want to insert into. database: Database schema and table is in, if not provided the default one will be used (Default value = None). @@ -315,7 +315,8 @@ def write_pandas( else: sql_use_logical_type = " USE_LOGICAL_TYPE = FALSE" - cursor = conn.cursor() + if isinstance(conn, SnowflakeConnection): + cursor = conn.cursor() stage_location = _create_temp_stage( cursor, database, diff --git a/test/integ/pandas/test_pandas_tools.py b/test/integ/pandas/test_pandas_tools.py index 897b0c969..f47e93617 100644 --- a/test/integ/pandas/test_pandas_tools.py +++ b/test/integ/pandas/test_pandas_tools.py @@ -509,6 +509,20 @@ def test_empty_dataframe_write_pandas( ), f"sucess: {success}, num_chunks: {num_chunks}, num_rows: {num_rows}" +def test_cursor_write_pandas( + conn_cnx: Callable[..., Generator[SnowflakeConnection, None, None]], +): + table_name = random_string(5, "empty_dataframe_") + df = pandas.DataFrame([], columns=["name", "balance"]) + with conn_cnx() as cnx: + success, num_chunks, num_rows, _ = write_pandas( + cnx.cursor(), df, table_name, auto_create_table=True, table_type="temp" + ) + assert ( + success and num_chunks == 1 and num_rows == 0 + ), f"sucess: {success}, num_chunks: {num_chunks}, num_rows: {num_rows}" + + @pytest.mark.parametrize( "database,schema,quote_identifiers,expected_location", [