Skip to content

Commit

Permalink
Fix UnboundLocalError when sql is empty list in SnowflakeHook
Browse files Browse the repository at this point in the history
  • Loading branch information
kazanzhy authored and potiuk committed May 19, 2022
1 parent 54aa234 commit e5c165c
Show file tree
Hide file tree
Showing 2 changed files with 17 additions and 5 deletions.
14 changes: 9 additions & 5 deletions airflow/providers/snowflake/hooks/snowflake.py
Original file line number Diff line number Diff line change
Expand Up @@ -308,14 +308,18 @@ def run(
"""
self.query_ids = []

if isinstance(sql, str):
split_statements_tuple = split_statements(StringIO(sql))
sql = [sql_string for sql_string, _ in split_statements_tuple if sql_string]

if sql:
self.log.debug("Executing %d statements against Snowflake DB", len(sql))
else:
raise ValueError("List of SQL statements is empty")

with closing(self.get_conn()) as conn:
self.set_autocommit(conn, autocommit)

if isinstance(sql, str):
split_statements_tuple = split_statements(StringIO(sql))
sql = [sql_string for sql_string, _ in split_statements_tuple if sql_string]

self.log.debug("Executing %d statements against Snowflake DB", len(sql))
# SnowflakeCursor does not extend ContextManager, so we have to ignore mypy error here
with closing(conn.cursor(DictCursor)) as cur: # type: ignore[type-var]

Expand Down
8 changes: 8 additions & 0 deletions tests/providers/snowflake/hooks/test_snowflake.py
Original file line number Diff line number Diff line change
Expand Up @@ -515,3 +515,11 @@ def test_connection_failure(self, mock_run):
assert status is False
assert msg == 'Connection Errors'
mock_run.assert_called_once_with(sql='select 1')

def test_empty_sql_parameter(self):
hook = SnowflakeHook()

for empty_statement in ([], '', '\n'):
with pytest.raises(ValueError) as err:
hook.run(sql=empty_statement)
assert err.value.args[0] == "List of SQL statements is empty"

0 comments on commit e5c165c

Please sign in to comment.