Test 1493 (#1515)

Co-authored-by: Kevin C <[email protected]>
Co-authored-by: Mark Keller <[email protected]>

Author: 3 people

DESCRIPTION.md

@@ -13,7 +13,7 @@ Source code is also available at: https://github.com/snowflakedb/snowflake-conne
   - Fixed a bug in which we cannot call `SnowflakeCursor.nextset` before fetching the result of the first query if the cursor runs an async multistatement query.
   - Bumped vendored library urllib3 to 1.26.15
   - Bumped vendored library requests to 2.29.0
-
+  - Fixed a bug when `_prefetch_hook()` was not called before yielding results of `execute_async()`.
 
 - v3.0.3(April 20, 2023)
   - Fixed a bug that prints error in logs for GET command on GCS.

src/snowflake/connector/cursor.py

@@ -21,7 +21,6 @@
     TYPE_CHECKING,
     Any,
     Callable,
-    Generator,
     Iterator,
     NamedTuple,
     NoReturn,
@@ -1242,37 +1241,30 @@ def executemany(
 
         return self
 
-    def _result_iterator(
-        self,
-    ) -> Generator[dict, None, None] | Generator[tuple, None, None]:
-        """Yields the elements from _result and raises an exception when appropriate."""
-        try:
-            for _next in self._result:
-                if isinstance(_next, Exception):
-                    Error.errorhandler_wrapper_from_ready_exception(
-                        self._connection,
-                        self,
-                        _next,
-                    )
-                self._rownumber += 1
-                yield _next
-        except TypeError as err:
-            if self._result_state == ResultState.DEFAULT:
-                raise err
-            else:
-                yield None
-
     def fetchone(self) -> dict | tuple | None:
         """Fetches one row."""
         if self._prefetch_hook is not None:
             self._prefetch_hook()
         if self._result is None and self._result_set is not None:
             self._result = iter(self._result_set)
             self._result_state = ResultState.VALID
+
         try:
-            return next(self._result_iterator())
-        except StopIteration:
-            return None
+            _next = next(self._result, None)
+            if isinstance(_next, Exception):
+                Error.errorhandler_wrapper_from_ready_exception(
+                    self._connection,
+                    self,
+                    _next,
+                )
+            if _next is not None:
+                self._rownumber += 1
+            return _next
+        except TypeError as err:
+            if self._result_state == ResultState.DEFAULT:
+                raise err
+            else:
+                return None
 
     def fetchmany(self, size: int | None = None) -> list[tuple] | list[dict]:
         """Fetches the number of specified rows."""
@@ -1369,11 +1361,11 @@ def reset(self, closing: bool = False) -> None:
 
     def __iter__(self) -> Iterator[dict] | Iterator[tuple]:
         """Iteration over the result set."""
-        # set _result if _result_set is not None
-        if self._result is None and self._result_set is not None:
-            self._result = iter(self._result_set)
-            self._result_state = ResultState.VALID
-        return self._result_iterator()
+        while True:
+            _next = self.fetchone()
+            if _next is None:
+                break
+            yield _next
 
     def __cancel_query(self, query) -> None:
         if self._sequence_counter >= 0 and not self.is_closed():

test/integ/test_async.py

@@ -35,6 +35,25 @@ def test_simple_async(conn_cnx):
             assert cur.description
 
 
+def test_async_result_iteration(conn_cnx):
+    """Test yielding results of an async query.
+
+    Ensures that wait_until_ready is also called in __iter__() via _prefetch_hook().
+    """
+
+    def result_generator(query):
+        with conn_cnx() as con:
+            with con.cursor() as cur:
+                cur.execute_async(query)
+                cur.get_results_from_sfqid(cur.sfqid)
+                yield from cur
+
+    gen = result_generator("select count(*) from table(generator(timeLimit => 5))")
+    assert next(gen)
+    with pytest.raises(StopIteration):
+        next(gen)
+
+
 def test_async_exec(conn_cnx):
     """Tests whether simple async query execution works.