Skip to content

Commit aefb841

Browse files
sfc-gh-mkellerandreqaugustosfc-gh-aalam
authored
Test 1792 (#1833)
* 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]>
1 parent 1e50704 commit aefb841

File tree

3 files changed

+27
-3
lines changed

3 files changed

+27
-3
lines changed

DESCRIPTION.md

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -13,6 +13,7 @@ Source code is also available at: https://github.com/snowflakedb/snowflake-conne
1313
- Bumped cryptography dependency from <42.0.0,>=3.1.0 to >=3.1.0,<43.0.0.
1414
- Bumped pyOpenSSL dependency from >=16.2.0,<24.0.0 to >=16.2.0,<25.0.0.
1515
- Fixed a memory leak in decimal data conversion.
16+
- Fixed a bug where `write_pandas` wasn't truncating the target table.
1617

1718
- v3.7.0(January 25,2024)
1819

src/snowflake/connector/pandas_tools.py

Lines changed: 9 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -389,7 +389,7 @@ def drop_object(name: str, object_type: str) -> None:
389389
target_table_location = build_location_helper(
390390
database,
391391
schema,
392-
random_string() if overwrite else table_name,
392+
random_string() if (overwrite and auto_create_table) else table_name,
393393
quote_identifiers,
394394
)
395395

@@ -417,6 +417,11 @@ def drop_object(name: str, object_type: str) -> None:
417417
)
418418

419419
try:
420+
if overwrite and (not auto_create_table):
421+
truncate_sql = f"TRUNCATE TABLE {target_table_location} /* Python:snowflake.connector.pandas_tools.write_pandas() */"
422+
logger.debug(f"truncating table with '{truncate_sql}'")
423+
cursor.execute(truncate_sql, _is_internal=True)
424+
420425
copy_into_sql = (
421426
f"COPY INTO {target_table_location} /* Python:snowflake.connector.pandas_tools.write_pandas() */ "
422427
f"({columns}) "
@@ -432,7 +437,7 @@ def drop_object(name: str, object_type: str) -> None:
432437
logger.debug(f"copying into with '{copy_into_sql}'")
433438
copy_results = cursor.execute(copy_into_sql, _is_internal=True).fetchall()
434439

435-
if overwrite:
440+
if overwrite and auto_create_table:
436441
original_table_location = build_location_helper(
437442
database=database,
438443
schema=schema,
@@ -444,7 +449,8 @@ def drop_object(name: str, object_type: str) -> None:
444449
logger.debug(f"rename table with '{rename_table_sql}'")
445450
cursor.execute(rename_table_sql, _is_internal=True)
446451
except ProgrammingError:
447-
if overwrite:
452+
if overwrite and auto_create_table:
453+
# drop table only if we created a new one with a random name
448454
drop_object(target_table_location, "table")
449455
raise
450456
finally:

test/integ/pandas/test_pandas_tools.py

Lines changed: 17 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -170,6 +170,23 @@ def test_write_pandas_with_overwrite(
170170
if quote_identifiers
171171
else "YEAR" in [col.name for col in result[0].description]
172172
)
173+
else:
174+
# Should fail because the table will be truncated and df3 schema doesn't match
175+
# (since df3 should at least have a subset of the columns of the target table)
176+
with pytest.raises(ProgrammingError, match="invalid identifier"):
177+
write_pandas(
178+
cnx,
179+
df3,
180+
random_table_name,
181+
quote_identifiers=quote_identifiers,
182+
auto_create_table=auto_create_table,
183+
overwrite=True,
184+
index=index,
185+
)
186+
187+
# Check that we have truncated the table but not dropped it in case or error.
188+
result = cnx.cursor(DictCursor).execute(select_count_sql).fetchone()
189+
assert result["COUNT(*)"] == 0
173190

174191
if not quote_identifiers:
175192
original_result = (

0 commit comments

Comments
 (0)