@@ -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 } \n unable 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 } \n unable to rollback"
26592680 )
0 commit comments