From aefb84175d96fef6a3f9c2802f1762863cef415c Mon Sep 17 00:00:00 2001
From: Mark Keller <mark.keller@snowflake.com>
Date: Mon, 5 Feb 2024 22:51:11 +0000
Subject: [PATCH] Test 1792 (#1833)
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

* 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 <andre.augusto@trela.com.br>
Co-authored-by: Afroz Alam <afroz.alam@snowflake.com>
---
 DESCRIPTION.md                          |  1 +
 src/snowflake/connector/pandas_tools.py | 12 +++++++++---
 test/integ/pandas/test_pandas_tools.py  | 17 +++++++++++++++++
 3 files changed, 27 insertions(+), 3 deletions(-)

diff --git a/DESCRIPTION.md b/DESCRIPTION.md
index b9f544f6c..04e936616 100644
--- a/DESCRIPTION.md
+++ b/DESCRIPTION.md
@@ -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)
 
diff --git a/src/snowflake/connector/pandas_tools.py b/src/snowflake/connector/pandas_tools.py
index 92ab5d084..f197cfafa 100644
--- a/src/snowflake/connector/pandas_tools.py
+++ b/src/snowflake/connector/pandas_tools.py
@@ -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,
         )
 
@@ -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}) "
@@ -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,
@@ -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:
diff --git a/test/integ/pandas/test_pandas_tools.py b/test/integ/pandas/test_pandas_tools.py
index 5b73a5787..897b0c969 100644
--- a/test/integ/pandas/test_pandas_tools.py
+++ b/test/integ/pandas/test_pandas_tools.py
@@ -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 = (