Skip to content

Commit 804fb3e

Browse files
committed
refactor: deprecate usage of cursor.execute statements in favor of the in class implementation of execute.
1 parent 4c3b968 commit 804fb3e

File tree

1 file changed

+41
-20
lines changed

1 file changed

+41
-20
lines changed

pandas/io/sql.py

+41-20
Original file line numberDiff line numberDiff line change
@@ -1651,10 +1651,18 @@ def run_transaction(self):
16511651

16521652
def execute(self, sql: str | Select | TextClause, params=None):
16531653
"""Simple passthrough to SQLAlchemy connectable"""
1654+
from sqlalchemy.exc import DBAPIError as SQLAlchemyDatabaseError
1655+
16541656
args = [] if params is None else [params]
16551657
if isinstance(sql, str):
1656-
return self.con.exec_driver_sql(sql, *args)
1657-
return self.con.execute(sql, *args)
1658+
execute_function = self.con.exec_driver_sql
1659+
else:
1660+
execute_function = self.con.execute
1661+
1662+
try:
1663+
return execute_function(sql, *args)
1664+
except SQLAlchemyDatabaseError as exc:
1665+
raise DatabaseError(f"Execution failed on sql '{sql}': {exc}") from exc
16581666

16591667
def read_table(
16601668
self,
@@ -2108,17 +2116,19 @@ def run_transaction(self):
21082116
self.con.commit()
21092117

21102118
def execute(self, sql: str | Select | TextClause, params=None):
2119+
from adbc_driver_manager import DatabaseError as ADBCDatabaseError
2120+
21112121
if not isinstance(sql, str):
21122122
raise TypeError("Query must be a string unless using sqlalchemy.")
21132123
args = [] if params is None else [params]
21142124
cur = self.con.cursor()
21152125
try:
21162126
cur.execute(sql, *args)
21172127
return cur
2118-
except Exception as exc:
2128+
except ADBCDatabaseError as exc:
21192129
try:
21202130
self.con.rollback()
2121-
except Exception as inner_exc: # pragma: no cover
2131+
except ADBCDatabaseError as inner_exc: # pragma: no cover
21222132
ex = DatabaseError(
21232133
f"Execution failed on sql: {sql}\n{exc}\nunable to rollback"
21242134
)
@@ -2207,8 +2217,7 @@ def read_table(
22072217
else:
22082218
stmt = f"SELECT {select_list} FROM {table_name}"
22092219

2210-
with self.con.cursor() as cur:
2211-
cur.execute(stmt)
2220+
with self.execute(stmt) as cur:
22122221
pa_table = cur.fetch_arrow_table()
22132222
df = arrow_table_to_pandas(pa_table, dtype_backend=dtype_backend)
22142223

@@ -2278,8 +2287,7 @@ def read_query(
22782287
if chunksize:
22792288
raise NotImplementedError("'chunksize' is not implemented for ADBC drivers")
22802289

2281-
with self.con.cursor() as cur:
2282-
cur.execute(sql)
2290+
with self.execute(sql) as cur:
22832291
pa_table = cur.fetch_arrow_table()
22842292
df = arrow_table_to_pandas(pa_table, dtype_backend=dtype_backend)
22852293

@@ -2335,6 +2343,9 @@ def to_sql(
23352343
engine : {'auto', 'sqlalchemy'}, default 'auto'
23362344
Raises NotImplementedError if not set to 'auto'
23372345
"""
2346+
from adbc_driver_manager import DatabaseError as ADBCDatabaseError
2347+
import pyarrow as pa
2348+
23382349
if index_label:
23392350
raise NotImplementedError(
23402351
"'index_label' is not implemented for ADBC drivers"
@@ -2364,22 +2375,25 @@ def to_sql(
23642375
if if_exists == "fail":
23652376
raise ValueError(f"Table '{table_name}' already exists.")
23662377
elif if_exists == "replace":
2367-
with self.con.cursor() as cur:
2368-
cur.execute(f"DROP TABLE {table_name}")
2378+
sql_statement = f"DROP TABLE {table_name}"
2379+
self.execute(sql_statement).close()
23692380
elif if_exists == "append":
23702381
mode = "append"
23712382

2372-
import pyarrow as pa
2373-
23742383
try:
23752384
tbl = pa.Table.from_pandas(frame, preserve_index=index)
23762385
except pa.ArrowNotImplementedError as exc:
23772386
raise ValueError("datatypes not supported") from exc
23782387

23792388
with self.con.cursor() as cur:
2380-
total_inserted = cur.adbc_ingest(
2381-
table_name=name, data=tbl, mode=mode, db_schema_name=schema
2382-
)
2389+
try:
2390+
total_inserted = cur.adbc_ingest(
2391+
table_name=name, data=tbl, mode=mode, db_schema_name=schema
2392+
)
2393+
except ADBCDatabaseError as exc:
2394+
raise DatabaseError(
2395+
f"Failed to insert records on table={name} with {mode=}"
2396+
) from exc
23832397

23842398
self.con.commit()
23852399
return total_inserted
@@ -2496,9 +2510,9 @@ def sql_schema(self) -> str:
24962510
return str(";\n".join(self.table))
24972511

24982512
def _execute_create(self) -> None:
2499-
with self.pd_sql.run_transaction() as conn:
2513+
with self.pd_sql.run_transaction():
25002514
for stmt in self.table:
2501-
conn.execute(stmt)
2515+
self.pd_sql.execute(stmt).close()
25022516

25032517
def insert_statement(self, *, num_rows: int) -> str:
25042518
names = list(map(str, self.frame.columns))
@@ -2520,8 +2534,13 @@ def insert_statement(self, *, num_rows: int) -> str:
25202534
return insert_statement
25212535

25222536
def _execute_insert(self, conn, keys, data_iter) -> int:
2537+
from sqlite3 import DatabaseError as SQLiteDatabaseError
2538+
25232539
data_list = list(data_iter)
2524-
conn.executemany(self.insert_statement(num_rows=1), data_list)
2540+
try:
2541+
conn.executemany(self.insert_statement(num_rows=1), data_list)
2542+
except SQLiteDatabaseError as exc:
2543+
raise DatabaseError("Execution failed") from exc
25252544
return conn.rowcount
25262545

25272546
def _execute_insert_multi(self, conn, keys, data_iter) -> int:
@@ -2643,17 +2662,19 @@ def run_transaction(self):
26432662
cur.close()
26442663

26452664
def execute(self, sql: str | Select | TextClause, params=None):
2665+
from sqlite3 import DatabaseError as SQLiteDatabaseError
2666+
26462667
if not isinstance(sql, str):
26472668
raise TypeError("Query must be a string unless using sqlalchemy.")
26482669
args = [] if params is None else [params]
26492670
cur = self.con.cursor()
26502671
try:
26512672
cur.execute(sql, *args)
26522673
return cur
2653-
except Exception as exc:
2674+
except SQLiteDatabaseError as exc:
26542675
try:
26552676
self.con.rollback()
2656-
except Exception as inner_exc: # pragma: no cover
2677+
except SQLiteDatabaseError as inner_exc: # pragma: no cover
26572678
ex = DatabaseError(
26582679
f"Execution failed on sql: {sql}\n{exc}\nunable to rollback"
26592680
)

0 commit comments

Comments
 (0)