Skip to content
This repository was archived by the owner on Mar 6, 2026. It is now read-only.

Commit afd5e83

Browse files
ericfe-googleericfe-google
committed
Fix schema view after merge, add unit tests.
1 parent 376a617 commit afd5e83

File tree

2 files changed

+146
-4
lines changed

2 files changed

+146
-4
lines changed

‎bigquery_magics/graph_server.py‎

Lines changed: 6 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -164,7 +164,7 @@ def process_table(table, kind):
164164
return json.dumps(output, indent=2)
165165

166166

167-
def _convert_graph_data(query_results: Dict[str, Dict[str, str]], schema: str = None):
167+
def _convert_graph_data(query_results: Dict[str, Dict[str, str]], schema: Dict = None):
168168
"""
169169
Converts graph data to the form expected by the visualization framework.
170170
@@ -185,7 +185,7 @@ def _convert_graph_data(query_results: Dict[str, Dict[str, str]], schema: str =
185185
for the current row/column. (Note: We only support graph
186186
visualization for columns of type JSON).
187187
schema:
188-
A JSON string containing the schema for the graph.
188+
A dictionary containing the schema for the graph.
189189
"""
190190
# Delay spanner imports until this function is called to avoid making
191191
# spanner-graph-notebook (and its dependencies) hard requirements for bigquery
@@ -245,7 +245,7 @@ def _convert_graph_data(query_results: Dict[str, Dict[str, str]], schema: str =
245245
"nodes": nodes_json,
246246
"edges": edges_json,
247247
# This populates the visualizer's schema view.
248-
"schema": json.loads(schema) if schema is not None else None,
248+
"schema": schema,
249249
# This field is used to populate the visualizer's tabular view.
250250
"query_result": tabular_data,
251251
}
@@ -269,8 +269,10 @@ def convert_graph_params(params: Dict[str, Any]):
269269
query_results = json.loads(
270270
bq_client.list_rows(table_ref).to_dataframe().to_json()
271271
)
272+
schema_json = params.get("schema")
273+
schema = json.loads(schema_json) if schema_json is not None else None
272274
return _convert_graph_data(
273-
query_results=query_results, schema=params.get("schema")
275+
query_results=query_results, schema=schema
274276
)
275277

276278

‎tests/unit/bigquery/test_bigquery.py‎

Lines changed: 140 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1115,6 +1115,146 @@ def test_bigquery_graph_missing_spanner_deps(monkeypatch):
11151115
display_mock.assert_not_called()
11161116

11171117

1118+
@pytest.mark.skipif(
1119+
graph_visualization is None or bigquery_storage is None,
1120+
reason="Requires `spanner-graph-notebook` and `google-cloud-bigquery-storage`",
1121+
)
1122+
def test_add_graph_widget_with_schema(monkeypatch):
1123+
"""Test _add_graph_widget with a valid graph query that retrieves a schema."""
1124+
mock_display = mock.patch("IPython.display.display", autospec=True)
1125+
mock_gen_html = mock.patch(
1126+
"spanner_graphs.graph_visualization.generate_visualization_html",
1127+
return_value="<html>generated_html</html>",
1128+
)
1129+
1130+
bq_client = mock.create_autospec(bigquery.Client, instance=True)
1131+
query_result = pandas.DataFrame([{"id": 1}], columns=["result"])
1132+
query_text = "GRAPH my_dataset.my_graph"
1133+
1134+
query_job = mock.create_autospec(bigquery.job.QueryJob, instance=True)
1135+
query_job.configuration.destination.project = "p"
1136+
query_job.configuration.destination.dataset_id = "d"
1137+
query_job.configuration.destination.table_id = "t"
1138+
1139+
args = mock.Mock()
1140+
args.bigquery_api_endpoint = "e"
1141+
args.project = "p"
1142+
args.location = "l"
1143+
1144+
# Mock INFORMATION_SCHEMA query for schema retrieval
1145+
schema_json = '{"propertyGraphReference": {"propertyGraphId": "my_graph"}}'
1146+
mock_schema_df = pandas.DataFrame(
1147+
[[schema_json]], columns=["PROPERTY_GRAPH_METADATA_JSON"]
1148+
)
1149+
bq_client.query.return_value.to_dataframe.return_value = mock_schema_df
1150+
1151+
with mock_display as display_mock, mock_gen_html as gen_html_mock:
1152+
magics._add_graph_widget(bq_client, query_result, query_text, query_job, args)
1153+
1154+
# Verify schema was retrieved and converted
1155+
assert bq_client.query.called
1156+
call_args = bq_client.query.call_args[0][0]
1157+
assert "INFORMATION_SCHEMA.PROPERTY_GRAPHS" in call_args
1158+
assert 'PROPERTY_GRAPH_NAME = "my_graph"' in call_args
1159+
1160+
# Verify generate_visualization_html was called with the converted schema
1161+
assert gen_html_mock.called
1162+
params_str = gen_html_mock.call_args[1]["params"]
1163+
params = json.loads(params_str.replace('\\"', '"').replace("\\\\", "\\"))
1164+
assert "schema" in params
1165+
schema_obj = json.loads(params["schema"])
1166+
assert schema_obj["name"] == "my_graph"
1167+
1168+
# Verify display was called
1169+
assert display_mock.called
1170+
1171+
1172+
@pytest.mark.skipif(
1173+
graph_visualization is None or bigquery_storage is None,
1174+
reason="Requires `spanner-graph-notebook` and `google-cloud-bigquery-storage`",
1175+
)
1176+
def test_add_graph_widget_no_graph_name(monkeypatch):
1177+
"""Test _add_graph_widget with a query that is not a GRAPH query."""
1178+
mock_display = mock.patch("IPython.display.display", autospec=True)
1179+
mock_gen_html = mock.patch(
1180+
"spanner_graphs.graph_visualization.generate_visualization_html",
1181+
return_value="<html>generated_html</html>",
1182+
)
1183+
1184+
bq_client = mock.create_autospec(bigquery.Client, instance=True)
1185+
query_result = pandas.DataFrame([{"id": 1}], columns=["result"])
1186+
query_text = "SELECT * FROM my_dataset.my_table"
1187+
1188+
query_job = mock.create_autospec(bigquery.job.QueryJob, instance=True)
1189+
query_job.configuration.destination.project = "p"
1190+
query_job.configuration.destination.dataset_id = "d"
1191+
query_job.configuration.destination.table_id = "t"
1192+
1193+
args = mock.Mock()
1194+
args.bigquery_api_endpoint = "e"
1195+
args.project = "p"
1196+
args.location = "l"
1197+
1198+
with mock_display as display_mock, mock_gen_html as gen_html_mock:
1199+
magics._add_graph_widget(bq_client, query_result, query_text, query_job, args)
1200+
1201+
# Verify schema retrieval was NOT attempted since graph name couldn't be parsed
1202+
assert not bq_client.query.called
1203+
1204+
# Verify generate_visualization_html was called without a schema
1205+
assert gen_html_mock.called
1206+
params_str = gen_html_mock.call_args[1]["params"]
1207+
params = json.loads(params_str.replace('\\"', '"').replace("\\\\", "\\"))
1208+
assert "schema" not in params
1209+
1210+
assert display_mock.called
1211+
1212+
1213+
@pytest.mark.skipif(
1214+
graph_visualization is None or bigquery_storage is None,
1215+
reason="Requires `spanner-graph-notebook` and `google-cloud-bigquery-storage`",
1216+
)
1217+
def test_add_graph_widget_schema_not_found(monkeypatch):
1218+
"""Test _add_graph_widget when the graph schema is not found in INFORMATION_SCHEMA."""
1219+
mock_display = mock.patch("IPython.display.display", autospec=True)
1220+
mock_gen_html = mock.patch(
1221+
"spanner_graphs.graph_visualization.generate_visualization_html",
1222+
return_value="<html>generated_html</html>",
1223+
)
1224+
1225+
bq_client = mock.create_autospec(bigquery.Client, instance=True)
1226+
query_result = pandas.DataFrame([{"id": 1}], columns=["result"])
1227+
query_text = "GRAPH my_dataset.my_graph"
1228+
1229+
query_job = mock.create_autospec(bigquery.job.QueryJob, instance=True)
1230+
query_job.configuration.destination.project = "p"
1231+
query_job.configuration.destination.dataset_id = "d"
1232+
query_job.configuration.destination.table_id = "t"
1233+
1234+
args = mock.Mock()
1235+
args.bigquery_api_endpoint = "e"
1236+
args.project = "p"
1237+
args.location = "l"
1238+
1239+
# Mock INFORMATION_SCHEMA query returning empty results
1240+
mock_schema_df = pandas.DataFrame([], columns=["PROPERTY_GRAPH_METADATA_JSON"])
1241+
bq_client.query.return_value.to_dataframe.return_value = mock_schema_df
1242+
1243+
with mock_display as display_mock, mock_gen_html as gen_html_mock:
1244+
magics._add_graph_widget(bq_client, query_result, query_text, query_job, args)
1245+
1246+
# Verify schema retrieval was attempted
1247+
assert bq_client.query.called
1248+
1249+
# Verify generate_visualization_html was called without a schema
1250+
assert gen_html_mock.called
1251+
params_str = gen_html_mock.call_args[1]["params"]
1252+
params = json.loads(params_str.replace('\\"', '"').replace("\\\\", "\\"))
1253+
assert "schema" not in params
1254+
1255+
assert display_mock.called
1256+
1257+
11181258
def test_bigquery_magic_default_connection_user_agent():
11191259
globalipapp.start_ipython()
11201260
ip = globalipapp.get_ipython()

0 commit comments

Comments
 (0)