Skip to content

Commit 4e93848

Browse files
committed
feat: Implement option 'truncate' of argument 'if_exists' in 'DataFrame.to_sql' API.
1 parent 642d244 commit 4e93848

File tree

3 files changed

+114
-13
lines changed

3 files changed

+114
-13
lines changed

Diff for: doc/source/whatsnew/v3.0.0.rst

+1
Original file line numberDiff line numberDiff line change
@@ -50,6 +50,7 @@ Other enhancements
5050
- :meth:`DataFrame.pivot_table` and :func:`pivot_table` now allow the passing of keyword arguments to ``aggfunc`` through ``**kwargs`` (:issue:`57884`)
5151
- :meth:`Series.cummin` and :meth:`Series.cummax` now supports :class:`CategoricalDtype` (:issue:`52335`)
5252
- :meth:`Series.plot` now correctly handle the ``ylabel`` parameter for pie charts, allowing for explicit control over the y-axis label (:issue:`58239`)
53+
- Add ``"truncate"`` option to ``if_exists`` argument in :meth:`DataFrame.to_sql` truncating the table before inserting data (:issue:`37210`).
5354
- Restore support for reading Stata 104-format and enable reading 103-format dta files (:issue:`58554`)
5455
- Support reading Stata 102-format (Stata 1) dta files (:issue:`58978`)
5556
- Support reading Stata 110-format (Stata 7) dta files (:issue:`47176`)

Diff for: pandas/io/sql.py

+55-11
Original file line numberDiff line numberDiff line change
@@ -733,7 +733,7 @@ def to_sql(
733733
name: str,
734734
con,
735735
schema: str | None = None,
736-
if_exists: Literal["fail", "replace", "append"] = "fail",
736+
if_exists: Literal["fail", "replace", "append", "truncate"] = "fail",
737737
index: bool = True,
738738
index_label: IndexLabel | None = None,
739739
chunksize: int | None = None,
@@ -759,10 +759,12 @@ def to_sql(
759759
schema : str, optional
760760
Name of SQL schema in database to write to (if database flavor
761761
supports this). If None, use default schema (default).
762-
if_exists : {'fail', 'replace', 'append'}, default 'fail'
762+
if_exists : {'fail', 'replace', 'append', 'truncate'}, default 'fail'
763763
- fail: If table exists, do nothing.
764764
- replace: If table exists, drop it, recreate it, and insert data.
765765
- append: If table exists, insert data. Create if does not exist.
766+
- truncate: If table exists, truncate it. Create if does not exist.
767+
Raises NotImplementedError if 'TRUNCATE TABLE' is not supported
766768
index : bool, default True
767769
Write DataFrame index as a column.
768770
index_label : str or sequence, optional
@@ -813,7 +815,7 @@ def to_sql(
813815
`sqlite3 <https://docs.python.org/3/library/sqlite3.html#sqlite3.Cursor.rowcount>`__ or
814816
`SQLAlchemy <https://docs.sqlalchemy.org/en/14/core/connections.html#sqlalchemy.engine.BaseCursorResult.rowcount>`__
815817
""" # noqa: E501
816-
if if_exists not in ("fail", "replace", "append"):
818+
if if_exists not in ("fail", "replace", "append", "truncate"):
817819
raise ValueError(f"'{if_exists}' is not valid for if_exists")
818820

819821
if isinstance(frame, Series):
@@ -921,7 +923,7 @@ def __init__(
921923
pandas_sql_engine,
922924
frame=None,
923925
index: bool | str | list[str] | None = True,
924-
if_exists: Literal["fail", "replace", "append"] = "fail",
926+
if_exists: Literal["fail", "replace", "append", "truncate"] = "fail",
925927
prefix: str = "pandas",
926928
index_label=None,
927929
schema=None,
@@ -969,11 +971,13 @@ def create(self) -> None:
969971
if self.exists():
970972
if self.if_exists == "fail":
971973
raise ValueError(f"Table '{self.name}' already exists.")
972-
if self.if_exists == "replace":
974+
elif self.if_exists == "replace":
973975
self.pd_sql.drop_table(self.name, self.schema)
974976
self._execute_create()
975977
elif self.if_exists == "append":
976978
pass
979+
elif self.if_exists == "truncate":
980+
self.pd_sql.truncate_table(self.name, self.schema)
977981
else:
978982
raise ValueError(f"'{self.if_exists}' is not valid for if_exists")
979983
else:
@@ -1465,7 +1469,7 @@ def to_sql(
14651469
self,
14661470
frame,
14671471
name: str,
1468-
if_exists: Literal["fail", "replace", "append"] = "fail",
1472+
if_exists: Literal["fail", "replace", "append", "truncate"] = "fail",
14691473
index: bool = True,
14701474
index_label=None,
14711475
schema=None,
@@ -1850,7 +1854,7 @@ def prep_table(
18501854
self,
18511855
frame,
18521856
name: str,
1853-
if_exists: Literal["fail", "replace", "append"] = "fail",
1857+
if_exists: Literal["fail", "replace", "append", "truncate"] = "fail",
18541858
index: bool | str | list[str] | None = True,
18551859
index_label=None,
18561860
schema=None,
@@ -1927,7 +1931,7 @@ def to_sql(
19271931
self,
19281932
frame,
19291933
name: str,
1930-
if_exists: Literal["fail", "replace", "append"] = "fail",
1934+
if_exists: Literal["fail", "replace", "append", "truncate"] = "fail",
19311935
index: bool = True,
19321936
index_label=None,
19331937
schema: str | None = None,
@@ -1945,10 +1949,12 @@ def to_sql(
19451949
frame : DataFrame
19461950
name : string
19471951
Name of SQL table.
1948-
if_exists : {'fail', 'replace', 'append'}, default 'fail'
1952+
if_exists : {'fail', 'replace', 'append', 'truncate'}, default 'fail'
19491953
- fail: If table exists, do nothing.
19501954
- replace: If table exists, drop it, recreate it, and insert data.
19511955
- append: If table exists, insert data. Create if does not exist.
1956+
- truncate: If table exists, truncate it. Create if does not exist.
1957+
Raises NotImplementedError if 'TRUNCATE TABLE' is not supported
19521958
index : boolean, default True
19531959
Write DataFrame index as a column.
19541960
index_label : string or sequence, default None
@@ -2045,6 +2051,26 @@ def drop_table(self, table_name: str, schema: str | None = None) -> None:
20452051
self.get_table(table_name, schema).drop(bind=self.con)
20462052
self.meta.clear()
20472053

2054+
def truncate_table(self, table_name: str, schema: str | None = None) -> None:
2055+
from sqlalchemy.exc import OperationalError
2056+
2057+
schema = schema or self.meta.schema
2058+
2059+
if self.has_table(table_name, schema):
2060+
self.meta.reflect(
2061+
bind=self.con, only=[table_name], schema=schema, views=True
2062+
)
2063+
with self.run_transaction():
2064+
table = self.get_table(table_name, schema)
2065+
try:
2066+
self.execute(f"TRUNCATE TABLE {table.name}")
2067+
except OperationalError as exc:
2068+
raise NotImplementedError(
2069+
"'TRUNCATE TABLE' is not supported by this database."
2070+
) from exc
2071+
2072+
self.meta.clear()
2073+
20482074
def _create_sql_schema(
20492075
self,
20502076
frame: DataFrame,
@@ -2301,7 +2327,7 @@ def to_sql(
23012327
self,
23022328
frame,
23032329
name: str,
2304-
if_exists: Literal["fail", "replace", "append"] = "fail",
2330+
if_exists: Literal["fail", "replace", "append", "truncate"] = "fail",
23052331
index: bool = True,
23062332
index_label=None,
23072333
schema: str | None = None,
@@ -2323,6 +2349,8 @@ def to_sql(
23232349
- fail: If table exists, do nothing.
23242350
- replace: If table exists, drop it, recreate it, and insert data.
23252351
- append: If table exists, insert data. Create if does not exist.
2352+
- truncate: If table exists, truncate it. Create if does not exist.
2353+
Raises NotImplementedError if 'TRUNCATE TABLE' is not supported
23262354
index : boolean, default True
23272355
Write DataFrame index as a column.
23282356
index_label : string or sequence, default None
@@ -2340,6 +2368,8 @@ def to_sql(
23402368
engine : {'auto', 'sqlalchemy'}, default 'auto'
23412369
Raises NotImplementedError if not set to 'auto'
23422370
"""
2371+
from adbc_driver_manager import ProgrammingError
2372+
23432373
if index_label:
23442374
raise NotImplementedError(
23452375
"'index_label' is not implemented for ADBC drivers"
@@ -2373,6 +2403,15 @@ def to_sql(
23732403
cur.execute(f"DROP TABLE {table_name}")
23742404
elif if_exists == "append":
23752405
mode = "append"
2406+
elif if_exists == "truncate":
2407+
mode = "append"
2408+
with self.con.cursor() as cur:
2409+
try:
2410+
cur.execute(f"TRUNCATE TABLE {table_name}")
2411+
except ProgrammingError as exc:
2412+
raise NotImplementedError(
2413+
"'TRUNCATE TABLE' is not supported by this database."
2414+
) from exc
23762415

23772416
import pyarrow as pa
23782417

@@ -2774,10 +2813,12 @@ def to_sql(
27742813
frame: DataFrame
27752814
name: string
27762815
Name of SQL table.
2777-
if_exists: {'fail', 'replace', 'append'}, default 'fail'
2816+
if_exists: {'fail', 'replace', 'append', 'truncate'}, default 'fail'
27782817
fail: If table exists, do nothing.
27792818
replace: If table exists, drop it, recreate it, and insert data.
27802819
append: If table exists, insert data. Create if it does not exist.
2820+
truncate: If table exists, truncate it. Create if does not exist.
2821+
Raises NotImplementedError if 'TRUNCATE TABLE' is not supported
27812822
index : bool, default True
27822823
Write DataFrame index as a column
27832824
index_label : string or sequence, default None
@@ -2853,6 +2894,9 @@ def drop_table(self, name: str, schema: str | None = None) -> None:
28532894
drop_sql = f"DROP TABLE {_get_valid_sqlite_name(name)}"
28542895
self.execute(drop_sql)
28552896

2897+
def truncate_table(self, name: str, schema: str | None = None) -> None:
2898+
raise NotImplementedError("'TRUNCATE TABLE' is not supported by this database.")
2899+
28562900
def _create_sql_schema(
28572901
self,
28582902
frame,

Diff for: pandas/tests/io/test_sql.py

+58-2
Original file line numberDiff line numberDiff line change
@@ -1067,12 +1067,27 @@ def test_to_sql(conn, method, test_frame1, request):
10671067

10681068

10691069
@pytest.mark.parametrize("conn", all_connectable)
1070-
@pytest.mark.parametrize("mode, num_row_coef", [("replace", 1), ("append", 2)])
1070+
@pytest.mark.parametrize(
1071+
"mode, num_row_coef", [("replace", 1), ("append", 2), ("truncate", 1)]
1072+
)
10711073
def test_to_sql_exist(conn, mode, num_row_coef, test_frame1, request):
1074+
connections_without_truncate = sqlite_connectable + [
1075+
"sqlite_buildin",
1076+
"sqlite_adbc_conn",
1077+
]
1078+
if conn in connections_without_truncate and mode == "truncate":
1079+
context = pytest.raises(
1080+
NotImplementedError,
1081+
match="'TRUNCATE TABLE' is not supported by this database.",
1082+
)
1083+
else:
1084+
context = contextlib.nullcontext()
10721085
conn = request.getfixturevalue(conn)
1086+
10731087
with pandasSQL_builder(conn, need_transaction=True) as pandasSQL:
10741088
pandasSQL.to_sql(test_frame1, "test_frame", if_exists="fail")
1075-
pandasSQL.to_sql(test_frame1, "test_frame", if_exists=mode)
1089+
with context:
1090+
pandasSQL.to_sql(test_frame1, "test_frame", if_exists=mode)
10761091
assert pandasSQL.has_table("test_frame")
10771092
assert count_rows(conn, "test_frame") == num_row_coef * len(test_frame1)
10781093

@@ -2697,6 +2712,47 @@ def test_drop_table(conn, request):
26972712
assert not insp.has_table("temp_frame")
26982713

26992714

2715+
@pytest.mark.parametrize("conn", mysql_connectable + postgresql_connectable)
2716+
def test_truncate_table_success(conn, test_frame1, request):
2717+
table_name = "temp_frame"
2718+
conn = request.getfixturevalue(conn)
2719+
2720+
with sql.SQLDatabase(conn) as pandasSQL:
2721+
with pandasSQL.run_transaction():
2722+
assert pandasSQL.to_sql(test_frame1, table_name, if_exists="replace") == 4
2723+
2724+
with pandasSQL.run_transaction():
2725+
pandasSQL.truncate_table(table_name)
2726+
assert count_rows(conn, table_name) == 0
2727+
2728+
2729+
@pytest.mark.parametrize("conn", sqlite_connectable)
2730+
def test_truncate_table_not_supported(conn, test_frame1, request):
2731+
table_name = "temp_frame"
2732+
conn = request.getfixturevalue(conn)
2733+
2734+
with sql.SQLDatabase(conn) as pandasSQL:
2735+
with pandasSQL.run_transaction():
2736+
assert pandasSQL.to_sql(test_frame1, table_name, if_exists="replace") == 4
2737+
2738+
with pandasSQL.run_transaction():
2739+
with pytest.raises(
2740+
NotImplementedError,
2741+
match="'TRUNCATE TABLE' is not supported by this database.",
2742+
):
2743+
pandasSQL.truncate_table(table_name)
2744+
assert count_rows(conn, table_name) == len(test_frame1)
2745+
2746+
2747+
def test_truncate_table_sqlite(sqlite_buildin):
2748+
with sql.SQLiteDatabase(sqlite_buildin) as pandasSQL:
2749+
with pytest.raises(
2750+
NotImplementedError,
2751+
match="'TRUNCATE TABLE' is not supported by this database.",
2752+
):
2753+
pandasSQL.truncate_table("table")
2754+
2755+
27002756
@pytest.mark.parametrize("conn", all_connectable)
27012757
def test_roundtrip(conn, request, test_frame1):
27022758
if conn == "sqlite_str":

0 commit comments

Comments
 (0)