Repository: apache/airflow

Issue:
Snowflake hook doesn't parameterize SQL passed as a string type, causing SnowflakeOperator to fail
<!--

Welcome to Apache Airflow!  For a smooth issue process, try to answer the following questions.
Don't worry if they're not all applicable; just try to include what you can :-)

If you need to include code snippets or logs, please put them in fenced code
blocks.  If they're super-long, please use the details tag like
<details><summary>super-long log</summary> lots of stuff </details>

Please delete these comment blocks before submitting the issue.

-->

<!--

IMPORTANT!!!

PLEASE CHECK "SIMILAR TO X EXISTING ISSUES" OPTION IF VISIBLE
NEXT TO "SUBMIT NEW ISSUE" BUTTON!!!

PLEASE CHECK IF THIS ISSUE HAS BEEN REPORTED PREVIOUSLY USING SEARCH!!!

Please complete the next sections or the issue will be closed.
These questions are the first thing we need to know to understand the context.

-->

**Apache Airflow version**: 2.1.0

**What happened**: The new Snowflake hook run method is not taking parameters into account when the SQL is passed as a string; it's using Snowflake connector's execute_string method, which does not support parameterization. So the only way to parameterize your query from a SnowflakeOperator is to put the SQL into a list. 

https://github.com/apache/airflow/blob/304e174674ff6921cb7ed79c0158949b50eff8fe/airflow/providers/snowflake/hooks/snowflake.py#L272-L279

https://docs.snowflake.com/en/user-guide/python-connector-api.html#execute_string

**How to reproduce it**: Pass a sql string and parameters to SnowflakeOperator; the query will not be parameterized, and will fail as a SQL syntax error on the parameterization characters, e.g. %(param)s.
<!---

As minimally and precisely as possible. Keep in mind we do not have access to your cluster or dags.

If you are using kubernetes, please attempt to recreate the issue using minikube or kind.

## Install minikube/kind

- Minikube https://minikube.sigs.k8s.io/docs/start/
- Kind https://kind.sigs.k8s.io/docs/user/quick-start/

If this is a UI bug, please provide a screenshot of the bug or a link to a youtube video of the bug in action

You can include images using the .md style of
![alt text](http://url/to/img.png)

To record a screencast, mac users can use QuickTime and then create an unlisted youtube video with the resulting .mov file.

--->


**Anything else we need to know**: Quick workaround is to put your sql string into a list, these are still being parameterized correctly. 

<!--

How often does this problem occur? Once? Every time etc?

Any relevant logs to include? Put them here in side a detail tag:
<details><summary>x.log</summary> lots of stuff </details>

-->



Discussion:
Thanks for opening your first issue here! Be sure to follow the issue template!

Would you be willing to submit a fix @grassten ? Happy to review and merge it,
> Would you be willing to submit a fix @grassten ? Happy to review and merge it,

Sure! I'll work on getting a PR in today or tomorrow morning.

Resolving patch:
diff --git a/airflow/providers/snowflake/hooks/snowflake.py b/airflow/providers/snowflake/hooks/snowflake.py
--- a/airflow/providers/snowflake/hooks/snowflake.py
+++ b/airflow/providers/snowflake/hooks/snowflake.py
@@ -16,6 +16,7 @@
 # specific language governing permissions and limitations
 # under the License.
 from contextlib import closing
+from io import StringIO
 from typing import Any, Dict, Optional, Tuple, Union
 
 from cryptography.hazmat.backends import default_backend
@@ -24,6 +25,7 @@
 # pylint: disable=no-name-in-module
 from snowflake import connector
 from snowflake.connector import SnowflakeConnection
+from snowflake.connector.util_text import split_statements
 
 from airflow.hooks.dbapi import DbApiHook
 
@@ -270,27 +272,21 @@ def run(self, sql: Union[str, list], autocommit: bool = False, parameters: Optio
             self.set_autocommit(conn, autocommit)
 
             if isinstance(sql, str):
-                cursors = conn.execute_string(sql, return_cursors=True)
-                for cur in cursors:
-                    self.query_ids.append(cur.sfqid)
-
+                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))
+            with closing(conn.cursor()) as cur:
+                for sql_statement in sql:
+
+                    self.log.info("Running statement: %s, parameters: %s", sql_statement, parameters)
+                    if parameters:
+                        cur.execute(sql_statement, parameters)
+                    else:
+                        cur.execute(sql_statement)
                     self.log.info("Rows affected: %s", cur.rowcount)
                     self.log.info("Snowflake query id: %s", cur.sfqid)
-                    cur.close()
-
-            elif isinstance(sql, list):
-                self.log.debug("Executing %d statements against Snowflake DB", len(sql))
-                with closing(conn.cursor()) as cur:
-                    for sql_statement in sql:
-
-                        self.log.info("Running statement: %s, parameters: %s", sql_statement, parameters)
-                        if parameters:
-                            cur.execute(sql_statement, parameters)
-                        else:
-                            cur.execute(sql_statement)
-                        self.log.info("Rows affected: %s", cur.rowcount)
-                        self.log.info("Snowflake query id: %s", cur.sfqid)
-                        self.query_ids.append(cur.sfqid)
+                    self.query_ids.append(cur.sfqid)
 
             # If autocommit was set to False for db that supports autocommit,
             # or if db does not supports autocommit, we do a manual commit.


Test patch:
diff --git a/tests/providers/snowflake/hooks/test_snowflake.py b/tests/providers/snowflake/hooks/test_snowflake.py
--- a/tests/providers/snowflake/hooks/test_snowflake.py
+++ b/tests/providers/snowflake/hooks/test_snowflake.py
@@ -17,12 +17,14 @@
 # under the License.
 #
 import os
+import re
 import unittest
 from unittest import mock
 
 from cryptography.hazmat.backends import default_backend
 from cryptography.hazmat.primitives import serialization
 from cryptography.hazmat.primitives.asymmetric import rsa
+from parameterized import parameterized
 
 from airflow.providers.snowflake.hooks.snowflake import SnowflakeHook
 
@@ -31,15 +33,7 @@ class TestSnowflakeHook(unittest.TestCase):
     def setUp(self):
         super().setUp()
 
-        self.cur = mock.MagicMock(rowcount=0)
-        self.cur2 = mock.MagicMock(rowcount=0)
-
-        self.cur.sfqid = 'uuid'
-        self.cur2.sfqid = 'uuid2'
-
         self.conn = conn = mock.MagicMock()
-        self.conn.cursor.return_value = self.cur
-        self.conn.execute_string.return_value = [self.cur, self.cur2]
 
         self.conn.login = 'user'
         self.conn.password = 'pw'
@@ -95,18 +89,25 @@ def test_get_uri(self):
         )
         assert uri_shouldbe == self.db_hook.get_uri()
 
-    def test_single_element_list_calls_execute(self):
-        self.db_hook.run(['select * from table'])
-        self.cur.execute.assert_called()
-        assert self.db_hook.query_ids == ['uuid']
-
-    def test_passed_string_calls_execute_string(self):
-        self.db_hook.run('select * from table; select * from table2')
-
-        assert self.db_hook.query_ids == ['uuid', 'uuid2']
-        self.conn.execute_string.assert_called()
-        self.cur.close.assert_called()
-        self.cur2.close.assert_called()
+    @parameterized.expand(
+        [
+            ('select * from table', ['uuid', 'uuid']),
+            ('select * from table;select * from table2', ['uuid', 'uuid', 'uuid2', 'uuid2']),
+            (['select * from table;'], ['uuid', 'uuid']),
+            (['select * from table;', 'select * from table2;'], ['uuid', 'uuid', 'uuid2', 'uuid2']),
+        ],
+    )
+    def test_run_storing_query_ids(self, sql, query_ids):
+        cur = mock.MagicMock(rowcount=0)
+        self.conn.cursor.return_value = cur
+        type(cur).sfqid = mock.PropertyMock(side_effect=query_ids)
+        mock_params = {"mock_param": "mock_param"}
+        self.db_hook.run(sql, parameters=mock_params)
+
+        sql_list = sql if isinstance(sql, list) else re.findall(".*?[;]", sql)
+        cur.execute.assert_has_calls([mock.call(query, mock_params) for query in sql_list])
+        assert self.db_hook.query_ids == query_ids[::2]
+        cur.close.assert_called()
 
     def test_get_conn_params(self):
         conn_params_shouldbe = {