fix(bigquery): close GAPIC storage transport correctly to release sockets

Author: chalmerlowe

packages/google-cloud-bigquery/google/cloud/bigquery/dbapi/connection.py

@@ -84,7 +84,7 @@ def close(self):
 
         if self._owns_bqstorage_client:
             # There is no close() on the BQ Storage client itself.
-            self._bqstorage_client._transport.grpc_channel.close()
+            self._bqstorage_client._transport.close()
 
         for cursor_ in self._cursors_created:
             if not cursor_._closed:

packages/google-cloud-bigquery/google/cloud/bigquery/magics/magics.py

@@ -773,4 +773,4 @@ def _close_transports(client, bqstorage_client):
     """
     client.close()
     if bqstorage_client is not None:
-        bqstorage_client._transport.grpc_channel.close()
+        bqstorage_client._transport.close()

packages/google-cloud-bigquery/google/cloud/bigquery/table.py

@@ -2353,7 +2353,7 @@ def to_arrow(
                 progress_bar.close()
         finally:
             if owns_bqstorage_client:
-                bqstorage_client._transport.grpc_channel.close()  # type: ignore
+                bqstorage_client._transport.close()
 
         if record_batches and bqstorage_client is not None:
             return pyarrow.Table.from_batches(record_batches)

packages/google-cloud-bigquery/tests/unit/test_dbapi_connection.py

@@ -40,7 +40,7 @@ def _mock_bqstorage_client(self):
         from google.cloud import bigquery_storage
 
         mock_client = mock.create_autospec(bigquery_storage.BigQueryReadClient)
-        mock_client._transport = mock.Mock(spec=["channel"])
+        mock_client._transport = mock.Mock(spec=["channel", "close"])
         mock_client._transport.grpc_channel = mock.Mock(spec=["close"])
         return mock_client
 
@@ -176,7 +176,7 @@ def test_close_closes_all_created_bigquery_clients(self):
         connection.close()
 
         self.assertTrue(client.close.called)
-        self.assertTrue(bqstorage_client._transport.grpc_channel.close.called)
+        self.assertTrue(bqstorage_client._transport.close.called)
 
     def test_close_does_not_close_bigquery_clients_passed_to_it(self):
         pytest.importorskip("google.cloud.bigquery_storage")
@@ -187,7 +187,7 @@ def test_close_does_not_close_bigquery_clients_passed_to_it(self):
         connection.close()
 
         self.assertFalse(client.close.called)
-        self.assertFalse(bqstorage_client._transport.grpc_channel.close.called)
+        self.assertFalse(bqstorage_client._transport.close.called)
 
     def test_close_closes_all_created_cursors(self):
         connection = self._make_one(client=self._mock_client())

packages/google-cloud-bigquery/tests/unit/test_table.py

@@ -3048,7 +3048,7 @@ def test_to_arrow_iterable_w_bqstorage(self):
             self.assertEqual(record_batch, expected_record_batch)
 
         # Don't close the client if it was passed in.
-        bqstorage_client._transport.grpc_channel.close.assert_not_called()
+        bqstorage_client._transport.close.assert_not_called()
 
     def test_to_arrow(self):
         pytest.importorskip("numpy")
@@ -3424,7 +3424,7 @@ def test_to_arrow_w_bqstorage(self):
         self.assertEqual(actual_tbl.num_rows, total_rows)
 
         # Don't close the client if it was passed in.
-        bqstorage_client._transport.grpc_channel.close.assert_not_called()
+        bqstorage_client._transport.close.assert_not_called()
 
     def test_to_arrow_w_bqstorage_creates_client(self):
         pytest.importorskip("numpy")
@@ -3458,7 +3458,7 @@ def test_to_arrow_w_bqstorage_creates_client(self):
         )
         row_iterator.to_arrow(create_bqstorage_client=True)
         mock_client._ensure_bqstorage_client.assert_called_once()
-        bqstorage_client._transport.grpc_channel.close.assert_called_once()
+        bqstorage_client._transport.close.assert_called_once()
 
     def test_to_arrow_ensure_bqstorage_client_wo_bqstorage(self):
         pytest.importorskip("numpy")
@@ -3741,7 +3741,7 @@ def test_to_dataframe_iterable_w_bqstorage(self):
         self.assertEqual(len(got), total_pages)
 
         # Don't close the client if it was passed in.
-        bqstorage_client._transport.grpc_channel.close.assert_not_called()
+        bqstorage_client._transport.close.assert_not_called()
 
     def test_to_dataframe_iterable_w_bqstorage_max_results_warning(self):
         pytest.importorskip("numpy")
@@ -4807,7 +4807,7 @@ def test_to_dataframe_w_bqstorage_creates_client(self):
         )
         row_iterator.to_dataframe(create_bqstorage_client=True)
         mock_client._ensure_bqstorage_client.assert_called_once()
-        bqstorage_client._transport.grpc_channel.close.assert_called_once()
+        bqstorage_client._transport.close.assert_called_once()
 
     def test_to_dataframe_w_bqstorage_no_streams(self):
         pytest.importorskip("numpy")
@@ -4999,7 +4999,7 @@ def test_to_dataframe_w_bqstorage_nonempty(self):
         self.assertEqual(len(got.index), total_rows)
 
         # Don't close the client if it was passed in.
-        bqstorage_client._transport.grpc_channel.close.assert_not_called()
+        bqstorage_client._transport.close.assert_not_called()
 
     def test_to_dataframe_w_bqstorage_multiple_streams_return_unique_index(self):
         pytest.importorskip("numpy")
@@ -5421,7 +5421,7 @@ def test_to_dataframe_concat_categorical_dtype_w_pyarrow(self):
         )
 
         # Don't close the client if it was passed in.
-        bqstorage_client._transport.grpc_channel.close.assert_not_called()
+        bqstorage_client._transport.close.assert_not_called()
 
     def test_to_dataframe_geography_as_object(self):
         pandas = pytest.importorskip("pandas")