Skip to content

Commit

Permalink
- Snowflake does not currently support variable binding in multi-stat…
Browse files Browse the repository at this point in the history
…ement SQL requests. (apache#42719)

- added warning for multi-statement query bindings in Snowflake hook
test: added unit test for multi-statement bindings warning

Related: apache#42033
  • Loading branch information
harjeevanmaan authored and joaopamaral committed Oct 21, 2024
1 parent ded7a12 commit 914dd23
Show file tree
Hide file tree
Showing 2 changed files with 43 additions and 2 deletions.
10 changes: 8 additions & 2 deletions airflow/providers/snowflake/hooks/snowflake_sql_api.py
Original file line number Diff line number Diff line change
Expand Up @@ -153,8 +153,14 @@ def execute_query(
url = f"{self.account_identifier}.snowflakecomputing.com/api/v2/statements"
params: dict[str, Any] | None = {"requestId": str(req_id), "async": True, "pageSize": 10}
headers = self.get_headers()
if bindings is None:
bindings = {}
sql_is_multi_stmt = ";" in sql.strip()
if not isinstance(bindings, dict) and bindings is not None:
raise AirflowException("Bindings should be a dictionary or None.")
if bindings and sql_is_multi_stmt:
self.log.warning(
"Bindings are not supported for multi-statement queries. Bindings will be ignored."
)
bindings = bindings or {}
data = {
"statement": sql,
"resultSetMetaData": {"format": "json"},
Expand Down
35 changes: 35 additions & 0 deletions tests/providers/snowflake/hooks/test_snowflake_sql_api.py
Original file line number Diff line number Diff line change
Expand Up @@ -234,6 +234,41 @@ def test_execute_query_exception_without_statement_handle(
hook.execute_query(sql, statement_count)
assert exception_info

@pytest.mark.parametrize(
"sql,statement_count,bindings",
[
(SQL_MULTIPLE_STMTS, 4, {"1": {"type": "FIXED", "value": "123"}}),
],
)
@mock.patch("airflow.providers.snowflake.hooks.snowflake_sql_api.requests")
@mock.patch(
"airflow.providers.snowflake.hooks.snowflake_sql_api.SnowflakeSqlApiHook._get_conn_params",
new_callable=PropertyMock,
)
@mock.patch("airflow.providers.snowflake.hooks.snowflake_sql_api.SnowflakeSqlApiHook.get_headers")
def test_execute_query_bindings_warning(
self,
mock_get_headers,
mock_conn_params,
mock_requests,
sql,
statement_count,
bindings,
):
"""Test execute_query method logs warning when bindings are provided for multi-statement queries"""
mock_conn_params.return_value = CONN_PARAMS
mock_get_headers.return_value = HEADERS
mock_requests.post.return_value = create_successful_response_mock(
{"statementHandles": ["uuid", "uuid1"]}
)

hook = SnowflakeSqlApiHook(snowflake_conn_id="mock_conn_id")
with mock.patch.object(hook.log, "warning") as mock_log_warning:
hook.execute_query(sql, statement_count, bindings=bindings)
mock_log_warning.assert_called_once_with(
"Bindings are not supported for multi-statement queries. Bindings will be ignored."
)

@pytest.mark.parametrize(
"query_ids",
[
Expand Down

0 comments on commit 914dd23

Please sign in to comment.