Skip to content

Commit bce011a

Browse files
committed
feat: Implement option 'truncate' of argument 'if_exists' in 'DataFrame.to_sql' API.
1 parent 078e11f commit bce011a

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
@@ -59,6 +59,7 @@ Other enhancements
5959
- Support passing a :class:`Iterable[Hashable]` input to :meth:`DataFrame.drop_duplicates` (:issue:`59237`)
6060
- Support reading Stata 102-format (Stata 1) dta files (:issue:`58978`)
6161
- Support reading Stata 110-format (Stata 7) dta files (:issue:`47176`)
62+
- Add ``"truncate"`` option to ``if_exists`` argument in :meth:`DataFrame.to_sql` truncating the table before inserting data (:issue:`37210`).
6263

6364
.. ---------------------------------------------------------------------------
6465
.. _whatsnew_300.notable_bug_fixes:

Diff for: pandas/io/sql.py

+55-11
Original file line numberDiff line numberDiff line change
@@ -736,7 +736,7 @@ def to_sql(
736736
name: str,
737737
con,
738738
schema: str | None = None,
739-
if_exists: Literal["fail", "replace", "append"] = "fail",
739+
if_exists: Literal["fail", "replace", "append", "truncate"] = "fail",
740740
index: bool = True,
741741
index_label: IndexLabel | None = None,
742742
chunksize: int | None = None,
@@ -762,10 +762,12 @@ def to_sql(
762762
schema : str, optional
763763
Name of SQL schema in database to write to (if database flavor
764764
supports this). If None, use default schema (default).
765-
if_exists : {'fail', 'replace', 'append'}, default 'fail'
765+
if_exists : {'fail', 'replace', 'append', 'truncate'}, default 'fail'
766766
- fail: If table exists, do nothing.
767767
- replace: If table exists, drop it, recreate it, and insert data.
768768
- append: If table exists, insert data. Create if does not exist.
769+
- truncate: If table exists, truncate it. Create if does not exist.
770+
Raises NotImplementedError if 'TRUNCATE TABLE' is not supported
769771
index : bool, default True
770772
Write DataFrame index as a column.
771773
index_label : str or sequence, optional
@@ -816,7 +818,7 @@ def to_sql(
816818
`sqlite3 <https://docs.python.org/3/library/sqlite3.html#sqlite3.Cursor.rowcount>`__ or
817819
`SQLAlchemy <https://docs.sqlalchemy.org/en/14/core/connections.html#sqlalchemy.engine.BaseCursorResult.rowcount>`__
818820
""" # noqa: E501
819-
if if_exists not in ("fail", "replace", "append"):
821+
if if_exists not in ("fail", "replace", "append", "truncate"):
820822
raise ValueError(f"'{if_exists}' is not valid for if_exists")
821823

822824
if isinstance(frame, Series):
@@ -924,7 +926,7 @@ def __init__(
924926
pandas_sql_engine,
925927
frame=None,
926928
index: bool | str | list[str] | None = True,
927-
if_exists: Literal["fail", "replace", "append"] = "fail",
929+
if_exists: Literal["fail", "replace", "append", "truncate"] = "fail",
928930
prefix: str = "pandas",
929931
index_label=None,
930932
schema=None,
@@ -972,11 +974,13 @@ def create(self) -> None:
972974
if self.exists():
973975
if self.if_exists == "fail":
974976
raise ValueError(f"Table '{self.name}' already exists.")
975-
if self.if_exists == "replace":
977+
elif self.if_exists == "replace":
976978
self.pd_sql.drop_table(self.name, self.schema)
977979
self._execute_create()
978980
elif self.if_exists == "append":
979981
pass
982+
elif self.if_exists == "truncate":
983+
self.pd_sql.truncate_table(self.name, self.schema)
980984
else:
981985
raise ValueError(f"'{self.if_exists}' is not valid for if_exists")
982986
else:
@@ -1468,7 +1472,7 @@ def to_sql(
14681472
self,
14691473
frame,
14701474
name: str,
1471-
if_exists: Literal["fail", "replace", "append"] = "fail",
1475+
if_exists: Literal["fail", "replace", "append", "truncate"] = "fail",
14721476
index: bool = True,
14731477
index_label=None,
14741478
schema=None,
@@ -1854,7 +1858,7 @@ def prep_table(
18541858
self,
18551859
frame,
18561860
name: str,
1857-
if_exists: Literal["fail", "replace", "append"] = "fail",
1861+
if_exists: Literal["fail", "replace", "append", "truncate"] = "fail",
18581862
index: bool | str | list[str] | None = True,
18591863
index_label=None,
18601864
schema=None,
@@ -1931,7 +1935,7 @@ def to_sql(
19311935
self,
19321936
frame,
19331937
name: str,
1934-
if_exists: Literal["fail", "replace", "append"] = "fail",
1938+
if_exists: Literal["fail", "replace", "append", "truncate"] = "fail",
19351939
index: bool = True,
19361940
index_label=None,
19371941
schema: str | None = None,
@@ -1949,10 +1953,12 @@ def to_sql(
19491953
frame : DataFrame
19501954
name : string
19511955
Name of SQL table.
1952-
if_exists : {'fail', 'replace', 'append'}, default 'fail'
1956+
if_exists : {'fail', 'replace', 'append', 'truncate'}, default 'fail'
19531957
- fail: If table exists, do nothing.
19541958
- replace: If table exists, drop it, recreate it, and insert data.
19551959
- append: If table exists, insert data. Create if does not exist.
1960+
- truncate: If table exists, truncate it. Create if does not exist.
1961+
Raises NotImplementedError if 'TRUNCATE TABLE' is not supported
19561962
index : boolean, default True
19571963
Write DataFrame index as a column.
19581964
index_label : string or sequence, default None
@@ -2049,6 +2055,26 @@ def drop_table(self, table_name: str, schema: str | None = None) -> None:
20492055
self.get_table(table_name, schema).drop(bind=self.con)
20502056
self.meta.clear()
20512057

2058+
def truncate_table(self, table_name: str, schema: str | None = None) -> None:
2059+
from sqlalchemy.exc import OperationalError
2060+
2061+
schema = schema or self.meta.schema
2062+
2063+
if self.has_table(table_name, schema):
2064+
self.meta.reflect(
2065+
bind=self.con, only=[table_name], schema=schema, views=True
2066+
)
2067+
with self.run_transaction():
2068+
table = self.get_table(table_name, schema)
2069+
try:
2070+
self.execute(f"TRUNCATE TABLE {table.name}")
2071+
except OperationalError as exc:
2072+
raise NotImplementedError(
2073+
"'TRUNCATE TABLE' is not supported by this database."
2074+
) from exc
2075+
2076+
self.meta.clear()
2077+
20522078
def _create_sql_schema(
20532079
self,
20542080
frame: DataFrame,
@@ -2306,7 +2332,7 @@ def to_sql(
23062332
self,
23072333
frame,
23082334
name: str,
2309-
if_exists: Literal["fail", "replace", "append"] = "fail",
2335+
if_exists: Literal["fail", "replace", "append", "truncate"] = "fail",
23102336
index: bool = True,
23112337
index_label=None,
23122338
schema: str | None = None,
@@ -2328,6 +2354,8 @@ def to_sql(
23282354
- fail: If table exists, do nothing.
23292355
- replace: If table exists, drop it, recreate it, and insert data.
23302356
- append: If table exists, insert data. Create if does not exist.
2357+
- truncate: If table exists, truncate it. Create if does not exist.
2358+
Raises NotImplementedError if 'TRUNCATE TABLE' is not supported
23312359
index : boolean, default True
23322360
Write DataFrame index as a column.
23332361
index_label : string or sequence, default None
@@ -2345,6 +2373,8 @@ def to_sql(
23452373
engine : {'auto', 'sqlalchemy'}, default 'auto'
23462374
Raises NotImplementedError if not set to 'auto'
23472375
"""
2376+
from adbc_driver_manager import ProgrammingError
2377+
23482378
if index_label:
23492379
raise NotImplementedError(
23502380
"'index_label' is not implemented for ADBC drivers"
@@ -2378,6 +2408,15 @@ def to_sql(
23782408
cur.execute(f"DROP TABLE {table_name}")
23792409
elif if_exists == "append":
23802410
mode = "append"
2411+
elif if_exists == "truncate":
2412+
mode = "append"
2413+
with self.con.cursor() as cur:
2414+
try:
2415+
cur.execute(f"TRUNCATE TABLE {table_name}")
2416+
except ProgrammingError as exc:
2417+
raise NotImplementedError(
2418+
"'TRUNCATE TABLE' is not supported by this database."
2419+
) from exc
23812420

23822421
import pyarrow as pa
23832422

@@ -2779,10 +2818,12 @@ def to_sql(
27792818
frame: DataFrame
27802819
name: string
27812820
Name of SQL table.
2782-
if_exists: {'fail', 'replace', 'append'}, default 'fail'
2821+
if_exists: {'fail', 'replace', 'append', 'truncate'}, default 'fail'
27832822
fail: If table exists, do nothing.
27842823
replace: If table exists, drop it, recreate it, and insert data.
27852824
append: If table exists, insert data. Create if it does not exist.
2825+
truncate: If table exists, truncate it. Create if does not exist.
2826+
Raises NotImplementedError if 'TRUNCATE TABLE' is not supported
27862827
index : bool, default True
27872828
Write DataFrame index as a column
27882829
index_label : string or sequence, default None
@@ -2858,6 +2899,9 @@ def drop_table(self, name: str, schema: str | None = None) -> None:
28582899
drop_sql = f"DROP TABLE {_get_valid_sqlite_name(name)}"
28592900
self.execute(drop_sql)
28602901

2902+
def truncate_table(self, name: str, schema: str | None = None) -> None:
2903+
raise NotImplementedError("'TRUNCATE TABLE' is not supported by this database.")
2904+
28612905
def _create_sql_schema(
28622906
self,
28632907
frame,

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

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

10641064

10651065
@pytest.mark.parametrize("conn", all_connectable)
1066-
@pytest.mark.parametrize("mode, num_row_coef", [("replace", 1), ("append", 2)])
1066+
@pytest.mark.parametrize(
1067+
"mode, num_row_coef", [("replace", 1), ("append", 2), ("truncate", 1)]
1068+
)
10671069
def test_to_sql_exist(conn, mode, num_row_coef, test_frame1, request):
1070+
connections_without_truncate = sqlite_connectable + [
1071+
"sqlite_buildin",
1072+
"sqlite_adbc_conn",
1073+
]
1074+
if conn in connections_without_truncate and mode == "truncate":
1075+
context = pytest.raises(
1076+
NotImplementedError,
1077+
match="'TRUNCATE TABLE' is not supported by this database.",
1078+
)
1079+
else:
1080+
context = contextlib.nullcontext()
10681081
conn = request.getfixturevalue(conn)
1082+
10691083
with pandasSQL_builder(conn, need_transaction=True) as pandasSQL:
10701084
pandasSQL.to_sql(test_frame1, "test_frame", if_exists="fail")
1071-
pandasSQL.to_sql(test_frame1, "test_frame", if_exists=mode)
1085+
with context:
1086+
pandasSQL.to_sql(test_frame1, "test_frame", if_exists=mode)
10721087
assert pandasSQL.has_table("test_frame")
10731088
assert count_rows(conn, "test_frame") == num_row_coef * len(test_frame1)
10741089

@@ -2693,6 +2708,47 @@ def test_drop_table(conn, request):
26932708
assert not insp.has_table("temp_frame")
26942709

26952710

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

0 commit comments

Comments
 (0)