Fix return type of async DictCursor queries

Author: sfc-gh-bchinn

src/snowflake/connector/cursor.py

@@ -414,7 +414,7 @@ def __init__(
         self._first_chunk_time = None
 
         self._log_max_query_length = connection.log_max_query_length
-        self._inner_cursor: SnowflakeCursor | None = None
+        self._inner_cursor: SnowflakeCursorBase | None = None
         self._prefetch_hook = None
         self._rownumber: int | None = None
 
@@ -1720,20 +1720,31 @@ def wait_until_ready() -> None:
             # Unset this function, so that we don't block anymore
             self._prefetch_hook = None
 
-            if (
-                self._inner_cursor._total_rowcount == 1
-                and self._inner_cursor.fetchall()
-                == [("Multiple statements executed successfully.",)]
+            if self._inner_cursor._total_rowcount == 1 and _is_successful_multi_stmt(
+                self._inner_cursor.fetchall()
             ):
                 url = f"/queries/{sfqid}/result"
                 ret = self._connection.rest.request(url=url, method="get")
                 if "data" in ret and "resultIds" in ret["data"]:
                     self._init_multi_statement_results(ret["data"])
 
+        def _is_successful_multi_stmt(rows: list[Any]) -> True:
+            if len(rows) != 1:
+                return False
+            row = rows[0]
+            if isinstance(row, tuple):
+                return row == ("Multiple statements executed successfully.",)
+            elif isinstance(row, dict):
+                return row == {
+                    "multiple statement execution": "Multiple statements executed successfully."
+                }
+            else:
+                return False
+
         self.connection.get_query_status_throw_if_error(
             sfqid
         )  # Trigger an exception if query failed
-        self._inner_cursor = SnowflakeCursor(self.connection)
+        self._inner_cursor = self.__class__(self.connection)
         self._sfqid = sfqid
         self._prefetch_hook = wait_until_ready
 

test/integ/test_async.py

@@ -18,8 +18,10 @@
     QueryStatus = None
 
 
-@pytest.mark.parametrize("cursor_class", [SnowflakeCursor, DictCursor])
-def test_simple_async(conn_cnx, cursor_class):
+@pytest.mark.parametrize(
+    "cursor_class, row_type", [(SnowflakeCursor, tuple), (DictCursor, dict)]
+)
+def test_simple_async(conn_cnx, cursor_class, row_type):
     """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
@@ -29,7 +31,9 @@ def test_simple_async(conn_cnx, cursor_class):
         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
+            rows = cur.fetchall()
+            assert len(rows) == 1
+            assert isinstance(rows[0], row_type)
             assert cur.rowcount
             assert cur.description