SNOW-725840 Async support for arrow fetch functions (#1399)

* SNOW-725840 Async support for arrow fetch functions
[New Feature] Enabled fetch_arrow_batches and fetch_arrow_all to work for async query results.

Author: Harsh Chouraria

src/snowflake/connector/cursor.py

@@ -995,6 +995,8 @@ def query_result(self, qid):
 
     def fetch_arrow_batches(self) -> Iterator[Table]:
         self.check_can_use_arrow_resultset()
+        if self._prefetch_hook is not None:
+            self._prefetch_hook()
         if self._query_result_format != "arrow":
             raise NotSupportedError
         self._log_telemetry_job_data(
@@ -1004,6 +1006,8 @@ def fetch_arrow_batches(self) -> Iterator[Table]:
 
     def fetch_arrow_all(self) -> Table | None:
         self.check_can_use_arrow_resultset()
+        if self._prefetch_hook is not None:
+            self._prefetch_hook()
         if self._query_result_format != "arrow":
             raise NotSupportedError
         self._log_telemetry_job_data(TelemetryField.ARROW_FETCH_ALL, TelemetryData.TRUE)

test/integ/pandas/test_arrow_pandas.py

@@ -982,7 +982,25 @@ def test_execute_async_and_fetch_pandas_batches(conn_cnx):
                 assert r_sync.values == r_async.values
 
 
-def test_simple_async_arrow(conn_cnx):
+def test_execute_async_and_fetch_arrow_batches(conn_cnx):
+    """Test fetching result of an asynchronous query as batches of arrow tables"""
+
+    with conn_cnx() as cnx:
+        with cnx.cursor() as cur:
+            cur.execute("select 1/2")
+            res_sync = cur.fetch_arrow_batches()
+
+            result = cur.execute_async("select 1/2")
+            cur.get_results_from_sfqid(result["queryId"])
+            res_async = cur.fetch_arrow_batches()
+
+            assert res_sync is not None
+            assert res_async is not None
+            for r_sync, r_async in zip(res_sync, res_async):
+                assert r_sync == r_async
+
+
+def test_simple_async_pandas(conn_cnx):
     """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
@@ -997,6 +1015,17 @@ def test_simple_async_arrow(conn_cnx):
             assert cur.description
 
 
+def test_simple_async_arrow(conn_cnx):
+    """Simple test for async fetch_arrow_all"""
+    with conn_cnx() as con:
+        with con.cursor() as cur:
+            cur.execute_async("select count(*) from table(generator(timeLimit => 5))")
+            cur.get_results_from_sfqid(cur.sfqid)
+            assert len(cur.fetch_arrow_all()) == 1
+            assert cur.rowcount
+            assert cur.description
+
+
 @pytest.mark.parametrize(
     "use_decimal,expected",
     [