Skip to content

Commit

Permalink
fix(snowflake): fix quoting across all apis
Browse files Browse the repository at this point in the history
  • Loading branch information
cpcloud authored and gforsyth committed Dec 8, 2023
1 parent 9b65b48 commit 7bf8e84
Show file tree
Hide file tree
Showing 2 changed files with 31 additions and 23 deletions.
8 changes: 7 additions & 1 deletion ibis/backends/base/sql/alchemy/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -991,6 +991,10 @@ def _get_table_identifier(self, *, name, namespace):
db=schema,
catalog=db,
quoted=self.compiler.translator_class._quote_table_names,
).transform(
lambda node: node.__class__(this=node.this, quoted=True)
if isinstance(node, sg.exp.Identifier)
else node
)
return table

Expand All @@ -1015,7 +1019,9 @@ def _get_sqla_table(
def drop_table(
self, name: str, database: str | None = None, force: bool = False
) -> None:
table = sg.table(name, db=database)
table = sg.table(
name, db=database, quoted=self.compiler.translator_class._quote_table_names
)
drop_table = sg.exp.Drop(kind="TABLE", exists=force, this=table)
drop_table_sql = drop_table.sql(dialect=self.name)
with self.begin() as con:
Expand Down
46 changes: 24 additions & 22 deletions ibis/backends/snowflake/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -551,47 +551,48 @@ def _get_temp_view_definition(
yield f"CREATE OR REPLACE TEMPORARY VIEW {name} AS {definition}"

def create_database(self, name: str, force: bool = False) -> None:
current_database = self.current_database
current_schema = self.current_schema
name = self._quote(name)
if_not_exists = "IF NOT EXISTS " * force
create_stmt = sg.exp.Create(
kind="DATABASE", this=sg.to_identifier(name, quoted=True), exists=force
).sql(self.name)
current_ident = sg.table(
self.current_schema, db=self.current_database, quoted=True
).sql(self.name)
with self.begin() as con:
con.exec_driver_sql(f"CREATE DATABASE {if_not_exists}{name}")
con.exec_driver_sql(create_stmt)
# Snowflake automatically switches to the new database after creating
# it per
# https://docs.snowflake.com/en/sql-reference/sql/create-database#general-usage-notes
# so we switch back to the original database and schema
con.exec_driver_sql(
f"USE SCHEMA {self._quote(current_database)}.{self._quote(current_schema)}"
)
con.exec_driver_sql(f"USE SCHEMA {current_ident}")

def drop_database(self, name: str, force: bool = False) -> None:
current_database = self.current_database
if name == current_database:
raise com.UnsupportedOperationError(
"Dropping the current database is not supported because its behavior is undefined"
)
name = self._quote(name)
if_exists = "IF EXISTS " * force
drop_stmt = sg.exp.Drop(
kind="DATABASE", this=sg.to_identifier(name, quoted=True), exists=force
).sql(self.name)
with self.begin() as con:
con.exec_driver_sql(f"DROP DATABASE {if_exists}{name}")
con.exec_driver_sql(drop_stmt)

def create_schema(
self, name: str, database: str | None = None, force: bool = False
) -> None:
name = ".".join(map(self._quote, filter(None, [database, name])))
if_not_exists = "IF NOT EXISTS " * force
current_database = self.current_database
current_schema = self.current_schema
create_stmt = sg.exp.Create(
kind="SCHEMA", this=sg.table(name, db=database, quoted=True), exists=force
).sql(self.name)
current_ident = sg.table(
self.current_schema, db=self.current_database, quoted=True
).sql(self.name)
with self.begin() as con:
con.exec_driver_sql(f"CREATE SCHEMA {if_not_exists}{name}")
con.exec_driver_sql(create_stmt)
# Snowflake automatically switches to the new schema after creating
# it per
# https://docs.snowflake.com/en/sql-reference/sql/create-schema#usage-notes
# so we switch back to the original schema
con.exec_driver_sql(
f"USE SCHEMA {self._quote(current_database)}.{self._quote(current_schema)}"
)
con.exec_driver_sql(f"USE SCHEMA {current_ident}")

def drop_schema(
self, name: str, database: str | None = None, force: bool = False
Expand All @@ -603,10 +604,11 @@ def drop_schema(
"Dropping the current schema is not supported because its behavior is undefined"
)

name = ".".join(map(self._quote, filter(None, [database, name])))
if_exists = "IF EXISTS " * force
drop_stmt = sg.exp.Drop(
kind="SCHEMA", this=sg.table(name, db=database, quoted=True), exists=force
).sql(self.name)
with self.begin() as con:
con.exec_driver_sql(f"DROP SCHEMA {if_exists}{name}")
con.exec_driver_sql(drop_stmt)

def create_table(
self,
Expand Down

0 comments on commit 7bf8e84

Please sign in to comment.