From cb872faa5c9700a176b2c0c2a5549d77254ed1a1 Mon Sep 17 00:00:00 2001 From: Eric Feiveson Date: Tue, 27 Jan 2026 15:27:21 -0800 Subject: [PATCH 1/5] Plumb query results to graph server directly, without going through Javascript. WARNING: Works in Jupyter only, not colab. --- bigquery_magics/bigquery.py | 4 +++- bigquery_magics/graph_server.py | 5 ++++- 2 files changed, 7 insertions(+), 2 deletions(-) diff --git a/bigquery_magics/bigquery.py b/bigquery_magics/bigquery.py index c6d1d00..53cd824 100644 --- a/bigquery_magics/bigquery.py +++ b/bigquery_magics/bigquery.py @@ -697,11 +697,13 @@ def _add_graph_widget(query_result): singleton_server_thread = graph_server.graph_server.init() port = graph_server.graph_server.port + graph_server.graph_server.query_result = query_result + # Create html to invoke the graph server html_content = generate_visualization_html( query="placeholder query", port=port, - params=query_result.to_json().replace("\\", "\\\\").replace('"', '\\"'), + params="{}", ) IPython.display.display(IPython.core.display.HTML(html_content)) diff --git a/bigquery_magics/graph_server.py b/bigquery_magics/graph_server.py index 2cf58bc..77780df 100644 --- a/bigquery_magics/graph_server.py +++ b/bigquery_magics/graph_server.py @@ -166,6 +166,7 @@ def __init__(self): self.port = None self.url = None self._server = None + self.query_result = None def build_route(self, endpoint): """ @@ -251,7 +252,9 @@ def handle_post_ping(self): def handle_post_query(self): data = self.parse_post_data() - response = convert_graph_data(query_results=json.loads(data["params"])) + + query_results = json.loads(graph_server.query_result.to_json()) + response = convert_graph_data(query_results=query_results) self.do_data_response(response) def handle_post_node_expansion(self): From 4ad1b71f330bed69f1ed279a5bad96c7db12ddd0 Mon Sep 17 00:00:00 2001 From: Eric Feiveson Date: Wed, 28 Jan 2026 14:29:28 -0800 Subject: [PATCH 2/5] Fix colab path --- bigquery_magics/bigquery.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/bigquery_magics/bigquery.py b/bigquery_magics/bigquery.py index 53cd824..f88061a 100644 --- a/bigquery_magics/bigquery.py +++ b/bigquery_magics/bigquery.py @@ -633,8 +633,9 @@ def _handle_result(result, args): def _colab_query_callback(query: str, params: str): + query_results = json.loads(graph_server.graph_server.query_result.to_json()) return IPython.core.display.JSON( - graph_server.convert_graph_data(query_results=json.loads(params)) + graph_server.convert_graph_data(query_results=query_results) ) From 78e1d852469afc8c93b9a0cf578f344355f8b827 Mon Sep 17 00:00:00 2001 From: Eric Feiveson Date: Wed, 28 Jan 2026 14:40:02 -0800 Subject: [PATCH 3/5] Fix lint --- bigquery_magics/graph_server.py | 2 -- 1 file changed, 2 deletions(-) diff --git a/bigquery_magics/graph_server.py b/bigquery_magics/graph_server.py index 77780df..de95a13 100644 --- a/bigquery_magics/graph_server.py +++ b/bigquery_magics/graph_server.py @@ -251,8 +251,6 @@ def handle_post_ping(self): self.do_data_response({"your_request": data}) def handle_post_query(self): - data = self.parse_post_data() - query_results = json.loads(graph_server.query_result.to_json()) response = convert_graph_data(query_results=query_results) self.do_data_response(response) From bbd0741c9ab2d6dc663060c2f20c69a6c3e3ce50 Mon Sep 17 00:00:00 2001 From: Eric Feiveson Date: Wed, 28 Jan 2026 16:40:09 -0800 Subject: [PATCH 4/5] Fix unit tests --- bigquery_magics/graph_server.py | 1 + tests/unit/bigquery/test_bigquery.py | 37 ++++------------------------ tests/unit/test_graph_server.py | 11 +++------ 3 files changed, 10 insertions(+), 39 deletions(-) diff --git a/bigquery_magics/graph_server.py b/bigquery_magics/graph_server.py index de95a13..869aa51 100644 --- a/bigquery_magics/graph_server.py +++ b/bigquery_magics/graph_server.py @@ -251,6 +251,7 @@ def handle_post_ping(self): self.do_data_response({"your_request": data}) def handle_post_query(self): + self.parse_post_data() query_results = json.loads(graph_server.query_result.to_json()) response = convert_graph_data(query_results=query_results) self.do_data_response(response) diff --git a/tests/unit/bigquery/test_bigquery.py b/tests/unit/bigquery/test_bigquery.py index d5206c1..431c247 100644 --- a/tests/unit/bigquery/test_bigquery.py +++ b/tests/unit/bigquery/test_bigquery.py @@ -680,18 +680,10 @@ def test_bigquery_graph_json_result(monkeypatch): html_content = display_mock.call_args_list[0][0][0].data assert "" in html_content - # Verify that the query results are embedded into the HTML, allowing them to be visualized. - # Due to escaping, it is not possible check for graph_json_rows exactly, so we check for a few - # sentinel strings within the query results, instead. + # Sanity check that query results are not embedded into the HTML. assert ( - "mUZpbkdyYXBoLlBlcnNvbgB4kQI=" in html_content + "mUZpbkdyYXBoLlBlcnNvbgB4kQI=" not in html_content ) # identifier in 1st row of query result - assert ( - "mUZpbkdyYXBoLlBlcnNvbgB4kQY=" in html_content - ) # identifier in 2nd row of query result - assert ( - "mUZpbkdyYXBoLlBlcnNvbgB4kQQ=" in html_content - ) # identifier in 3rd row of query result # Make sure we can run a second graph query, after the graph server is already running. try: @@ -704,18 +696,6 @@ def test_bigquery_graph_json_result(monkeypatch): html_content = display_mock.call_args_list[0][0][0].data assert "" in html_content - # Verify that the query results are embedded into the HTML, allowing them to be visualized. - # Due to escaping, it is not possible check for graph_json_rows exactly, so we check for a few - # sentinel strings within the query results, instead. - assert ( - "mUZpbkdyYXBoLlBlcnNvbgB4kQI=" in html_content - ) # identifier in 1st row of query result - assert ( - "mUZpbkdyYXBoLlBlcnNvbgB4kQY=" in html_content - ) # identifier in 2nd row of query result - assert ( - "mUZpbkdyYXBoLlBlcnNvbgB4kQQ=" in html_content - ) # identifier in 3rd row of query result assert bqstorage_mock.called # BQ storage client was used assert isinstance(return_value, pandas.DataFrame) @@ -791,18 +771,10 @@ def test_bigquery_graph_colab(monkeypatch): html_content = display_mock.call_args_list[0][0][0].data assert "" in html_content - # Verify that the query results are embedded into the HTML, allowing them to be visualized. - # Due to escaping, it is not possible check for graph_json_rows exactly, so we check for a few - # sentinel strings within the query results, instead. + # Verify that the query results are not embedded into the HTML. assert ( - "mUZpbkdyYXBoLlBlcnNvbgB4kQI=" in html_content + "mUZpbkdyYXBoLlBlcnNvbgB4kQI=" not in html_content ) # identifier in 1st row of query result - assert ( - "mUZpbkdyYXBoLlBlcnNvbgB4kQY=" in html_content - ) # identifier in 2nd row of query result - assert ( - "mUZpbkdyYXBoLlBlcnNvbgB4kQQ=" in html_content - ) # identifier in 3rd row of query result # Make sure we actually used colab path, not GraphServer path. assert sys.modules["google.colab"].output.register_callback.called @@ -818,6 +790,7 @@ def test_bigquery_graph_colab(monkeypatch): reason="Requires `spanner-graph-notebook` and `google-cloud-bigquery-storage`", ) def test_colab_query_callback(): + graph_server.graph_server.query_result = pandas.DataFrame([], columns=["result"]) result = bigquery_magics.bigquery._colab_query_callback( "query", json.dumps({"result": {}}) ) diff --git a/tests/unit/test_graph_server.py b/tests/unit/test_graph_server.py index 7c9845b..3102a9f 100644 --- a/tests/unit/test_graph_server.py +++ b/tests/unit/test_graph_server.py @@ -15,6 +15,7 @@ import json import unittest +import pandas as pd import pytest import requests @@ -469,18 +470,14 @@ def test_post_ping(self): @pytest.mark.skipif( graph_visualization is None, reason="Requires `spanner-graph-notebook`" ) - def test_post_query(self): + def test_post_query(self): self.assertTrue(self.server_thread.is_alive()) + graph_server.graph_server.query_result = pd.DataFrame([json.dumps(row_alex_owns_account)], columns=["result"]) route = graph_server.graph_server.build_route( graph_server.GraphServer.endpoints["post_query"] ) - data = { - "result": { - "0": json.dumps(row_alex_owns_account), - } - } - response = requests.post(route, json={"params": json.dumps(data)}) + response = requests.post(route, json={"params": "{}"}) self.assertEqual(response.status_code, 200) response_data = response.json()["response"] From c4f337f12770f8e6ca13a1aeb2b9c3ba70b8f7d0 Mon Sep 17 00:00:00 2001 From: Eric Feiveson Date: Wed, 28 Jan 2026 16:51:40 -0800 Subject: [PATCH 5/5] Fix lint --- tests/unit/test_graph_server.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/tests/unit/test_graph_server.py b/tests/unit/test_graph_server.py index 3102a9f..2048847 100644 --- a/tests/unit/test_graph_server.py +++ b/tests/unit/test_graph_server.py @@ -470,9 +470,11 @@ def test_post_ping(self): @pytest.mark.skipif( graph_visualization is None, reason="Requires `spanner-graph-notebook`" ) - def test_post_query(self): + def test_post_query(self): self.assertTrue(self.server_thread.is_alive()) - graph_server.graph_server.query_result = pd.DataFrame([json.dumps(row_alex_owns_account)], columns=["result"]) + graph_server.graph_server.query_result = pd.DataFrame( + [json.dumps(row_alex_owns_account)], columns=["result"] + ) route = graph_server.graph_server.build_route( graph_server.GraphServer.endpoints["post_query"] )