From 67a61c61f52ff647d61dd3c8dd0b024e60a37099 Mon Sep 17 00:00:00 2001 From: Shuowei Li Date: Thu, 30 Jul 2026 20:56:09 +0000 Subject: [PATCH 1/3] fix(bigquery): harden test fixtures and prevent socket leaks --- .../google/cloud/bigquery/dbapi/connection.py | 12 ++-- .../tests/system/test_client.py | 35 ++++++---- .../tests/unit/test_magics.py | 67 ++++++++++++------- 3 files changed, 72 insertions(+), 42 deletions(-) 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..0cc029c49530 100644 --- a/packages/google-cloud-bigquery/tests/system/test_client.py +++ b/packages/google-cloud-bigquery/tests/system/test_client.py @@ -204,12 +204,16 @@ def _still_in_use(bad_request): tag_key = key_values.pop() # Delete tag values first - [ - tag_values_client.delete_tag_value(name=tag_value.name).result() - for tag_value in key_values - ] + for tag_value in key_values: + try: + tag_values_client.delete_tag_value(name=tag_value.name).result() + except NotFound: + pass - tag_keys_client.delete_tag_key(name=tag_key.name).result() + try: + tag_keys_client.delete_tag_key(name=tag_key.name).result() + except NotFound: + pass def test_get_service_account_email(self): client = Config.CLIENT @@ -2184,6 +2188,7 @@ 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_start_addrs = {c.laddr for c in conn_start if c.laddr} conn_count_start = len(conn_start) with helpers.patch_tracked_requests(): @@ -2202,24 +2207,29 @@ def test_dbapi_connection_does_not_leak_sockets(self): rows = cursor.fetchall() self.assertEqual(len(rows), 100000) + cursor.close() connection.close() + + del connection, cursor, rows import gc gc.collect() - for _ in range(30): # Wait up to 3 seconds + for _ in range(60): # Wait up to 6 seconds for background socket cleanup + gc.collect() conn_end = current_process.net_connections() conn_count_end = len(conn_end) - if conn_count_end <= conn_count_start: + new_conns_remaining = [ + c for c in conn_end if c.laddr and c.laddr not in conn_start_addrs + ] + if conn_count_end <= conn_count_start or len(new_conns_remaining) == 0: break time.sleep(0.1) try: - self.assertLessEqual(conn_count_end, conn_count_start) + self.assertTrue( + conn_count_end <= conn_count_start or len(new_conns_remaining) == 0 + ) 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() @@ -2231,6 +2241,7 @@ def test_dbapi_connection_does_not_leak_sockets(self): f"--- Socket Leak Debug Info ---\n" f"Start Count: {conn_count_start}\n" f"End Count: {conn_count_end}\n" + f"New Sockets Remaining: {len(new_conns_remaining)}\n" f"Current Connections:\n{debug_msg}" ) diff --git a/packages/google-cloud-bigquery/tests/unit/test_magics.py b/packages/google-cloud-bigquery/tests/unit/test_magics.py index 03a3a2dbbdba..e2eed9b35793 100644 --- a/packages/google-cloud-bigquery/tests/unit/test_magics.py +++ b/packages/google-cloud-bigquery/tests/unit/test_magics.py @@ -37,13 +37,13 @@ bigquery_storage = pytest.importorskip("google.cloud.bigquery_storage") IPython = pytest.importorskip("IPython") -interactiveshell = pytest.importorskip("IPython.terminal.interactiveshell") +interactiveshell = pytest.importorskip("IPython.core.interactiveshell") tools = pytest.importorskip("IPython.testing.tools") io = pytest.importorskip("IPython.utils.io") pandas = pytest.importorskip("pandas") -@pytest.fixture() +@pytest.fixture(autouse=True) def use_local_magics_context(monkeypatch): if magics is not None: # pragma: NO COVER local_context = magics.Context() @@ -58,8 +58,7 @@ def use_local_magics_context(monkeypatch): @pytest.fixture(scope="session") def ipython(): config = tools.default_config() - config.TerminalInteractiveShell.simple_prompt = True - shell = interactiveshell.TerminalInteractiveShell.instance(config=config) + shell = interactiveshell.InteractiveShell.instance(config=config) return shell @@ -147,6 +146,8 @@ def test_context_with_default_credentials(): """When Application Default Credentials are set, the context credentials will be created the first time it is called """ + magics.context._credentials = None + magics.context._project = None assert magics.context._credentials is None assert magics.context._project is None @@ -164,6 +165,16 @@ def test_context_with_default_credentials(): assert default_mock.call_count == 2 +def test_context_fallback_when_bigquery_magics_none(): + ctx = magics.Context() + credentials_mock = mock.create_autospec( + google.auth.credentials.Credentials, instance=True + ) + with mock.patch("google.auth.default", return_value=(credentials_mock, "proj-123")): + assert ctx.credentials is credentials_mock + assert ctx.project == "proj-123" + + @pytest.mark.usefixtures("ipython_interactive") @pytest.mark.skipif(pandas is None, reason="Requires `pandas`") def test_context_with_default_connection(monkeypatch): @@ -674,9 +685,11 @@ def test_bigquery_magic_with_bqstorage_from_argument( google.cloud.bigquery.job.QueryJob, instance=True ) query_job_mock.to_dataframe.return_value = result - with run_query_patch as run_query_mock, ( - bqstorage_client_patch - ), warnings.catch_warnings(record=True) as warned: + with ( + run_query_patch as run_query_mock, + bqstorage_client_patch, + warnings.catch_warnings(record=True) as warned, + ): run_query_mock.return_value = query_job_mock return_value = ip.run_cell_magic("bigquery", "--use_bqstorage_api", sql) @@ -842,11 +855,12 @@ def test_bigquery_magic_w_max_results_query_job_results_fails(monkeypatch): ) query_job_mock.result.side_effect = [[], OSError] - with pytest.raises( - OSError - ), client_query_patch as client_query_mock, ( - default_patch - ), close_transports_patch as close_transports: + with ( + pytest.raises(OSError), + client_query_patch as client_query_mock, + default_patch, + close_transports_patch as close_transports, + ): client_query_mock.return_value = query_job_mock ip.run_cell_magic("bigquery", "--max_results=5", sql) @@ -1965,9 +1979,10 @@ def test_bigquery_magic_nonexisting_query_variable(monkeypatch): ip.user_ns.pop("custom_query", None) # Make sure the variable does NOT exist. cell_body = "$custom_query" # Referring to a non-existing variable name. - with pytest.raises( - NameError, match=r".*custom_query does not exist.*" - ), run_query_patch as run_query_mock: + with ( + pytest.raises(NameError, match=r".*custom_query does not exist.*"), + run_query_patch as run_query_mock, + ): ip.run_cell_magic("bigquery", "", cell_body) run_query_mock.assert_not_called() @@ -1988,9 +2003,10 @@ def test_bigquery_magic_empty_query_variable_name(monkeypatch): ) cell_body = "$" # Not referring to any variable (name omitted). - with pytest.raises( - NameError, match=r"(?i).*missing query variable name.*" - ), run_query_patch as run_query_mock: + with ( + pytest.raises(NameError, match=r"(?i).*missing query variable name.*"), + run_query_patch as run_query_mock, + ): ip.run_cell_magic("bigquery", "", cell_body) run_query_mock.assert_not_called() @@ -2016,9 +2032,10 @@ def test_bigquery_magic_query_variable_non_string(ipython_ns_cleanup, monkeypatc ip.user_ns["custom_query"] = object() cell_body = "$custom_query" # Referring to a non-string variable. - with pytest.raises( - TypeError, match=r".*must be a string or a bytes-like.*" - ), run_query_patch as run_query_mock: + with ( + pytest.raises(TypeError, match=r".*must be a string or a bytes-like.*"), + run_query_patch as run_query_mock, + ): ip.run_cell_magic("bigquery", "", cell_body) run_query_mock.assert_not_called() @@ -2183,9 +2200,11 @@ def test_bigquery_magic_create_dataset_fails(monkeypatch): autospec=True, ) - with pytest.raises( - OSError - ), create_dataset_if_necessary_patch, close_transports_patch as close_transports: + with ( + pytest.raises(OSError), + create_dataset_if_necessary_patch, + close_transports_patch as close_transports, + ): ip.run_cell_magic( "bigquery", "--destination_table dataset_id.table_id", From d8df9119c5fe535017b3bdd4c3c2276d2074eb24 Mon Sep 17 00:00:00 2001 From: Shuowei Li Date: Thu, 30 Jul 2026 14:06:57 -0700 Subject: [PATCH 2/3] Update packages/google-cloud-bigquery/google/cloud/bigquery/dbapi/connection.py Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com> --- .../google/cloud/bigquery/dbapi/connection.py | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) 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 450e4188beef..40ae6e75d87b 100644 --- a/packages/google-cloud-bigquery/google/cloud/bigquery/dbapi/connection.py +++ b/packages/google-cloud-bigquery/google/cloud/bigquery/dbapi/connection.py @@ -79,8 +79,11 @@ def close(self): self._closed = True for cursor_ in list(self._cursors_created): - if not cursor_._closed: - cursor_.close() + try: + if not cursor_._closed: + cursor_.close() + except Exception: + pass if self._owns_client and self._client is not None: self._client.close() From 623276b4bf428000c28f012704a934894b0a75b6 Mon Sep 17 00:00:00 2001 From: Shuowei Li Date: Thu, 30 Jul 2026 21:17:09 +0000 Subject: [PATCH 3/3] fix(bigquery): remove unexercised try-except in Connection.close --- .../google/cloud/bigquery/dbapi/connection.py | 7 ++----- 1 file changed, 2 insertions(+), 5 deletions(-) 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 40ae6e75d87b..450e4188beef 100644 --- a/packages/google-cloud-bigquery/google/cloud/bigquery/dbapi/connection.py +++ b/packages/google-cloud-bigquery/google/cloud/bigquery/dbapi/connection.py @@ -79,11 +79,8 @@ def close(self): self._closed = True for cursor_ in list(self._cursors_created): - try: - if not cursor_._closed: - cursor_.close() - except Exception: - pass + if not cursor_._closed: + cursor_.close() if self._owns_client and self._client is not None: self._client.close()