diff --git a/packages/google-cloud-bigquery/google/cloud/bigquery/dbapi/connection.py b/packages/google-cloud-bigquery/google/cloud/bigquery/dbapi/connection.py index 53f32ce387ce..450e4188beef 100644 --- a/packages/google-cloud-bigquery/google/cloud/bigquery/dbapi/connection.py +++ b/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.""" diff --git a/packages/google-cloud-bigquery/tests/system/test_client.py b/packages/google-cloud-bigquery/tests/system/test_client.py index 9ddec48428b1..6f14cc1ed6a0 100644 --- a/packages/google-cloud-bigquery/tests/system/test_client.py +++ b/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 (