Skip to content

Commit ad21c3b

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

File tree

3 files changed

+68
-16
lines changed

3 files changed

+68
-16
lines changed

doc/source/whatsnew/v3.0.0.rst

+1-1
Original file line numberDiff line numberDiff line change
@@ -53,7 +53,7 @@ Other enhancements
5353
- Restore support for reading Stata 104-format and enable reading 103-format dta files (:issue:`58554`)
5454
- Support reading Stata 102-format (Stata 1) dta files (:issue:`58978`)
5555
- Support reading Stata 110-format (Stata 7) dta files (:issue:`47176`)
56-
56+
- Add ``"truncate"`` option to ``if_exists`` argument in :meth:`DataFrame.to_sql` truncating the table before inserting data (:issue:`37210`).
5757
.. ---------------------------------------------------------------------------
5858
.. _whatsnew_300.notable_bug_fixes:
5959

pandas/io/sql.py

+52-10
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 TABLE'}, 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 TABLE'}, 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,25 @@ 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:
2068+
raise NotImplementedError("'TRUNCATE TABLE' is not supported by this database.")
2069+
2070+
self.meta.clear()
2071+
2072+
20482073
def _create_sql_schema(
20492074
self,
20502075
frame: DataFrame,
@@ -2301,7 +2326,7 @@ def to_sql(
23012326
self,
23022327
frame,
23032328
name: str,
2304-
if_exists: Literal["fail", "replace", "append"] = "fail",
2329+
if_exists: Literal["fail", "replace", "append", "truncate"] = "fail",
23052330
index: bool = True,
23062331
index_label=None,
23072332
schema: str | None = None,
@@ -2323,6 +2348,8 @@ def to_sql(
23232348
- fail: If table exists, do nothing.
23242349
- replace: If table exists, drop it, recreate it, and insert data.
23252350
- append: If table exists, insert data. Create if does not exist.
2351+
- truncate: If table exists, truncate it. Create if does not exist.
2352+
Raises NotImplementedError if 'TRUNCATE TABLE' is not supported
23262353
index : boolean, default True
23272354
Write DataFrame index as a column.
23282355
index_label : string or sequence, default None
@@ -2340,6 +2367,8 @@ def to_sql(
23402367
engine : {'auto', 'sqlalchemy'}, default 'auto'
23412368
Raises NotImplementedError if not set to 'auto'
23422369
"""
2370+
from adbc_driver_manager import ProgrammingError
2371+
23432372
if index_label:
23442373
raise NotImplementedError(
23452374
"'index_label' is not implemented for ADBC drivers"
@@ -2373,6 +2402,14 @@ def to_sql(
23732402
cur.execute(f"DROP TABLE {table_name}")
23742403
elif if_exists == "append":
23752404
mode = "append"
2405+
elif if_exists == "truncate":
2406+
mode = "append"
2407+
with self.con.cursor() as cur:
2408+
try:
2409+
cur.execute(f"TRUNCATE TABLE {table_name}")
2410+
except ProgrammingError:
2411+
raise NotImplementedError("'TRUNCATE TABLE' is not supported by this database.")
2412+
23762413

23772414
import pyarrow as pa
23782415

@@ -2778,6 +2815,8 @@ def to_sql(
27782815
fail: If table exists, do nothing.
27792816
replace: If table exists, drop it, recreate it, and insert data.
27802817
append: If table exists, insert data. Create if it does not exist.
2818+
truncate: If table exists, truncate it. Create if does not exist.
2819+
Raises NotImplementedError if 'TRUNCATE TABLE' is not supported
27812820
index : bool, default True
27822821
Write DataFrame index as a column
27832822
index_label : string or sequence, default None
@@ -2853,6 +2892,9 @@ def drop_table(self, name: str, schema: str | None = None) -> None:
28532892
drop_sql = f"DROP TABLE {_get_valid_sqlite_name(name)}"
28542893
self.execute(drop_sql)
28552894

2895+
def truncate_table(self, name:str, schema: str | None) -> None:
2896+
raise NotImplementedError("'TRUNCATE TABLE' is not supported by this database.")
2897+
28562898
def _create_sql_schema(
28572899
self,
28582900
frame,

pandas/tests/io/test_sql.py

+15-5
Original file line numberDiff line numberDiff line change
@@ -606,7 +606,7 @@ def mysql_pymysql_engine():
606606
sqlalchemy = pytest.importorskip("sqlalchemy")
607607
pymysql = pytest.importorskip("pymysql")
608608
engine = sqlalchemy.create_engine(
609-
"mysql+pymysql://root@localhost:3306/pandas",
609+
"mysql+pymysql://root@mysql:3306/pandas",
610610
connect_args={"client_flag": pymysql.constants.CLIENT.MULTI_STATEMENTS},
611611
poolclass=sqlalchemy.pool.NullPool,
612612
)
@@ -654,7 +654,7 @@ def postgresql_psycopg2_engine():
654654
sqlalchemy = pytest.importorskip("sqlalchemy")
655655
pytest.importorskip("psycopg2")
656656
engine = sqlalchemy.create_engine(
657-
"postgresql+psycopg2://postgres:postgres@localhost:5432/pandas",
657+
"postgresql+psycopg2://postgres:postgres@postgres:5432/pandas",
658658
poolclass=sqlalchemy.pool.NullPool,
659659
)
660660
yield engine
@@ -689,7 +689,7 @@ def postgresql_adbc_conn():
689689
pytest.importorskip("adbc_driver_postgresql")
690690
from adbc_driver_postgresql import dbapi
691691

692-
uri = "postgresql://postgres:postgres@localhost:5432/pandas"
692+
uri = "postgresql://postgres:postgres@postgres:5432/pandas"
693693
with dbapi.connect(uri) as conn:
694694
yield conn
695695
for view in get_all_views(conn):
@@ -1067,12 +1067,22 @@ 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("mode, num_row_coef", [("replace", 1), ("append", 2), ("truncate", 1)])
10711071
def test_to_sql_exist(conn, mode, num_row_coef, test_frame1, request):
1072+
connections_without_truncate = sqlite_connectable + ["sqlite_buildin", "sqlite_adbc_conn"]
1073+
if conn in connections_without_truncate and mode == "truncate":
1074+
context = pytest.raises(
1075+
NotImplementedError,
1076+
match="'TRUNCATE TABLE' is not supported by this database."
1077+
)
1078+
else:
1079+
context = contextlib.nullcontext()
10721080
conn = request.getfixturevalue(conn)
1081+
10731082
with pandasSQL_builder(conn, need_transaction=True) as pandasSQL:
10741083
pandasSQL.to_sql(test_frame1, "test_frame", if_exists="fail")
1075-
pandasSQL.to_sql(test_frame1, "test_frame", if_exists=mode)
1084+
with context:
1085+
pandasSQL.to_sql(test_frame1, "test_frame", if_exists=mode)
10761086
assert pandasSQL.has_table("test_frame")
10771087
assert count_rows(conn, "test_frame") == num_row_coef * len(test_frame1)
10781088

0 commit comments

Comments
 (0)