feat: Implement mode 'truncate' in 'to_sql' API

Author: gmcrocetti

pandas/io/sql.py

@@ -970,11 +970,13 @@ def create(self) -> None:
         if self.exists():
             if self.if_exists == "fail":
                 raise ValueError(f"Table '{self.name}' already exists.")
-            if self.if_exists == "replace":
+            elif self.if_exists == "replace":
                 self.pd_sql.drop_table(self.name, self.schema)
                 self._execute_create()
             elif self.if_exists == "append":
                 pass
+            elif self.if_exists == "truncate":
+                self.pd_sql.truncate_table(self.name, self.schema)
             else:
                 raise ValueError(f"'{self.if_exists}' is not valid for if_exists")
         else:
@@ -2047,6 +2049,25 @@ def drop_table(self, table_name: str, schema: str | None = None) -> None:
                 self.get_table(table_name, schema).drop(bind=self.con)
             self.meta.clear()
 
+    def truncate_table(self, table_name: str, schema: str | None = None) -> None:
+        from sqlalchemy.exc import OperationalError
+
+        schema = schema or self.meta.schema
+
+        if self.has_table(table_name, schema):
+            self.meta.reflect(
+                bind=self.con, only=[table_name], schema=schema, views=True
+            )
+            with self.run_transaction():
+                table = self.get_table(table_name, schema)
+                try:
+                    self.execute(f"TRUNCATE TABLE {table.name}")
+                except OperationalError:
+                    raise NotImplementedError("'TRUNCATE' is not supported by this database.")
+
+            self.meta.clear()
+
+
     def _create_sql_schema(
         self,
         frame: DataFrame,
@@ -2343,6 +2364,8 @@ def to_sql(
         engine : {'auto', 'sqlalchemy'}, default 'auto'
             Raises NotImplementedError if not set to 'auto'
         """
+        from adbc_driver_manager import ProgrammingError
+
         if index_label:
             raise NotImplementedError(
                 "'index_label' is not implemented for ADBC drivers"
@@ -2376,6 +2399,14 @@ def to_sql(
                     cur.execute(f"DROP TABLE {table_name}")
             elif if_exists == "append":
                 mode = "append"
+            elif if_exists == "truncate":
+                mode = "append"
+                with self.con.cursor() as cur:
+                    try:
+                        cur.execute(f"TRUNCATE TABLE {table_name}")
+                    except ProgrammingError:
+                        raise NotImplementedError("'TRUNCATE' is not supported by this database.")
+
 
         import pyarrow as pa
 
@@ -2857,6 +2888,9 @@ def drop_table(self, name: str, schema: str | None = None) -> None:
         drop_sql = f"DROP TABLE {_get_valid_sqlite_name(name)}"
         self.execute(drop_sql)
 
+    def truncate_table(self, name:str, schema: str | None) -> None:
+        raise NotImplementedError("'TRUNCATE' is not supported by this database.")
+
     def _create_sql_schema(
         self,
         frame,

pandas/tests/io/test_sql.py

@@ -1067,12 +1067,19 @@ def test_to_sql(conn, method, test_frame1, request):
 
 
 @pytest.mark.parametrize("conn", all_connectable)
-@pytest.mark.parametrize("mode, num_row_coef", [("replace", 1), ("append", 2)])
+@pytest.mark.parametrize("mode, num_row_coef", [("replace", 1), ("append", 2), ("truncate", 1)])
 def test_to_sql_exist(conn, mode, num_row_coef, test_frame1, request):
+    connections_without_truncate = sqlite_connectable + ["sqlite_buildin", "sqlite_adbc_conn"]
+    if conn in connections_without_truncate and mode == "truncate":
+        context = pytest.raises(NotImplementedError)
+    else:
+        context = contextlib.nullcontext()
     conn = request.getfixturevalue(conn)
+
     with pandasSQL_builder(conn, need_transaction=True) as pandasSQL:
         pandasSQL.to_sql(test_frame1, "test_frame", if_exists="fail")
-        pandasSQL.to_sql(test_frame1, "test_frame", if_exists=mode)
+        with context:
+            pandasSQL.to_sql(test_frame1, "test_frame", if_exists=mode)
         assert pandasSQL.has_table("test_frame")
     assert count_rows(conn, "test_frame") == num_row_coef * len(test_frame1)