From 2b94edc2d3e0666cf53b3b2421f561f8b4cbb26a Mon Sep 17 00:00:00 2001 From: Brandon Chinn Date: Thu, 11 Sep 2025 12:48:51 -0700 Subject: [PATCH 1/2] Add test for get_results_from_sfqid with DictCursor --- test/integ/test_async.py | 6 ++++-- test/integ/test_multi_statement.py | 25 +++++++++++++++++++++---- 2 files changed, 25 insertions(+), 6 deletions(-) diff --git a/test/integ/test_async.py b/test/integ/test_async.py index 41047b5f35..eec0861f13 100644 --- a/test/integ/test_async.py +++ b/test/integ/test_async.py @@ -7,6 +7,7 @@ import pytest from snowflake.connector import DatabaseError, ProgrammingError +from snowflake.connector.cursor import DictCursor, SnowflakeCursor # Mark all tests in this file to time out after 2 minutes to prevent hanging forever pytestmark = [pytest.mark.timeout(120), pytest.mark.skipolddriver] @@ -17,14 +18,15 @@ QueryStatus = None -def test_simple_async(conn_cnx): +@pytest.mark.parametrize("cursor_class", [SnowflakeCursor, DictCursor]) +def test_simple_async(conn_cnx, cursor_class): """Simple test to that shows the most simple usage of fire and forget. This test also makes sure that wait_until_ready function's sleeping is tested and that some fields are copied over correctly from the original query. """ with conn_cnx() as con: - with con.cursor() as cur: + with con.cursor(cursor_class) as cur: cur.execute_async("select count(*) from table(generator(timeLimit => 5))") cur.get_results_from_sfqid(cur.sfqid) assert len(cur.fetchall()) == 1 diff --git a/test/integ/test_multi_statement.py b/test/integ/test_multi_statement.py index 3fd80485d1..1dff738f20 100644 --- a/test/integ/test_multi_statement.py +++ b/test/integ/test_multi_statement.py @@ -14,6 +14,7 @@ import snowflake.connector.cursor from snowflake.connector import ProgrammingError, errors +from snowflake.connector.cursor import DictCursor, SnowflakeCursor try: # pragma: no cover from snowflake.connector.constants import ( @@ -153,10 +154,11 @@ def test_binding_multi(conn_cnx, style: str, skip_to_last_set: bool): ) -def test_async_exec_multi(conn_cnx, skip_to_last_set: bool): +@pytest.mark.parametrize("cursor_class", [SnowflakeCursor, DictCursor]) +def test_async_exec_multi(conn_cnx, cursor_class, skip_to_last_set: bool): """Tests whether async execution query works within a multi-statement""" with conn_cnx() as con: - with con.cursor() as cur: + with con.cursor(cursor_class) as cur: cur.execute_async( "select 1; select 2; select count(*) from table(generator(timeLimit => 1)); select 'b';", num_statements=4, @@ -165,14 +167,29 @@ def test_async_exec_multi(conn_cnx, skip_to_last_set: bool): assert con.is_still_running(con.get_query_status(q_id)) _wait_while_query_running(con, q_id, sleep_time=1) with conn_cnx() as con: - with con.cursor() as cur: + with con.cursor(cursor_class) as cur: _wait_until_query_success(con, q_id, num_checks=3, sleep_per_check=1) assert con.get_query_status_throw_if_error(q_id) == QueryStatus.SUCCESS + if cursor_class == SnowflakeCursor: + expected = [ + [(1,)], + [(2,)], + lambda x: len(x) == 1 and len(x[0]) == 1 and x[0][0] > 0, + [("b",)], + ] + elif cursor_class == DictCursor: + expected = [ + [{"1": 1}], + [{"2": 2}], + lambda x: len(x) == 1 and len(x[0]) == 1 and x[0]["COUNT(*)"] > 0, + [{"'B'": "b"}], + ] + cur.get_results_from_sfqid(q_id) _check_multi_statement_results( cur, - checks=[[(1,)], [(2,)], lambda x: x > [(0,)], [("b",)]], + checks=expected, skip_to_last_set=skip_to_last_set, ) From 90c833a86edebfe9f8926f0ceaa94f8b6198244b Mon Sep 17 00:00:00 2001 From: Brandon Chinn Date: Thu, 11 Sep 2025 12:49:32 -0700 Subject: [PATCH 2/2] Fix get_results_from_sfqid for DictCursor --- DESCRIPTION.md | 1 + src/snowflake/connector/cursor.py | 3 +-- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/DESCRIPTION.md b/DESCRIPTION.md index af9dc93307..3e02e72660 100644 --- a/DESCRIPTION.md +++ b/DESCRIPTION.md @@ -9,6 +9,7 @@ Source code is also available at: https://github.com/snowflakedb/snowflake-conne # Release Notes - v3.18.0(TBD) - Added the `workload_identity_impersonation_path` parameter to support service account impersonation for Workload Identity Federation on GCP and AWS workloads only + - Fixed `get_results_from_sfqid` when using `DictCursor` and executing multiple statements at once - v3.17.3(September 02,2025) - Enhanced configuration file permission warning messages. diff --git a/src/snowflake/connector/cursor.py b/src/snowflake/connector/cursor.py index a442ad102d..6ade7f3d8e 100644 --- a/src/snowflake/connector/cursor.py +++ b/src/snowflake/connector/cursor.py @@ -1741,8 +1741,7 @@ def wait_until_ready() -> None: self.connection.get_query_status_throw_if_error( sfqid ) # Trigger an exception if query failed - klass = self.__class__ - self._inner_cursor = klass(self.connection) + self._inner_cursor = SnowflakeCursor(self.connection) self._sfqid = sfqid self._prefetch_hook = wait_until_ready