diff --git a/bigquery_magics/bigquery.py b/bigquery_magics/bigquery.py index a10de40..28f9e60 100644 --- a/bigquery_magics/bigquery.py +++ b/bigquery_magics/bigquery.py @@ -648,7 +648,10 @@ def _get_graph_name(query_text: str): """ match = re.match(r"\s*GRAPH\s+(\S+)\.(\S+)", query_text, re.IGNORECASE) if match: - return (match.group(1), match.group(2)) + (dataset_id, graph_id) = (match.group(1)), match.group(2) + if "`" in dataset_id or "`" in graph_id: + return None # Backticks in graph name not support for schema view + return (dataset_id, graph_id) return None @@ -668,10 +671,15 @@ def _get_graph_schema( job_config = bigquery.QueryJobConfig( query_parameters=[bigquery.ScalarQueryParameter("graph_id", "STRING", graph_id)] ) - info_schema_results = bq_client.query( - info_schema_query, job_config=job_config - ).to_dataframe() - + job_config.use_legacy_sql = False + try: + info_schema_results = bq_client.query( + info_schema_query, job_config=job_config + ).to_dataframe() + except Exception: + # If the INFORMATION_SCHEMA query fails for some reason, disable only schema + # view, not the entire visualizer. + return None if info_schema_results.shape == (1, 1): return graph_server._convert_schema(info_schema_results.iloc[0, 0]) return None @@ -733,7 +741,7 @@ def _add_graph_widget( "Error: The query result is too large for graph visualization." ) ) - return + return False schema = _get_graph_schema(bq_client, query_text, query_job) @@ -764,6 +772,7 @@ def _add_graph_widget( '"bigquery.graph_visualization.NodeExpansion"', ) IPython.display.display(IPython.core.display.HTML(html_content)) + return True def _is_valid_json(s: str): @@ -870,7 +879,12 @@ def _make_bq_query( result = result.to_dataframe(**dataframe_kwargs) if args.graph and _supports_graph_widget(result): - _add_graph_widget(bq_client, result, query, query_job, args) + if _add_graph_widget(bq_client, result, query, query_job, args): + # Invoke _handle_result() in case the result is saved to a variable, + # but return None to suppress the default table view, which is redundant + # with the table view in the graph visualizer. + _handle_result(result, args) + return None return _handle_result(result, args) diff --git a/tests/unit/bigquery/test_bigquery.py b/tests/unit/bigquery/test_bigquery.py index d101616..efb2c2f 100644 --- a/tests/unit/bigquery/test_bigquery.py +++ b/tests/unit/bigquery/test_bigquery.py @@ -241,6 +241,23 @@ def test__run_query_dry_run_without_errors_is_silent(): assert len(captured.stdout) == 0 +def test__get_graph_name(): + assert magics._get_graph_name("GRAPH foo.bar") == ("foo", "bar") + assert magics._get_graph_name("GRAPH `foo.bar`") is None + assert magics._get_graph_name("GRAPH `foo`.bar") is None + assert magics._get_graph_name("SELECT 1") is None + + +def test__get_graph_schema_exception(): + bq_client = mock.create_autospec(bigquery.Client, instance=True) + bq_client.query.side_effect = Exception("error") + query_text = "GRAPH foo.bar" + query_job = mock.Mock() + query_job.configuration.destination.project = "my-project" + + assert magics._get_graph_schema(bq_client, query_text, query_job) is None + + @pytest.mark.skipif( bigquery_storage is None, reason="Requires `google-cloud-bigquery-storage`" ) @@ -614,9 +631,7 @@ def test_bigquery_graph_json_json_result(monkeypatch): display_mock.assert_called() assert bqstorage_mock.called # BQ storage client was used - assert isinstance(return_value, pandas.DataFrame) - assert len(return_value) == len(result) # verify row count - assert list(return_value) == list(result) # verify column names + assert return_value is None @pytest.mark.skipif( @@ -735,9 +750,7 @@ def test_bigquery_graph_json_result(monkeypatch): assert '\\"location\\": null' in html_content assert bqstorage_mock.called # BQ storage client was used - assert isinstance(return_value, pandas.DataFrame) - assert len(return_value) == len(result) # verify row count - assert list(return_value) == list(result) # verify column names + assert return_value is None @pytest.mark.skipif( @@ -1015,9 +1028,7 @@ def test_bigquery_graph_colab(monkeypatch): assert sys.modules["google.colab"].output.register_callback.called assert bqstorage_mock.called # BQ storage client was used - assert isinstance(return_value, pandas.DataFrame) - assert len(return_value) == len(result) # verify row count - assert list(return_value) == list(result) # verify column names + assert return_value is None @pytest.mark.skipif(