Skip to content

Commit

Permalink
feat: Implement mode 'truncate' in 'to_sql' API
Browse files Browse the repository at this point in the history
  • Loading branch information
gmcrocetti committed Aug 2, 2024
1 parent 2abbcf3 commit d262bb4
Show file tree
Hide file tree
Showing 2 changed files with 44 additions and 3 deletions.
36 changes: 35 additions & 1 deletion pandas/io/sql.py
Original file line number Diff line number Diff line change
Expand Up @@ -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:
Expand Down Expand Up @@ -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,
Expand Down Expand Up @@ -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"
Expand Down Expand Up @@ -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

Expand Down Expand Up @@ -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,
Expand Down
11 changes: 9 additions & 2 deletions pandas/tests/io/test_sql.py
Original file line number Diff line number Diff line change
Expand Up @@ -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)

Expand Down

0 comments on commit d262bb4

Please sign in to comment.