Skip to content

Commit

Permalink
chore(mysql): try using drop temporary
Browse files Browse the repository at this point in the history
  • Loading branch information
cpcloud committed Sep 10, 2024
1 parent c37d5c3 commit 2c5f4dc
Showing 1 changed file with 34 additions and 6 deletions.
40 changes: 34 additions & 6 deletions ibis/backends/mysql/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -5,6 +5,7 @@
import contextlib
import re
import warnings
import weakref
from functools import cached_property
from operator import itemgetter
from typing import TYPE_CHECKING, Any
Expand Down Expand Up @@ -255,18 +256,19 @@ def drop_database(self, name: str, force: bool = False) -> None:
def begin(self):
con = self.con
cur = con.cursor()
autocommit = con.autocommit_mode

if not con.autocommit_mode:
if not autocommit:
con.begin()

Check warning on line 262 in ibis/backends/mysql/__init__.py

View check run for this annotation

Codecov / codecov/patch

ibis/backends/mysql/__init__.py#L262

Added line #L262 was not covered by tests

try:
yield cur
except Exception:
if not con.autocommit_mode:
if not autocommit:
con.rollback()

Check warning on line 268 in ibis/backends/mysql/__init__.py

View check run for this annotation

Codecov / codecov/patch

ibis/backends/mysql/__init__.py#L268

Added line #L268 was not covered by tests
raise
else:
if not con.autocommit_mode:
if not autocommit:
con.commit()

Check warning on line 272 in ibis/backends/mysql/__init__.py

View check run for this annotation

Codecov / codecov/patch

ibis/backends/mysql/__init__.py#L272

Added line #L272 was not covered by tests
finally:
cur.close()
Expand All @@ -283,20 +285,22 @@ def raw_sql(self, query: str | sg.Expression, **kwargs: Any) -> Any:
query = query.sql(dialect=self.name)

con = self.con
autocommit = con.autocommit_mode

cursor = con.cursor()

if not con.autocommit_mode:
if not autocommit:
con.begin()

Check warning on line 293 in ibis/backends/mysql/__init__.py

View check run for this annotation

Codecov / codecov/patch

ibis/backends/mysql/__init__.py#L293

Added line #L293 was not covered by tests

try:
cursor.execute(query, **kwargs)
except Exception:
if not con.autocommit_mode:
if not autocommit:
con.rollback()

Check warning on line 299 in ibis/backends/mysql/__init__.py

View check run for this annotation

Codecov / codecov/patch

ibis/backends/mysql/__init__.py#L299

Added line #L299 was not covered by tests
cursor.close()
raise
else:
if not con.autocommit_mode:
if not autocommit:
con.commit()

Check warning on line 304 in ibis/backends/mysql/__init__.py

View check run for this annotation

Codecov / codecov/patch

ibis/backends/mysql/__init__.py#L304

Added line #L304 was not covered by tests
return cursor

Expand Down Expand Up @@ -504,6 +508,30 @@ def _register_in_memory_table(self, op: ops.InMemoryTable) -> None:
if not df.empty:
cur.executemany(sql, data)

def _register_memtable_finalizer(self, op: ops.InMemoryTable):
weakref.finalize(op, self.drop_table, op.name, force=True, temp=True)

def drop_table(
self,
name: str,
database: tuple[str, str] | str | None = None,
force: bool = False,
temp: bool = False,
) -> None:
table_loc = self._warn_and_create_table_loc(database, None)
catalog, db = self._to_catalog_db_tuple(table_loc)

drop_stmt = sg.exp.Drop(
kind="TABLE",
this=sg.table(name, db=db, catalog=catalog, quoted=self.compiler.quoted),
exists=force,
temporary=temp,
)
drop_stmt_sql = drop_stmt.sql(self.dialect)

with self.con.cursor() as cur:
cur.execute(drop_stmt_sql)

@util.experimental
def to_pyarrow_batches(
self,
Expand Down

0 comments on commit 2c5f4dc

Please sign in to comment.