From 8eebb9dbcd35aa7eda4d2958f678b6b7d965d3c9 Mon Sep 17 00:00:00 2001
From: Kien Truong <duckientruong@gmail.com>
Date: Sat, 28 Sep 2024 10:09:55 +0700
Subject: [PATCH] fix: avoid blocking in download thread when using BQ Storage
 API

This prevents a deadlock between the main thead and download threads
when the threadpool is shutdown prematurely.
---
 google/cloud/bigquery/_pandas_helpers.py | 11 ++++++++---
 1 file changed, 8 insertions(+), 3 deletions(-)

diff --git a/google/cloud/bigquery/_pandas_helpers.py b/google/cloud/bigquery/_pandas_helpers.py
index bf7d10c0f..050672531 100644
--- 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):