Skip to content

Commit

Permalink
fix: log DDL SQL in verbose mode
Browse files Browse the repository at this point in the history
Before (and still), we log the SQL inside Backend.compile().
I *think* what this means is that we are logging all SELECT statements.
But, any DDL statements like from create_table() were not logged.
Now they are.
Some SQL statements may be logged twice now, in .compile() and in .raw_sql(), but I don't think that's a big problem?
  • Loading branch information
NickCrews committed Aug 11, 2024
1 parent 9165255 commit 22e79bd
Show file tree
Hide file tree
Showing 12 changed files with 27 additions and 22 deletions.
4 changes: 1 addition & 3 deletions ibis/backends/bigquery/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -666,6 +666,7 @@ def raw_sql(self, query: str, params=None, page_size: int | None = None):
]
with contextlib.suppress(AttributeError):
query = query.sql(self.dialect)
self._log(query)

Check warning on line 669 in ibis/backends/bigquery/__init__.py

View check run for this annotation

Codecov / codecov/patch

ibis/backends/bigquery/__init__.py#L669

Added line #L669 was not covered by tests

job_config = bq.job.QueryJobConfig(query_parameters=query_parameters or [])
return self.client.query_and_wait(
Expand Down Expand Up @@ -737,7 +738,6 @@ def execute(self, expr, params=None, limit="default", **kwargs):
schema = expr.as_table().schema() - ibis.schema({"_TABLE_SUFFIX": "string"})

sql = self.compile(expr, limit=limit, params=params, **kwargs)
self._log(sql)
query = self.raw_sql(sql, params=params, **kwargs)

arrow_t = query.to_arrow(
Expand Down Expand Up @@ -799,7 +799,6 @@ def to_pyarrow(
self._import_pyarrow()
self._register_in_memory_tables(expr)
sql = self.compile(expr, limit=limit, params=params, **kwargs)
self._log(sql)
query = self.raw_sql(sql, params=params, **kwargs)
table = query.to_arrow(
progress_bar_type=None, bqstorage_client=self.storage_client
Expand All @@ -822,7 +821,6 @@ def to_pyarrow_batches(

self._register_in_memory_tables(expr)
sql = self.compile(expr, limit=limit, params=params, **kwargs)
self._log(sql)
query = self.raw_sql(sql, params=params, page_size=chunk_size, **kwargs)
batch_iter = query.to_arrow_iterable(bqstorage_client=self.storage_client)
return pa.ipc.RecordBatchReader.from_batches(schema.to_pyarrow(), batch_iter)
Expand Down
30 changes: 15 additions & 15 deletions ibis/backends/duckdb/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -91,6 +91,7 @@ def current_database(self) -> str:
def raw_sql(self, query: str | sg.Expression, **kwargs: Any) -> Any:
with contextlib.suppress(AttributeError):
query = query.sql(dialect=self.name)
self._log(query)
return self.con.execute(query, **kwargs)

def create_table(
Expand Down Expand Up @@ -203,43 +204,42 @@ def create_table(
# This is the same table as initial_table unless overwrite == True
final_table = sge.Table(
this=sg.to_identifier(name, quoted=self.compiler.quoted),
catalog=catalog,
db=database,
catalog=sg.to_identifier(catalog, quoted=self.compiler.quoted),
db=sg.to_identifier(database, quoted=self.compiler.quoted),
)
with self._safe_raw_sql(create_stmt) as cur:

def cur_exec(stmt):
sql = stmt.sql(self.name)
self._log(sql)
return cur.execute(sql)

if query is not None:
insert_stmt = sge.insert(query, into=initial_table).sql(self.name)
cur.execute(insert_stmt).fetchall()
cur_exec(sge.insert(query, into=initial_table)).fetchall()

if overwrite:
cur.execute(
sge.Drop(kind="TABLE", this=final_table, exists=True).sql(self.name)
)
cur_exec(sge.Drop(kind="TABLE", this=final_table, exists=True))
# TODO: This branching should be removed once DuckDB >=0.9.3 is
# our lower bound (there's an upstream bug in 0.9.2 that
# disallows renaming temp tables)
# We should (pending that release) be able to remove the if temp
# branch entirely.
if temp:
cur.execute(
cur_exec(
sge.Create(
kind="TABLE",
this=final_table,
expression=sg.select(STAR).from_(initial_table),
properties=sge.Properties(expressions=properties),
).sql(self.name)
)
cur.execute(
sge.Drop(kind="TABLE", this=initial_table, exists=True).sql(
self.name
)
)
cur_exec(sge.Drop(kind="TABLE", this=initial_table, exists=True))
else:
cur.execute(
cur_exec(
sge.AlterTable(
this=initial_table,
actions=[sge.RenameTable(this=final_table)],
).sql(self.name)
)
)

if temp_memtable_view is not None:
Expand Down
1 change: 1 addition & 0 deletions ibis/backends/flink/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -88,6 +88,7 @@ def disconnect(self) -> None:
pass

def raw_sql(self, query: str) -> TableResult:
self._log(query)
return self._table_env.execute_sql(query)

def _get_schema_using_query(self, query: str) -> sch.Schema:
Expand Down
4 changes: 2 additions & 2 deletions ibis/backends/impala/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -247,12 +247,12 @@ def raw_sql(self, query: str):
try:
for k, v in self.options.items():
q = f"SET {k} = {v!r}"
util.log(q)
self._log(q)
cursor.execute_async(q)

cursor._wait_to_finish()

util.log(query)
self._log(query)
cursor.execute_async(query)

cursor._wait_to_finish()
Expand Down
1 change: 1 addition & 0 deletions ibis/backends/mssql/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -338,6 +338,7 @@ def _safe_ddl(self, query, *args, **kwargs):
def raw_sql(self, query: str | sg.Expression, **kwargs: Any) -> Any:
with contextlib.suppress(AttributeError):
query = query.sql(self.dialect)
self._log(query)

Check warning on line 341 in ibis/backends/mssql/__init__.py

View check run for this annotation

Codecov / codecov/patch

ibis/backends/mssql/__init__.py#L341

Added line #L341 was not covered by tests

con = self.con
cursor = con.cursor()
Expand Down
1 change: 1 addition & 0 deletions ibis/backends/mysql/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -275,6 +275,7 @@ def _safe_raw_sql(self, *args, **kwargs):
def raw_sql(self, query: str | sg.Expression, **kwargs: Any) -> Any:
with contextlib.suppress(AttributeError):
query = query.sql(dialect=self.name)
self._log(query)

con = self.con
cursor = con.cursor()
Expand Down
1 change: 1 addition & 0 deletions ibis/backends/oracle/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -219,6 +219,7 @@ def _safe_raw_sql(self, *args, **kwargs):
def raw_sql(self, query: str | sg.Expression, **kwargs: Any) -> Any:
with contextlib.suppress(AttributeError):
query = query.sql(dialect=self.name)
self._log(query)

con = self.con
cursor = con.cursor()
Expand Down
2 changes: 1 addition & 1 deletion ibis/backends/postgres/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -747,7 +747,7 @@ def raw_sql(self, query: str | sg.Expression, **kwargs: Any) -> Any:

with contextlib.suppress(AttributeError):
query = query.sql(dialect=self.dialect)

self._log(query)
con = self.con
cursor = con.cursor()

Expand Down
1 change: 1 addition & 0 deletions ibis/backends/pyspark/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -410,6 +410,7 @@ def _safe_raw_sql(self, query: str) -> Any:
def raw_sql(self, query: str | sg.Expression, **kwargs: Any) -> Any:
with contextlib.suppress(AttributeError):
query = query.sql(dialect=self.dialect)
self._log(query)
return self._session.sql(query, **kwargs)

def execute(
Expand Down
1 change: 1 addition & 0 deletions ibis/backends/snowflake/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -728,6 +728,7 @@ def _safe_raw_sql(self, query: str | sg.Expression, **kwargs: Any) -> Any:
def raw_sql(self, query: str | sg.Expression, **kwargs: Any) -> Any:
with contextlib.suppress(AttributeError):
query = query.sql(dialect=self.name)
self._log(query)

Check warning on line 731 in ibis/backends/snowflake/__init__.py

View check run for this annotation

Codecov / codecov/patch

ibis/backends/snowflake/__init__.py#L731

Added line #L731 was not covered by tests
cur = self.con.cursor()
try:
cur.execute(query, **kwargs)
Expand Down
1 change: 1 addition & 0 deletions ibis/backends/sqlite/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -126,6 +126,7 @@ def _post_connect(
def raw_sql(self, query: str | sg.Expression, **kwargs: Any) -> Any:
if not isinstance(query, str):
query = query.sql(dialect=self.name)
self._log(query)
return self.con.execute(query, **kwargs)

@contextlib.contextmanager
Expand Down
2 changes: 1 addition & 1 deletion ibis/backends/trino/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -57,7 +57,7 @@ def raw_sql(self, query: str | sg.Expression) -> Any:
"""Execute a raw SQL query."""
with contextlib.suppress(AttributeError):
query = query.sql(dialect=self.name, pretty=True)

self._log(query)
con = self.con
cur = con.cursor()
try:
Expand Down

0 comments on commit 22e79bd

Please sign in to comment.