Skip to content

Commit

Permalink
Test 1792 (#1833)
Browse files Browse the repository at this point in the history
* fixing the `auto_create_table` behavior with `overwrite` option

* adding a test

* pre-commit fix

* adding changelog

* SNOW-733213: fix test

* drop table in case of failure only if we created one using random_string

* revert unwanted change

* changelog updates

---------

Co-authored-by: André Augusto <[email protected]>
Co-authored-by: Afroz Alam <[email protected]>
  • Loading branch information
3 people authored Feb 5, 2024
1 parent 1e50704 commit aefb841
Show file tree
Hide file tree
Showing 3 changed files with 27 additions and 3 deletions.
1 change: 1 addition & 0 deletions DESCRIPTION.md
Original file line number Diff line number Diff line change
Expand Up @@ -13,6 +13,7 @@ Source code is also available at: https://github.com/snowflakedb/snowflake-conne
- Bumped cryptography dependency from <42.0.0,>=3.1.0 to >=3.1.0,<43.0.0.
- Bumped pyOpenSSL dependency from >=16.2.0,<24.0.0 to >=16.2.0,<25.0.0.
- Fixed a memory leak in decimal data conversion.
- Fixed a bug where `write_pandas` wasn't truncating the target table.

- v3.7.0(January 25,2024)

Expand Down
12 changes: 9 additions & 3 deletions src/snowflake/connector/pandas_tools.py
Original file line number Diff line number Diff line change
Expand Up @@ -389,7 +389,7 @@ def drop_object(name: str, object_type: str) -> None:
target_table_location = build_location_helper(
database,
schema,
random_string() if overwrite else table_name,
random_string() if (overwrite and auto_create_table) else table_name,
quote_identifiers,
)

Expand Down Expand Up @@ -417,6 +417,11 @@ 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)

copy_into_sql = (
f"COPY INTO {target_table_location} /* Python:snowflake.connector.pandas_tools.write_pandas() */ "
f"({columns}) "
Expand All @@ -432,7 +437,7 @@ def drop_object(name: str, object_type: str) -> None:
logger.debug(f"copying into with '{copy_into_sql}'")
copy_results = cursor.execute(copy_into_sql, _is_internal=True).fetchall()

if overwrite:
if overwrite and auto_create_table:
original_table_location = build_location_helper(
database=database,
schema=schema,
Expand All @@ -444,7 +449,8 @@ def drop_object(name: str, object_type: str) -> None:
logger.debug(f"rename table with '{rename_table_sql}'")
cursor.execute(rename_table_sql, _is_internal=True)
except ProgrammingError:
if overwrite:
if overwrite and auto_create_table:
# drop table only if we created a new one with a random name
drop_object(target_table_location, "table")
raise
finally:
Expand Down
17 changes: 17 additions & 0 deletions test/integ/pandas/test_pandas_tools.py
Original file line number Diff line number Diff line change
Expand Up @@ -170,6 +170,23 @@ def test_write_pandas_with_overwrite(
if quote_identifiers
else "YEAR" in [col.name for col in result[0].description]
)
else:
# Should fail because the table will be truncated and df3 schema doesn't match
# (since df3 should at least have a subset of the columns of the target table)
with pytest.raises(ProgrammingError, match="invalid identifier"):
write_pandas(
cnx,
df3,
random_table_name,
quote_identifiers=quote_identifiers,
auto_create_table=auto_create_table,
overwrite=True,
index=index,
)

# Check that we have truncated the table but not dropped it in case or error.
result = cnx.cursor(DictCursor).execute(select_count_sql).fetchone()
assert result["COUNT(*)"] == 0

if not quote_identifiers:
original_result = (
Expand Down

0 comments on commit aefb841

Please sign in to comment.