fix: restore parameters support when sql passed to SnowflakeHook as str

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.

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 = {