fix(bigquery): Fix bigquery socket leak

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

@@ -78,10 +78,14 @@ def close(self):
         """
         self._closed = True
 
-        if self._owns_client:
+        for cursor_ in list(self._cursors_created):
+            if not cursor_._closed:
+                cursor_.close()
+
+        if self._owns_client and self._client is not None:
             self._client.close()
 
-        if self._owns_bqstorage_client:
+        if self._owns_bqstorage_client and self._bqstorage_client is not None:
             # There is no close() on the BQ Storage client itself.
             transport = self._bqstorage_client.transport
             transport.close()
@@ -95,10 +99,6 @@ def close(self):
             if channel is not None:
                 channel.close()
 
-        for cursor_ in self._cursors_created:
-            if not cursor_._closed:
-                cursor_.close()
-
     def commit(self):
         """No-op, but for consistency raise an error if connection is closed."""
 

packages/google-cloud-bigquery/tests/system/test_client.py

@@ -2182,58 +2182,101 @@ def test_dbapi_dry_run_query(self):
 
     def test_dbapi_connection_does_not_leak_sockets(self):
         pytest.importorskip("google.cloud.bigquery_storage")
-        current_process = psutil.Process()
-        conn_start = current_process.net_connections()
-        conn_count_start = len(conn_start)
 
-        with helpers.patch_tracked_requests():
-            # Provide no explicit clients, so that the connection will create and own them.
-            connection = dbapi.connect()
-            cursor = connection.cursor()
+        import google.api_core.grpc_helpers
+        import requests
 
-            cursor.execute(
-                """
-                SELECT id, `by`, timestamp
-                FROM `bigquery-public-data.hacker_news.full`
-                ORDER BY `id` ASC
-                LIMIT 100000
-            """
-            )
-            rows = cursor.fetchall()
-            self.assertEqual(len(rows), 100000)
+        # Track HTTP Sessions
+        created_sessions = set()
+        closed_sessions = set()
 
-            connection.close()
-        import gc
+        original_session_init = requests.Session.__init__
 
-        gc.collect()
-        for _ in range(30):  # Wait up to 3 seconds
-            conn_end = current_process.net_connections()
-            conn_count_end = len(conn_end)
-            if conn_count_end <= conn_count_start:
-                break
-            time.sleep(0.1)
+        def patched_session_init(self, *args, **kwargs):
+            original_session_init(self, *args, **kwargs)
+            created_sessions.add(self)
+
+            original_close = self.close
+
+            def patched_close(*s_args, **s_kwargs):
+                original_close(*s_args, **s_kwargs)
+                closed_sessions.add(self)
+
+            self.close = patched_close
+
+        # Track gRPC Channels
+        created_channels = set()
+        closed_channels = set()
+
+        original_create_channel = google.api_core.grpc_helpers.create_channel
+
+        def patched_create_channel(*args, **kwargs):
+            channel = original_create_channel(*args, **kwargs)
+            created_channels.add(channel)
+
+            original_channel_close = channel.close
+
+            def patched_channel_close(*c_args, **c_kwargs):
+                original_channel_close(*c_args, **c_kwargs)
+                closed_channels.add(channel)
+
+            channel.close = patched_channel_close
+            return channel
+
+        # Apply patches
+        requests.Session.__init__ = patched_session_init
+        google.api_core.grpc_helpers.create_channel = patched_create_channel
 
         try:
-            self.assertLessEqual(conn_count_end, conn_count_start)
-        except AssertionError as e:
-            # Due to flakiness in this test (likely caused by OS cleanup delays or
-            # non-deterministic garbage collection of sockets), we want to capture
-            # the detailed state of connections in future failing runs to help
-            # decrease false positives and identify the root cause.
-            conn_debug = [
-                f"Status: {c.status}, Laddr: {c.laddr}, Raddr: {c.raddr}"
-                for c in current_process.net_connections()
-            ]
-            debug_msg = "\n".join(conn_debug)
-
-            raise AssertionError(
-                f"{e}\n\n"
-                f"--- Socket Leak Debug Info ---\n"
-                f"Start Count: {conn_count_start}\n"
-                f"End Count: {conn_count_end}\n"
-                f"Current Connections:\n{debug_msg}"
+            with helpers.patch_tracked_requests():
+                # Provide no explicit clients, so that the connection will create and own them.
+                connection = dbapi.connect()
+                cursor = connection.cursor()
+
+                cursor.execute(
+                    """
+                    SELECT id, `by`, timestamp
+                    FROM `bigquery-public-data.hacker_news.full`
+                    ORDER BY `id` ASC
+                    LIMIT 100000
+                """
+                )
+                rows = cursor.fetchall()
+                self.assertEqual(len(rows), 100000)
+
+                connection.close()
+
+            # Assertions
+            self.assertEqual(
+                created_sessions,
+                closed_sessions,
+                f"HTTP Sessions leak detected! Leaked: {created_sessions - closed_sessions}",
+            )
+
+            self.assertEqual(
+                created_channels,
+                closed_channels,
+                f"gRPC Channels leak detected! Leaked: {created_channels - closed_channels}",
             )
 
+        finally:
+            # Revert patches
+            requests.Session.__init__ = original_session_init
+            google.api_core.grpc_helpers.create_channel = original_create_channel
+
+            # Clean up any unclosed sessions/channels to avoid leaking in the test runner
+            for s in created_sessions - closed_sessions:
+                try:
+                    s.close()
+                except Exception:
+                    pass
+
+            for c in created_channels - closed_channels:
+                try:
+                    c.close()
+                except Exception:
+                    pass
+
     def _load_table_for_dml(self, rows, dataset_id, table_id):
         from google.cloud._testing import _NamedTemporaryFile
         from google.cloud.bigquery.job import (