fix: close bqstorage client if created by to_dataframe/to_arrow

Author: tswast

bigquery/google/cloud/bigquery/table.py

@@ -1519,25 +1519,31 @@ def to_arrow(
         if pyarrow is None:
             raise ValueError(_NO_PYARROW_ERROR)
 
+        owns_bqstorage_client = False
         if not bqstorage_client and create_bqstorage_client:
+            owns_bqstorage_client = True
             bqstorage_client = self.client._create_bqstorage_client()
 
-        progress_bar = self._get_progress_bar(progress_bar_type)
+        try:
+            progress_bar = self._get_progress_bar(progress_bar_type)
 
-        record_batches = []
-        for record_batch in self._to_arrow_iterable(bqstorage_client=bqstorage_client):
-            record_batches.append(record_batch)
+            record_batches = []
+            for record_batch in self._to_arrow_iterable(bqstorage_client=bqstorage_client):
+                record_batches.append(record_batch)
 
-            if progress_bar is not None:
-                # In some cases, the number of total rows is not populated
-                # until the first page of rows is fetched. Update the
-                # progress bar's total to keep an accurate count.
-                progress_bar.total = progress_bar.total or self.total_rows
-                progress_bar.update(record_batch.num_rows)
+                if progress_bar is not None:
+                    # In some cases, the number of total rows is not populated
+                    # until the first page of rows is fetched. Update the
+                    # progress bar's total to keep an accurate count.
+                    progress_bar.total = progress_bar.total or self.total_rows
+                    progress_bar.update(record_batch.num_rows)
 
-        if progress_bar is not None:
-            # Indicate that the download has finished.
-            progress_bar.close()
+            if progress_bar is not None:
+                # Indicate that the download has finished.
+                progress_bar.close()
+        finally:
+            if owns_bqstorage_client:
+                bqstorage_client.transport.channel.close()
 
         if record_batches:
             return pyarrow.Table.from_batches(record_batches)
@@ -1655,35 +1661,42 @@ def to_dataframe(
         if dtypes is None:
             dtypes = {}
 
-        if not bqstorage_client and create_bqstorage_client:
-            bqstorage_client = self.client._create_bqstorage_client()
-
-        if bqstorage_client and self.max_results is not None:
+        if (bqstorage_client or create_bqstorage_client) and self.max_results is not None:
             warnings.warn(
                 "Cannot use bqstorage_client if max_results is set, "
                 "reverting to fetching data with the tabledata.list endpoint.",
                 stacklevel=2,
             )
+            create_bqstorage_client = False
             bqstorage_client = None
 
-        progress_bar = self._get_progress_bar(progress_bar_type)
+        owns_bqstorage_client = False
+        if not bqstorage_client and create_bqstorage_client:
+            owns_bqstorage_client = True
+            bqstorage_client = self.client._create_bqstorage_client()
 
-        frames = []
-        for frame in self._to_dataframe_iterable(
-            bqstorage_client=bqstorage_client, dtypes=dtypes
-        ):
-            frames.append(frame)
+        try:
+            progress_bar = self._get_progress_bar(progress_bar_type)
+
+            frames = []
+            for frame in self._to_dataframe_iterable(
+                bqstorage_client=bqstorage_client, dtypes=dtypes
+            ):
+                frames.append(frame)
+
+                if progress_bar is not None:
+                    # In some cases, the number of total rows is not populated
+                    # until the first page of rows is fetched. Update the
+                    # progress bar's total to keep an accurate count.
+                    progress_bar.total = progress_bar.total or self.total_rows
+                    progress_bar.update(len(frame))
 
             if progress_bar is not None:
-                # In some cases, the number of total rows is not populated
-                # until the first page of rows is fetched. Update the
-                # progress bar's total to keep an accurate count.
-                progress_bar.total = progress_bar.total or self.total_rows
-                progress_bar.update(len(frame))
-
-        if progress_bar is not None:
-            # Indicate that the download has finished.
-            progress_bar.close()
+                # Indicate that the download has finished.
+                progress_bar.close()
+        finally:
+            if owns_bqstorage_client:
+                bqstorage_client.transport.channel.close()
 
         # Avoid concatting an empty list.
         if not frames:

bigquery/tests/unit/test_table.py

@@ -26,8 +26,12 @@
 
 try:
     from google.cloud import bigquery_storage_v1beta1
+    from google.cloud.bigquery_storage_v1beta1.gapic.transports import (
+        big_query_storage_grpc_transport,
+    )
 except ImportError:  # pragma: NO COVER
     bigquery_storage_v1beta1 = None
+    big_query_storage_grpc_transport = None
 
 try:
     import pandas
@@ -1817,6 +1821,9 @@ def test_to_arrow_w_bqstorage(self):
         bqstorage_client = mock.create_autospec(
             bigquery_storage_v1beta1.BigQueryStorageClient
         )
+        bqstorage_client.transport = mock.create_autospec(
+            big_query_storage_grpc_transport.BigQueryStorageGrpcTransport
+        )
         streams = [
             # Use two streams we want to check frames are read from each stream.
             {"name": "/projects/proj/dataset/dset/tables/tbl/streams/1234"},
@@ -1882,6 +1889,9 @@ def test_to_arrow_w_bqstorage(self):
         total_rows = expected_num_rows * total_pages
         self.assertEqual(actual_tbl.num_rows, total_rows)
 
+        # Don't close the client if it was passed in.
+        bqstorage_client.transport.channel.close.assert_not_called()
+
     @unittest.skipIf(pyarrow is None, "Requires `pyarrow`")
     @unittest.skipIf(
         bigquery_storage_v1beta1 is None, "Requires `google-cloud-bigquery-storage`"
@@ -1894,6 +1904,9 @@ def test_to_arrow_w_bqstorage_creates_client(self):
         bqstorage_client = mock.create_autospec(
             bigquery_storage_v1beta1.BigQueryStorageClient
         )
+        bqstorage_client.transport = mock.create_autospec(
+            big_query_storage_grpc_transport.BigQueryStorageGrpcTransport
+        )
         mock_client._create_bqstorage_client.return_value = bqstorage_client
         session = bigquery_storage_v1beta1.types.ReadSession()
         bqstorage_client.create_read_session.return_value = session
@@ -1910,6 +1923,7 @@ def test_to_arrow_w_bqstorage_creates_client(self):
         )
         row_iterator.to_arrow(create_bqstorage_client=True)
         mock_client._create_bqstorage_client.assert_called_once()
+        bqstorage_client.transport.channel.close.assert_called_once()
 
     @unittest.skipIf(pyarrow is None, "Requires `pyarrow`")
     @unittest.skipIf(
@@ -2321,6 +2335,43 @@ def test_to_dataframe_max_results_w_bqstorage_warning(self):
         ]
         self.assertEqual(len(matches), 1, msg="User warning was not emitted.")
 
+    @unittest.skipIf(pandas is None, "Requires `pandas`")
+    def test_to_dataframe_max_results_w_create_bqstorage_warning(self):
+        from google.cloud.bigquery.schema import SchemaField
+
+        schema = [
+            SchemaField("name", "STRING", mode="REQUIRED"),
+            SchemaField("age", "INTEGER", mode="REQUIRED"),
+        ]
+        rows = [
+            {"f": [{"v": "Phred Phlyntstone"}, {"v": "32"}]},
+            {"f": [{"v": "Bharney Rhubble"}, {"v": "33"}]},
+        ]
+        path = "/foo"
+        api_request = mock.Mock(return_value={"rows": rows})
+        mock_client = _mock_client()
+
+        row_iterator = self._make_one(
+            client=mock_client,
+            api_request=api_request,
+            path=path,
+            schema=schema,
+            max_results=42,
+        )
+
+        with warnings.catch_warnings(record=True) as warned:
+            row_iterator.to_dataframe(create_bqstorage_client=True)
+
+        matches = [
+            warning
+            for warning in warned
+            if warning.category is UserWarning
+            and "cannot use bqstorage_client" in str(warning).lower()
+            and "tabledata.list" in str(warning)
+        ]
+        self.assertEqual(len(matches), 1, msg="User warning was not emitted.")
+        mock_client._create_bqstorage_client.assert_not_called()
+
     @unittest.skipIf(pandas is None, "Requires `pandas`")
     @unittest.skipIf(
         bigquery_storage_v1beta1 is None, "Requires `google-cloud-bigquery-storage`"
@@ -2333,6 +2384,9 @@ def test_to_dataframe_w_bqstorage_creates_client(self):
         bqstorage_client = mock.create_autospec(
             bigquery_storage_v1beta1.BigQueryStorageClient
         )
+        bqstorage_client.transport = mock.create_autospec(
+            big_query_storage_grpc_transport.BigQueryStorageGrpcTransport
+        )
         mock_client._create_bqstorage_client.return_value = bqstorage_client
         session = bigquery_storage_v1beta1.types.ReadSession()
         bqstorage_client.create_read_session.return_value = session
@@ -2349,6 +2403,7 @@ def test_to_dataframe_w_bqstorage_creates_client(self):
         )
         row_iterator.to_dataframe(create_bqstorage_client=True)
         mock_client._create_bqstorage_client.assert_called_once()
+        bqstorage_client.transport.channel.close.assert_called_once()
 
     @unittest.skipIf(pandas is None, "Requires `pandas`")
     @unittest.skipIf(
@@ -2485,6 +2540,9 @@ def test_to_dataframe_w_bqstorage_nonempty(self):
         bqstorage_client = mock.create_autospec(
             bigquery_storage_v1beta1.BigQueryStorageClient
         )
+        bqstorage_client.transport = mock.create_autospec(
+            big_query_storage_grpc_transport.BigQueryStorageGrpcTransport
+        )
         streams = [
             # Use two streams we want to check frames are read from each stream.
             {"name": "/projects/proj/dataset/dset/tables/tbl/streams/1234"},
@@ -2539,6 +2597,9 @@ def test_to_dataframe_w_bqstorage_nonempty(self):
         total_rows = len(page_items) * total_pages
         self.assertEqual(len(got.index), total_rows)
 
+        # Don't close the client if it was passed in.
+        bqstorage_client.transport.channel.close.assert_not_called()
+
     @unittest.skipIf(pandas is None, "Requires `pandas`")
     @unittest.skipIf(
         bigquery_storage_v1beta1 is None, "Requires `google-cloud-bigquery-storage`"