From a70f1d8e8d1388d63b97bbdc572f2a6db3b05d49 Mon Sep 17 00:00:00 2001 From: Anthonios Partheniou Date: Mon, 19 Oct 2026 15:21:25 +0000 Subject: [PATCH 1/6] fix(bigquery-magics): wait for graph server startup before returning from init() --- .../bigquery_magics/graph_server.py | 9 +++- .../tests/unit/test_graph_server.py | 46 +++++++++++++++++++ 2 files changed, 54 insertions(+), 1 deletion(-) diff --git a/packages/bigquery-magics/bigquery_magics/graph_server.py b/packages/bigquery-magics/bigquery_magics/graph_server.py index fac2bb281122..11200ba6a440 100644 --- a/packages/bigquery-magics/bigquery_magics/graph_server.py +++ b/packages/bigquery-magics/bigquery_magics/graph_server.py @@ -297,6 +297,7 @@ def __init__(self): self.port = None self.url = None self._server = None + self._server_ready = threading.Event() def build_route(self, endpoint): """ @@ -317,7 +318,11 @@ class ThreadedTCPServer(socketserver.TCPServer): with ThreadedTCPServer(("127.0.0.1", self.port), GraphServerHandler) as httpd: self._server = httpd - self._server.serve_forever() + self._server_ready.set() + try: + self._server.serve_forever() + finally: + self._server_ready.clear() def init(self): """ @@ -327,9 +332,11 @@ def init(self): self.port = portpicker.pick_unused_port() self.url = f"{GraphServer.host}:{self.port}" + self._server_ready.clear() server_thread = threading.Thread(target=self._start_server) server_thread.start() + self._server_ready.wait() return server_thread def stop_server(self): diff --git a/packages/bigquery-magics/tests/unit/test_graph_server.py b/packages/bigquery-magics/tests/unit/test_graph_server.py index 28f4ba58f8ec..719aea97438c 100644 --- a/packages/bigquery-magics/tests/unit/test_graph_server.py +++ b/packages/bigquery-magics/tests/unit/test_graph_server.py @@ -13,6 +13,7 @@ # limitations under the License. import json +import time import unittest from unittest import mock @@ -747,3 +748,48 @@ def test_convert_schema_shared_label(): labels = {label["name"]: label for label in result["labels"]} assert "Person" in labels assert set(labels["Person"]["propertyDeclarationNames"]) == {"id", "name"} + + +def test_stop_server_immediately_after_init(): + """Verifies that GraphServer.init() waits until the underlying TCP server + is initialized before returning, preventing a race condition with stop_server(). + """ + server = graph_server.GraphServer() + server_thread = server.init() + + try: + # 1. Immediately call stop_server() after init() without making any HTTP + # requests. If init() returns before `self._server` is assigned in + # _start_server(), `if self._server:` inside stop_server() evaluates to + # False, making stop_server() a no-op. + server.stop_server() + + # 2. Wait up to 1.0 second for the background server thread to exit. + # When properly synchronized, init() blocks until `self._server` is + # assigned, so stop_server() calls `self._server.shutdown()` and the + # thread exits almost immediately. If the shutdown signal is missed, + # the background thread remains blocked in `serve_forever()`. Using + # `timeout=1.0` prevents `join()` from hanging indefinitely. + server_thread.join(timeout=1.0) + + # 3. Assert that the server thread actually terminated. + assert not server_thread.is_alive(), ( + "GraphServer thread is still alive after stop_server(); " + "stop_server() likely raced with _start_server() before self._server was set." + ) + finally: + # 4. Cleanup in case the assertion above fails: + # Because `server_thread` is a non-daemon thread, Python's interpreter + # shutdown (`threading._shutdown()`) will hang at the end of the pytest + # run if `server_thread` is left running. + if server_thread.is_alive(): + # Poll with a short `time.sleep()` (releasing the GIL so the background + # thread can finish binding the socket and setting `server._server`) + # up to a bounded deadline. + deadline = time.monotonic() + 5.0 + while server._server is None and time.monotonic() < deadline: + time.sleep(0.01) + + # Shut down the server so the background thread exits cleanly. + server.stop_server() + server_thread.join(timeout=2.0) From 4fd9a6d06e66764431603df8ddf42c4fe5ae7cc3 Mon Sep 17 00:00:00 2001 From: Anthonios Partheniou Date: Mon, 19 Oct 2026 11:51:36 -0400 Subject: [PATCH 2/6] Update packages/bigquery-magics/bigquery_magics/graph_server.py Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com> --- packages/bigquery-magics/bigquery_magics/graph_server.py | 8 +++++++- 1 file changed, 7 insertions(+), 1 deletion(-) diff --git a/packages/bigquery-magics/bigquery_magics/graph_server.py b/packages/bigquery-magics/bigquery_magics/graph_server.py index 11200ba6a440..d16b063c3000 100644 --- a/packages/bigquery-magics/bigquery_magics/graph_server.py +++ b/packages/bigquery-magics/bigquery_magics/graph_server.py @@ -336,7 +336,13 @@ def init(self): server_thread = threading.Thread(target=self._start_server) server_thread.start() - self._server_ready.wait() + for _ in range(50): + if self._server_ready.wait(timeout=0.1): + break + if not server_thread.is_alive(): + raise RuntimeError("GraphServer thread terminated unexpectedly during startup.") + else: + raise RuntimeError("Timed out waiting for GraphServer to start.") return server_thread def stop_server(self): From 947ae9bb6b9550bfd8c942b5fc5f4cfdff9510e6 Mon Sep 17 00:00:00 2001 From: Anthonios Partheniou Date: Mon, 19 Oct 2026 16:31:52 +0000 Subject: [PATCH 3/6] fix build --- packages/bigquery-magics/tests/unit/test_graph_server.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/packages/bigquery-magics/tests/unit/test_graph_server.py b/packages/bigquery-magics/tests/unit/test_graph_server.py index 719aea97438c..ecf3c8f21d5f 100644 --- a/packages/bigquery-magics/tests/unit/test_graph_server.py +++ b/packages/bigquery-magics/tests/unit/test_graph_server.py @@ -750,6 +750,9 @@ def test_convert_schema_shared_label(): assert set(labels["Person"]["propertyDeclarationNames"]) == {"id", "name"} +@pytest.mark.skipif( + graph_visualization is None, reason="Requires `spanner-graph-notebook`" +) def test_stop_server_immediately_after_init(): """Verifies that GraphServer.init() waits until the underlying TCP server is initialized before returning, preventing a race condition with stop_server(). From e8af22ad64ae7bfd9f193c2eda058bfc27107a18 Mon Sep 17 00:00:00 2001 From: Anthonios Partheniou Date: Mon, 19 Oct 2026 16:32:30 +0000 Subject: [PATCH 4/6] lint --- packages/bigquery-magics/bigquery_magics/graph_server.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/packages/bigquery-magics/bigquery_magics/graph_server.py b/packages/bigquery-magics/bigquery_magics/graph_server.py index d16b063c3000..49e36e645f88 100644 --- a/packages/bigquery-magics/bigquery_magics/graph_server.py +++ b/packages/bigquery-magics/bigquery_magics/graph_server.py @@ -340,7 +340,9 @@ def init(self): if self._server_ready.wait(timeout=0.1): break if not server_thread.is_alive(): - raise RuntimeError("GraphServer thread terminated unexpectedly during startup.") + raise RuntimeError( + "GraphServer thread terminated unexpectedly during startup." + ) else: raise RuntimeError("Timed out waiting for GraphServer to start.") return server_thread From 85ea9ad481ba4413d8d99f9e244912b7c07187c7 Mon Sep 17 00:00:00 2001 From: Anthonios Partheniou Date: Mon, 19 Oct 2026 17:58:03 +0000 Subject: [PATCH 5/6] fix cover --- .../tests/unit/test_graph_server.py | 86 ++++++++++--------- 1 file changed, 45 insertions(+), 41 deletions(-) diff --git a/packages/bigquery-magics/tests/unit/test_graph_server.py b/packages/bigquery-magics/tests/unit/test_graph_server.py index ecf3c8f21d5f..2f4d10283560 100644 --- a/packages/bigquery-magics/tests/unit/test_graph_server.py +++ b/packages/bigquery-magics/tests/unit/test_graph_server.py @@ -750,49 +750,53 @@ def test_convert_schema_shared_label(): assert set(labels["Person"]["propertyDeclarationNames"]) == {"id", "name"} -@pytest.mark.skipif( - graph_visualization is None, reason="Requires `spanner-graph-notebook`" -) def test_stop_server_immediately_after_init(): """Verifies that GraphServer.init() waits until the underlying TCP server is initialized before returning, preventing a race condition with stop_server(). """ - server = graph_server.GraphServer() - server_thread = server.init() - - try: - # 1. Immediately call stop_server() after init() without making any HTTP - # requests. If init() returns before `self._server` is assigned in - # _start_server(), `if self._server:` inside stop_server() evaluates to - # False, making stop_server() a no-op. - server.stop_server() - - # 2. Wait up to 1.0 second for the background server thread to exit. - # When properly synchronized, init() blocks until `self._server` is - # assigned, so stop_server() calls `self._server.shutdown()` and the - # thread exits almost immediately. If the shutdown signal is missed, - # the background thread remains blocked in `serve_forever()`. Using - # `timeout=1.0` prevents `join()` from hanging indefinitely. - server_thread.join(timeout=1.0) - - # 3. Assert that the server thread actually terminated. - assert not server_thread.is_alive(), ( - "GraphServer thread is still alive after stop_server(); " - "stop_server() likely raced with _start_server() before self._server was set." - ) - finally: - # 4. Cleanup in case the assertion above fails: - # Because `server_thread` is a non-daemon thread, Python's interpreter - # shutdown (`threading._shutdown()`) will hang at the end of the pytest - # run if `server_thread` is left running. - if server_thread.is_alive(): - # Poll with a short `time.sleep()` (releasing the GIL so the background - # thread can finish binding the socket and setting `server._server`) - # up to a bounded deadline. - deadline = time.monotonic() + 5.0 - while server._server is None and time.monotonic() < deadline: - time.sleep(0.01) - - # Shut down the server so the background thread exits cleanly. + # Mock portpicker to return port 0 (OS-assigned ephemeral port) so this test + # runs in all unit test environments, even when optional dependencies like + # `spanner-graph-notebook` and `portpicker` are not installed. + mock_portpicker = mock.MagicMock() + mock_portpicker.pick_unused_port.return_value = 0 + + with mock.patch.dict("sys.modules", {"portpicker": mock_portpicker}): + server = graph_server.GraphServer() + server_thread = server.init() + + try: + # 1. Immediately call stop_server() after init() without making any HTTP + # requests. If init() returns before `self._server` is assigned in + # _start_server(), `if self._server:` inside stop_server() evaluates to + # False, making stop_server() a no-op. server.stop_server() - server_thread.join(timeout=2.0) + + # 2. Wait up to 1.0 second for the background server thread to exit. + # When properly synchronized, init() blocks until `self._server` is + # assigned, so stop_server() calls `self._server.shutdown()` and the + # thread exits almost immediately. If the shutdown signal is missed, + # the background thread remains blocked in `serve_forever()`. Using + # `timeout=1.0` prevents `join()` from hanging indefinitely. + server_thread.join(timeout=1.0) + + # 3. Assert that the server thread actually terminated. + assert not server_thread.is_alive(), ( + "GraphServer thread is still alive after stop_server(); " + "stop_server() likely raced with _start_server() before self._server was set." + ) + finally: + # 4. Cleanup in case the assertion above fails: + # Because `server_thread` is a non-daemon thread, Python's interpreter + # shutdown (`threading._shutdown()`) will hang at the end of the pytest + # run if `server_thread` is left running. + if server_thread.is_alive(): # pragma: NO COVER + # Poll with a short `time.sleep()` (releasing the GIL so the background + # thread can finish binding the socket and setting `server._server`) + # up to a bounded deadline. + deadline = time.monotonic() + 5.0 + while server._server is None and time.monotonic() < deadline: + time.sleep(0.01) + + # Shut down the server so the background thread exits cleanly. + server.stop_server() + server_thread.join(timeout=2.0) From b4d8d4d16865423a815b34936e0dc41315898860 Mon Sep 17 00:00:00 2001 From: Anthonios Partheniou Date: Mon, 19 Oct 2026 19:21:10 +0000 Subject: [PATCH 6/6] fix build --- .../bigquery_magics/graph_server.py | 15 ++------------- 1 file changed, 2 insertions(+), 13 deletions(-) diff --git a/packages/bigquery-magics/bigquery_magics/graph_server.py b/packages/bigquery-magics/bigquery_magics/graph_server.py index 49e36e645f88..74d3e937e173 100644 --- a/packages/bigquery-magics/bigquery_magics/graph_server.py +++ b/packages/bigquery-magics/bigquery_magics/graph_server.py @@ -319,10 +319,7 @@ class ThreadedTCPServer(socketserver.TCPServer): with ThreadedTCPServer(("127.0.0.1", self.port), GraphServerHandler) as httpd: self._server = httpd self._server_ready.set() - try: - self._server.serve_forever() - finally: - self._server_ready.clear() + self._server.serve_forever() def init(self): """ @@ -336,15 +333,7 @@ def init(self): server_thread = threading.Thread(target=self._start_server) server_thread.start() - for _ in range(50): - if self._server_ready.wait(timeout=0.1): - break - if not server_thread.is_alive(): - raise RuntimeError( - "GraphServer thread terminated unexpectedly during startup." - ) - else: - raise RuntimeError("Timed out waiting for GraphServer to start.") + self._server_ready.wait() return server_thread def stop_server(self):