From 02fd6a40a05f7841ddc8ffb063cb3fb2085a71db Mon Sep 17 00:00:00 2001 From: Oliver Mannion <125105+tekumara@users.noreply.github.com> Date: Mon, 29 May 2023 10:22:07 +1000 Subject: [PATCH] feat: support write_pandas with partial columns --- fakesnow/fakes.py | 7 ++++--- tests/test_fakes.py | 18 ++++++++++++++++++ 2 files changed, 22 insertions(+), 3 deletions(-) diff --git a/fakesnow/fakes.py b/fakesnow/fakes.py index 13fa631..a8d0cd1 100644 --- a/fakesnow/fakes.py +++ b/fakesnow/fakes.py @@ -394,10 +394,11 @@ def execute_string( ] return cursors if return_cursors else [] - def insert_df( + def _insert_df( self, df: pd.DataFrame, table_name: str, database: str | None = None, schema: str | None = None ) -> int: - self._duck_conn.execute(f"INSERT INTO {table_name} SELECT * FROM df") + cols = ",".join(df.columns.to_list()) + self._duck_conn.execute(f"INSERT INTO {table_name}({cols}) SELECT * FROM df") return self._duck_conn.fetchall()[0][0] @@ -463,7 +464,7 @@ def write_pandas( table_type: Literal["", "temp", "temporary", "transient"] = "", **kwargs: Any, ) -> WritePandasResult: - count = conn.insert_df(df, table_name, database, schema) + count = conn._insert_df(df, table_name, database, schema) # noqa: SLF001 # mocks https://docs.snowflake.com/en/sql-reference/sql/copy-into-table.html#output mock_copy_results = [("fakesnow/file0.txt", "LOADED", count, count, 1, 0, None, None, None, None)] diff --git a/tests/test_fakes.py b/tests/test_fakes.py index 57adc15..de6cbb8 100644 --- a/tests/test_fakes.py +++ b/tests/test_fakes.py @@ -491,3 +491,21 @@ def test_write_pandas(conn: snowflake.connector.SnowflakeConnection): cur.execute("select id, first_name, last_name from customers") assert cur.fetchall() == [(1, "Jenny", "P"), (2, "Jasper", "M")] + + +def test_write_pandas_partial_columns(conn: snowflake.connector.SnowflakeConnection): + with conn.cursor() as cur: + cur.execute("create table customers (ID int, FIRST_NAME varchar, LAST_NAME varchar)") + + df = pd.DataFrame.from_records( + [ + {"ID": 1, "FIRST_NAME": "Jenny"}, + {"ID": 2, "FIRST_NAME": "Jasper"}, + ] + ) + snowflake.connector.pandas_tools.write_pandas(conn, df, "customers") + + cur.execute("select id, first_name, last_name from customers") + + # columns not in dataframe will receive their default value + assert cur.fetchall() == [(1, "Jenny", None), (2, "Jasper", None)]