fix(bigquery): dbapi socket leak (#17921)

## Problem
The system test `test_dbapi_connection_does_not_leak_sockets` fails
intermittently because network connections remain in the `ESTABLISHED`
state even after `connection.close()` is called. This is likely due to
the underlying gRPC (remote procedure call) transport not being
immediately closed, or delays in Garbage Collection (cleaning up unused
memory).

## Solution
This Pull Request updates `Connection.close()` to explicitly close the
underlying gRPC channel (`transport._grpc_channel`) if it exists on the
BigQuery Storage client. This ensures that the connection is closed
immediately and resources are released deterministically.

## Notes to Reviewers
- This change explicitly targets the underlying gRPC channel to bypass
potential delays caused by transport interceptors or non-deterministic
garbage collection.
- It addresses intermittent test failures caused by lingering
`ESTABLISHED` sockets.
- Also includes several linting revisions (sorry).

Sample debug info from traceback:
```
E           --- Socket Leak Debug Info ---
E           Start Count: 6
E           End Count: 7
E           Current Connections:
E           Status: ESTABLISHED, Laddr: addr(ip='10.138.0.14', port=48550), Raddr: addr(ip='74.125.195.95', port=443)
E           Status: ESTABLISHED, Laddr: addr(ip='10.138.0.14', port=53976), Raddr: addr(ip='173.194.43.95', port=443)
E           Status: ESTABLISHED, Laddr: addr(ip='10.138.0.14', port=54818), Raddr: addr(ip='74.125.199.95', port=443)
E           Status: ESTABLISHED, Laddr: addr(ip='10.138.0.14', port=43762), Raddr: addr(ip='142.251.188.95', port=443)
E           Status: ESTABLISHED, Laddr: addr(ip='10.138.0.14', port=53946), Raddr: addr(ip='173.194.43.95', port=443)
E           Status: ESTABLISHED, Laddr: addr(ip='10.138.0.14', port=54792), Raddr: addr(ip='74.125.199.95', port=443)
E           Status: ESTABLISHED, Laddr: addr(ip='10.138.0.14', port=48524), Raddr: addr(ip='74.125.195.95', port=443)
```

Author: chalmerlowe

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

@@ -17,8 +17,7 @@
 import weakref
 
 from google.cloud import bigquery
-from google.cloud.bigquery.dbapi import cursor
-from google.cloud.bigquery.dbapi import _helpers
+from google.cloud.bigquery.dbapi import _helpers, cursor
 
 
 @_helpers.raise_on_closed("Operating on a closed connection.")
@@ -84,7 +83,17 @@ def close(self):
 
         if self._owns_bqstorage_client:
             # There is no close() on the BQ Storage client itself.
-            self._bqstorage_client._transport.close()
+            transport = self._bqstorage_client.transport
+            transport.close()
+
+            # Ensure the underlying gRPC channel is closed explicitly.
+            # This is important because the transport might be wrapped by interceptors
+            # which may not propagate the close call to the underlying channel immediately,
+            # or Garbage Collection delays might keep sockets open longer than expected,
+            # leading to intermittent socket leaks in tests.
+            channel = getattr(transport, "_grpc_channel", None)
+            if channel is not None:
+                channel.close()
 
         for cursor_ in self._cursors_created:
             if not cursor_._closed:

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

@@ -17,8 +17,8 @@
 from __future__ import annotations
 
 import collections
-from collections import abc as collections_abc
 import re
+from collections import abc as collections_abc
 from typing import Optional
 
 try:
@@ -29,11 +29,9 @@
     # Having BQ Storage available implies that pyarrow >=1.0.0 is available, too.
     _ARROW_COMPRESSION_SUPPORT = True
 
-from google.cloud.bigquery import job
-from google.cloud.bigquery.dbapi import _helpers
-from google.cloud.bigquery.dbapi import exceptions
 import google.cloud.exceptions  # type: ignore
-
+from google.cloud.bigquery import job
+from google.cloud.bigquery.dbapi import _helpers, exceptions
 
 # Per PEP 249: A 7-item sequence containing information describing one result
 # column. The first two items (name and type_code) are mandatory, the other
@@ -105,6 +103,15 @@ def close(self):
         """Mark the cursor as closed, preventing its further use."""
         self._closed = True
 
+        # Explicitly release references to query data and rows.
+        # This is important because these objects may hold references to underlying
+        # transports or gRPC streams (especially when using BQ Storage API).
+        # Clearing them ensures that resources can be garbage collected deterministically,
+        # preventing intermittent socket leaks (e.g., ESTABLISHED connections) in tests
+        # and long-running applications.
+        self._query_data = None
+        self._query_rows = None
+
     def _set_description(self, schema):
         """Set description from schema.
 

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

@@ -13,10 +13,11 @@
 #  limitations under the License.
 
 import gc
-import pytest
 import unittest
 from unittest import mock
 
+import pytest
+
 
 class TestConnection(unittest.TestCase):
     @staticmethod
@@ -40,8 +41,8 @@ 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", "close"])
-        mock_client._transport.grpc_channel = mock.Mock(spec=["close"])
+        mock_client.transport.close = mock.Mock()
+        mock_client.transport._grpc_channel = mock.Mock(spec=["close"])
         return mock_client
 
     def test_ctor_wo_bqstorage_client(self):
@@ -77,8 +78,7 @@ def test_ctor_w_bqstorage_client(self):
 
     @mock.patch("google.cloud.bigquery.Client", autospec=True)
     def test_connect_wo_client(self, mock_client):
-        from google.cloud.bigquery.dbapi import connect
-        from google.cloud.bigquery.dbapi import Connection
+        from google.cloud.bigquery.dbapi import Connection, connect
 
         connection = connect()
         self.assertIsInstance(connection, Connection)
@@ -87,8 +87,7 @@ def test_connect_wo_client(self, mock_client):
 
     def test_connect_w_client(self):
         pytest.importorskip("google.cloud.bigquery_storage")
-        from google.cloud.bigquery.dbapi import connect
-        from google.cloud.bigquery.dbapi import Connection
+        from google.cloud.bigquery.dbapi import Connection, connect
 
         mock_client = self._mock_client()
         mock_bqstorage_client = self._mock_bqstorage_client()
@@ -103,8 +102,7 @@ def test_connect_w_client(self):
 
     def test_connect_w_both_clients(self):
         pytest.importorskip("google.cloud.bigquery_storage")
-        from google.cloud.bigquery.dbapi import connect
-        from google.cloud.bigquery.dbapi import Connection
+        from google.cloud.bigquery.dbapi import Connection, connect
 
         mock_client = self._mock_client()
         mock_bqstorage_client = self._mock_bqstorage_client()
@@ -124,8 +122,7 @@ def test_connect_w_both_clients(self):
 
     def test_connect_prefer_bqstorage_client_false(self):
         pytest.importorskip("google.cloud.bigquery_storage")
-        from google.cloud.bigquery.dbapi import connect
-        from google.cloud.bigquery.dbapi import Connection
+        from google.cloud.bigquery.dbapi import Connection, connect
 
         mock_client = self._mock_client()
         mock_bqstorage_client = self._mock_bqstorage_client()
@@ -155,11 +152,16 @@ def test_raises_error_if_closed(self):
             ):
                 getattr(connection, method)()
 
-    def test_close_closes_all_created_bigquery_clients(self):
+    def _verify_close_closes_all_created_bigquery_clients(
+        self, simulate_no_grpc_channel=False
+    ):
         pytest.importorskip("google.cloud.bigquery_storage")
         client = self._mock_client()
         bqstorage_client = self._mock_bqstorage_client()
 
+        if simulate_no_grpc_channel:
+            bqstorage_client.transport._grpc_channel = None
+
         client_patcher = mock.patch(
             "google.cloud.bigquery.dbapi.connection.bigquery.Client",
             return_value=client,
@@ -176,7 +178,20 @@ def test_close_closes_all_created_bigquery_clients(self):
         connection.close()
 
         self.assertTrue(client.close.called)
-        self.assertTrue(bqstorage_client._transport.close.called)
+        self.assertTrue(bqstorage_client.transport.close.called)
+
+        if not simulate_no_grpc_channel:
+            self.assertTrue(bqstorage_client.transport._grpc_channel.close.called)
+
+    def test_close_closes_all_created_bigquery_clients(self):
+        self._verify_close_closes_all_created_bigquery_clients(
+            simulate_no_grpc_channel=False
+        )
+
+    def test_close_closes_all_created_bigquery_clients_no_grpc_channel(self):
+        self._verify_close_closes_all_created_bigquery_clients(
+            simulate_no_grpc_channel=True
+        )
 
     def test_close_does_not_close_bigquery_clients_passed_to_it(self):
         pytest.importorskip("google.cloud.bigquery_storage")
@@ -187,7 +202,8 @@ def test_close_does_not_close_bigquery_clients_passed_to_it(self):
         connection.close()
 
         self.assertFalse(client.close.called)
-        self.assertFalse(bqstorage_client._transport.close.called)
+        self.assertFalse(bqstorage_client.transport.close.called)
+        self.assertFalse(bqstorage_client.transport._grpc_channel.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_dbapi_cursor.py

@@ -17,10 +17,8 @@
 import unittest
 from unittest import mock
 
-import pytest
-
 import google.cloud.bigquery.table as bq_table
-
+import pytest
 from google.api_core import exceptions
 
 from tests.unit.helpers import _to_pyarrow
@@ -180,8 +178,7 @@ def _mock_results(self, total_rows=0, schema=None, num_dml_affected_rows=None):
         return mock_results
 
     def test_ctor(self):
-        from google.cloud.bigquery.dbapi import connect
-        from google.cloud.bigquery.dbapi import Cursor
+        from google.cloud.bigquery.dbapi import Cursor, connect
 
         connection = connect(self._mock_client())
         cursor = self._make_one(connection)
@@ -193,9 +190,17 @@ def test_close(self):
 
         connection = connect(self._mock_client())
         cursor = connection.cursor()
-        # close() is a no-op, there is nothing to test.
+
+        # Simulate that cursor has data/references
+        cursor._query_data = mock.Mock()
+        cursor._query_rows = mock.Mock()
+
         cursor.close()
 
+        self.assertTrue(cursor._closed)
+        self.assertIsNone(cursor._query_data)
+        self.assertIsNone(cursor._query_rows)
+
     def test_raises_error_if_closed(self):
         from google.cloud.bigquery.dbapi import connect
         from google.cloud.bigquery.dbapi.exceptions import ProgrammingError
@@ -519,8 +524,8 @@ def test_execute_w_default_config(self):
         self.assertIsNone(used_config)
 
     def test_execute_custom_job_config_wo_default_config(self):
-        from google.cloud.bigquery.dbapi import connect
         from google.cloud.bigquery import job
+        from google.cloud.bigquery.dbapi import connect
 
         config = job.QueryJobConfig(use_legacy_sql=True)
         client = self._mock_client(rows=[], num_dml_affected_rows=0)
@@ -533,8 +538,8 @@ def test_execute_custom_job_config_wo_default_config(self):
         self.assertEqual(kwargs["job_config"], config)
 
     def test_execute_custom_job_config_w_default_config(self):
-        from google.cloud.bigquery.dbapi import connect
         from google.cloud.bigquery import job
+        from google.cloud.bigquery.dbapi import connect
 
         client = self._mock_client(rows=[], num_dml_affected_rows=0)
         connection = connect(client)
@@ -563,8 +568,8 @@ def test_execute_w_dml(self):
         self.assertEqual(rows, [])
 
     def test_execute_w_query(self):
-        from google.cloud.bigquery.schema import SchemaField
         from google.cloud.bigquery import dbapi
+        from google.cloud.bigquery.schema import SchemaField
 
         connection = dbapi.connect(
             self._mock_client(
@@ -607,9 +612,9 @@ def test_execute_w_query(self):
         self.assertIsNone(row)
 
     def test_execute_w_query_dry_run(self):
+        from google.cloud.bigquery import dbapi
         from google.cloud.bigquery.job import QueryJobConfig
         from google.cloud.bigquery.schema import SchemaField
-        from google.cloud.bigquery import dbapi
 
         connection = dbapi.connect(
             self._mock_client(
@@ -637,10 +642,8 @@ def test_execute_w_query_dry_run(self):
 
     def test_execute_raises_if_result_raises(self):
         import google.cloud.exceptions
-
         from google.cloud.bigquery import client
-        from google.cloud.bigquery.dbapi import connect
-        from google.cloud.bigquery.dbapi import exceptions
+        from google.cloud.bigquery.dbapi import connect, exceptions
 
         client = mock.create_autospec(client.Client)
         client.query_and_wait.side_effect = google.cloud.exceptions.GoogleCloudError("")
@@ -702,7 +705,7 @@ def test_query_job_wo_execute(self):
         self.assertIsNone(cursor.query_job)
 
     def test_query_job_w_execute(self):
-        from google.cloud.bigquery import dbapi, QueryJob
+        from google.cloud.bigquery import QueryJob, dbapi
 
         connection = dbapi.connect(self._mock_client())
         cursor = connection.cursor()
@@ -722,7 +725,7 @@ def test_query_job_w_execute_no_job(self):
         self.assertIsNone(cursor.query_job)
 
     def test_query_job_w_executemany(self):
-        from google.cloud.bigquery import dbapi, QueryJob
+        from google.cloud.bigquery import QueryJob, dbapi
 
         connection = dbapi.connect(self._mock_client())
         cursor = connection.cursor()
@@ -932,8 +935,8 @@ def test__extract_types(inp, expect):
     ],
 )
 def test__extract_types_fail(match, inp):
-    from google.cloud.bigquery.dbapi.cursor import _extract_types as et
     from google.cloud.bigquery.dbapi import exceptions
+    from google.cloud.bigquery.dbapi.cursor import _extract_types as et
 
     with pytest.raises(exceptions.ProgrammingError, match=match):
         et(inp)