fix: avoid blocking in download thread when using BQ Storage API (#2034)

kien-truong · chalmerlowe · tswast · web-flow · commit 54c8d07f06a8 · 2025-02-03T14:37:26.000-06:00
This prevents a deadlock between the main thead and download threads
when the threadpool is shutdown prematurely.

Co-authored-by: Chalmer Lowe &lt;chalmerlowe@google.com&gt;
Co-authored-by: Tim Sweña (Swast) &lt;swast@google.com&gt;
diff --git a/google/cloud/bigquery/_pandas_helpers.py b/google/cloud/bigquery/_pandas_helpers.py
@@ -796,10 +796,15 @@ def _download_table_bqstorage_stream(
         rowstream = reader.rows(session)
 
     for page in rowstream.pages:
-        if download_state.done:
-            return
         item = page_to_item(page)
-        worker_queue.put(item)
+        while True:
+            if download_state.done:
+                return
+            try:
+                worker_queue.put(item, timeout=_PROGRESS_INTERVAL)
+                break
+            except queue.Full:  # pragma: NO COVER
+                continue
 
 
 def _nowait(futures):
