Skip to content

Commit 29b5b7e

Browse files
Fix get_results_from_sfqid with DictCursor + multi statements (#2531)
Port PR: snowflakedb/Stored-Proc-Python-Connector#225 I noticed that get_results_from_sfqid assumes that fetchall returns a tuple when that's not necessarily the case. Fixing it here + adding a test that fails without the fix. ``` E AssertionError: assert [{'multiple s...ccessfully.'}] == [{'1': 1}] E At index 0 diff: {'multiple statement execution': 'Multiple statements executed successfully.'} != {'1': 1} E Full diff: E - [{'1': 1}] E + [{'multiple statement execution': 'Multiple statements executed successfully.'}] ```
1 parent 2226c90 commit 29b5b7e

File tree

4 files changed

+27
-8
lines changed

4 files changed

+27
-8
lines changed

DESCRIPTION.md

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -9,6 +9,7 @@ Source code is also available at: https://github.com/snowflakedb/snowflake-conne
99
# Release Notes
1010
- v3.18.0(TBD)
1111
- Added the `workload_identity_impersonation_path` parameter to support service account impersonation for Workload Identity Federation on GCP and AWS workloads only
12+
- Fixed `get_results_from_sfqid` when using `DictCursor` and executing multiple statements at once
1213

1314
- v3.17.3(September 02,2025)
1415
- Enhanced configuration file permission warning messages.

src/snowflake/connector/cursor.py

Lines changed: 1 addition & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1741,8 +1741,7 @@ def wait_until_ready() -> None:
17411741
self.connection.get_query_status_throw_if_error(
17421742
sfqid
17431743
) # Trigger an exception if query failed
1744-
klass = self.__class__
1745-
self._inner_cursor = klass(self.connection)
1744+
self._inner_cursor = SnowflakeCursor(self.connection)
17461745
self._sfqid = sfqid
17471746
self._prefetch_hook = wait_until_ready
17481747

test/integ/test_async.py

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -7,6 +7,7 @@
77
import pytest
88

99
from snowflake.connector import DatabaseError, ProgrammingError
10+
from snowflake.connector.cursor import DictCursor, SnowflakeCursor
1011

1112
# Mark all tests in this file to time out after 2 minutes to prevent hanging forever
1213
pytestmark = [pytest.mark.timeout(120), pytest.mark.skipolddriver]
@@ -17,14 +18,15 @@
1718
QueryStatus = None
1819

1920

20-
def test_simple_async(conn_cnx):
21+
@pytest.mark.parametrize("cursor_class", [SnowflakeCursor, DictCursor])
22+
def test_simple_async(conn_cnx, cursor_class):
2123
"""Simple test to that shows the most simple usage of fire and forget.
2224
2325
This test also makes sure that wait_until_ready function's sleeping is tested and
2426
that some fields are copied over correctly from the original query.
2527
"""
2628
with conn_cnx() as con:
27-
with con.cursor() as cur:
29+
with con.cursor(cursor_class) as cur:
2830
cur.execute_async("select count(*) from table(generator(timeLimit => 5))")
2931
cur.get_results_from_sfqid(cur.sfqid)
3032
assert len(cur.fetchall()) == 1

test/integ/test_multi_statement.py

Lines changed: 21 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -14,6 +14,7 @@
1414

1515
import snowflake.connector.cursor
1616
from snowflake.connector import ProgrammingError, errors
17+
from snowflake.connector.cursor import DictCursor, SnowflakeCursor
1718

1819
try: # pragma: no cover
1920
from snowflake.connector.constants import (
@@ -153,10 +154,11 @@ def test_binding_multi(conn_cnx, style: str, skip_to_last_set: bool):
153154
)
154155

155156

156-
def test_async_exec_multi(conn_cnx, skip_to_last_set: bool):
157+
@pytest.mark.parametrize("cursor_class", [SnowflakeCursor, DictCursor])
158+
def test_async_exec_multi(conn_cnx, cursor_class, skip_to_last_set: bool):
157159
"""Tests whether async execution query works within a multi-statement"""
158160
with conn_cnx() as con:
159-
with con.cursor() as cur:
161+
with con.cursor(cursor_class) as cur:
160162
cur.execute_async(
161163
"select 1; select 2; select count(*) from table(generator(timeLimit => 1)); select 'b';",
162164
num_statements=4,
@@ -165,14 +167,29 @@ def test_async_exec_multi(conn_cnx, skip_to_last_set: bool):
165167
assert con.is_still_running(con.get_query_status(q_id))
166168
_wait_while_query_running(con, q_id, sleep_time=1)
167169
with conn_cnx() as con:
168-
with con.cursor() as cur:
170+
with con.cursor(cursor_class) as cur:
169171
_wait_until_query_success(con, q_id, num_checks=3, sleep_per_check=1)
170172
assert con.get_query_status_throw_if_error(q_id) == QueryStatus.SUCCESS
171173

174+
if cursor_class == SnowflakeCursor:
175+
expected = [
176+
[(1,)],
177+
[(2,)],
178+
lambda x: len(x) == 1 and len(x[0]) == 1 and x[0][0] > 0,
179+
[("b",)],
180+
]
181+
elif cursor_class == DictCursor:
182+
expected = [
183+
[{"1": 1}],
184+
[{"2": 2}],
185+
lambda x: len(x) == 1 and len(x[0]) == 1 and x[0]["COUNT(*)"] > 0,
186+
[{"'B'": "b"}],
187+
]
188+
172189
cur.get_results_from_sfqid(q_id)
173190
_check_multi_statement_results(
174191
cur,
175-
checks=[[(1,)], [(2,)], lambda x: x > [(0,)], [("b",)]],
192+
checks=expected,
176193
skip_to_last_set=skip_to_last_set,
177194
)
178195

0 commit comments

Comments
 (0)