From b5302e53879ae93d9f7e50bcb99a2bd45cead873 Mon Sep 17 00:00:00 2001
From: Eric Feiveson <ericfe+github@google.com>
Date: Mon, 2 Feb 2026 13:45:27 -0800
Subject: [PATCH 1/7] chore: Make graph visualization more scalable

We now only include query results directly in the html when the query results are less than 100 KB. For larger query results, we store only the reference to the destination table in the HTML, and have the python code re-read the query results from the destination table during the callback.

Also, added a hard limit of 5 MB in the query result size, beyond which, graph visualization is not supported altogether.
---
 bigquery_magics/bigquery.py          | 119 ++++++++-----
 bigquery_magics/core.py              |  73 ++++++++
 bigquery_magics/graph_server.py      |  24 ++-
 tests/unit/bigquery/test_bigquery.py | 244 ++++++++++++++++++++++++++-
 tests/unit/test_graph_server.py      |  51 +++++-
 5 files changed, 465 insertions(+), 46 deletions(-)
 create mode 100644 bigquery_magics/core.py

diff --git a/bigquery_magics/bigquery.py b/bigquery_magics/bigquery.py
index c6d1d00..411e816 100644
--- a/bigquery_magics/bigquery.py
+++ b/bigquery_magics/bigquery.py
@@ -117,7 +117,6 @@
 import IPython  # type: ignore
 from IPython.core import magic_arguments  # type: ignore
 from IPython.core.getipython import get_ipython
-from google.api_core import client_info
 from google.api_core.exceptions import NotFound
 from google.cloud import bigquery
 from google.cloud.bigquery import exceptions
@@ -126,13 +125,12 @@
 from google.cloud.bigquery.job import QueryJobConfig
 import pandas
 
-from bigquery_magics import environment
 from bigquery_magics import line_arg_parser as lap
 import bigquery_magics._versions_helpers
 import bigquery_magics.config
 import bigquery_magics.graph_server as graph_server
+from bigquery_magics import core
 import bigquery_magics.pyformat
-import bigquery_magics.version
 
 try:
     from google.cloud import bigquery_storage  # type: ignore
@@ -147,24 +145,6 @@
 context = bigquery_magics.config.context
 
 
-def _get_user_agent():
-    identities = [
-        f"ipython-{IPython.__version__}",
-        f"bigquery-magics/{bigquery_magics.version.__version__}",
-    ]
-
-    if environment.is_vscode():
-        identities.append("vscode")
-        if environment.is_vscode_google_cloud_code_extension_installed():
-            identities.append(environment.GOOGLE_CLOUD_CODE_EXTENSION_NAME)
-    elif environment.is_jupyter():
-        identities.append("jupyter")
-        if environment.is_jupyter_bigquery_plugin_installed():
-            identities.append(environment.BIGQUERY_JUPYTER_PLUGIN_NAME)
-
-    return " ".join(identities)
-
-
 def _handle_error(error, destination_var=None):
     """Process a query execution error.
 
@@ -565,23 +545,9 @@ def _query_with_pandas(query: str, params: List[Any], args: Any):
 
 
 def _create_clients(args: Any) -> Tuple[bigquery.Client, Any]:
-    bigquery_client_options = copy.deepcopy(context.bigquery_client_options)
-    if args.bigquery_api_endpoint:
-        if isinstance(bigquery_client_options, dict):
-            bigquery_client_options["api_endpoint"] = args.bigquery_api_endpoint
-        else:
-            bigquery_client_options.api_endpoint = args.bigquery_api_endpoint
-
-    bq_client = bigquery.Client(
-        project=args.project or context.project,
-        credentials=context.credentials,
-        default_query_job_config=context.default_query_job_config,
-        client_info=client_info.ClientInfo(user_agent=_get_user_agent()),
-        client_options=bigquery_client_options,
-        location=args.location,
+    bq_client = core.create_bq_client(
+        args.project, args.bigquery_api_endpoint, args.location
     )
-    if context._connection:
-        bq_client._connection = context._connection
 
     # Check and instantiate bq storage client
     if args.use_bqstorage_api is not None:
@@ -633,8 +599,9 @@ def _handle_result(result, args):
 
 
 def _colab_query_callback(query: str, params: str):
+    parsed_params = json.loads(params)
     return IPython.core.display.JSON(
-        graph_server.convert_graph_data(query_results=json.loads(params))
+        graph_server.convert_graph_data(query_results=parsed_params["query_result"])
     )
 
 
@@ -663,7 +630,49 @@ def _colab_node_expansion_callback(request: dict, params_str: str):
 singleton_server_thread: threading.Thread = None
 
 
-def _add_graph_widget(query_result):
+MAX_GRAPH_VISUALIZATION_SIZE = 5000000
+MAX_GRAPH_VISUALIZATION_QUERY_RESULT_SIZE = 100000
+
+
+def _estimate_json_size(df: pandas.DataFrame) -> int:
+    """Approximates the length of df.to_json(orient='records')
+    without materializing the string.
+    """
+    num_rows, num_cols = df.shape
+    if num_rows == 0:
+        return 2  # "[]"
+
+    # 1. Key overhead: "column_name": (repeated for every row)
+    # Includes quotes, colon, and comma separator per field
+    key_overhead = sum(len(f'"{col}":') + 1 for col in df.columns) * num_rows
+
+    # 2. Row structural overhead: { } per row and [ ] for the list
+    # Plus commas between rows (num_rows - 1)
+    structural_overhead = (2 * num_rows) + 2 + (num_rows - 1)
+
+    # 3. Value lengths
+    total_val_len = 0
+    for col in df.columns:
+        series = df[col]
+
+        if pandas.api.types.is_bool_dtype(series):
+            # true (4) or false (5)
+            total_val_len += series.map({True: 4, False: 5}).sum()
+        elif pandas.api.types.is_numeric_dtype(series):
+            # Numeric values (no quotes). Sample for average length to save memory.
+            sample_size = min(len(series), 1000)
+            avg_len = series.sample(sample_size).astype(str).str.len().mean()
+            total_val_len += avg_len * num_rows
+        else:
+            # Strings/Objects: "value" + quotes (2) + rough escaping factor
+            # .str.len() is relatively memory-efficient
+            val_chars = series.astype(str).str.len().sum()
+            total_val_len += val_chars + (2 * num_rows)
+
+    return int(key_overhead + structural_overhead + total_val_len)
+
+
+def _add_graph_widget(query_result: pandas.DataFrame, query_job: Any, args: Any):
     try:
         from spanner_graphs.graph_visualization import generate_visualization_html
     except ImportError as err:
@@ -698,10 +707,36 @@ def _add_graph_widget(query_result):
         port = graph_server.graph_server.port
 
     # Create html to invoke the graph server
+    args_dict = {
+        "bigquery_api_endpoint": args.bigquery_api_endpoint,
+        "project": args.project,
+        "location": args.location,
+    }
+
+    estimated_size = _estimate_json_size(query_result)
+    if estimated_size > MAX_GRAPH_VISUALIZATION_SIZE:
+        IPython.display.display(
+            IPython.core.display.HTML(
+                "<big><b>Error:</b> The query result is too large for graph visualization.</big>"
+            )
+        )
+        return
+
+    table_dict = {
+        "projectId": query_job.configuration.destination.project,
+        "datasetId": query_job.configuration.destination.dataset_id,
+        "tableId": query_job.configuration.destination.table_id,
+    }
+
+    params_dict = {"destination_table": table_dict, "args": args_dict}
+    if estimated_size < MAX_GRAPH_VISUALIZATION_QUERY_RESULT_SIZE:
+        params_dict["query_result"] = json.loads(query_result.to_json())
+
+    params_str = json.dumps(params_dict)
     html_content = generate_visualization_html(
         query="placeholder query",
         port=port,
-        params=query_result.to_json().replace("\\", "\\\\").replace('"', '\\"'),
+        params=params_str.replace("\\", "\\\\").replace('"', '\\"'),
     )
     IPython.display.display(IPython.core.display.HTML(html_content))
 
@@ -810,7 +845,7 @@ def _make_bq_query(
         result = result.to_dataframe(**dataframe_kwargs)
 
     if args.graph and _supports_graph_widget(result):
-        _add_graph_widget(result)
+        _add_graph_widget(result, query_job, args)
     return _handle_result(result, args)
 
 
@@ -904,7 +939,7 @@ def _make_bqstorage_client(client, client_options):
 
     return client._ensure_bqstorage_client(
         client_options=client_options,
-        client_info=gapic_client_info.ClientInfo(user_agent=_get_user_agent()),
+        client_info=gapic_client_info.ClientInfo(user_agent=core._get_user_agent()),
     )
 
 
diff --git a/bigquery_magics/core.py b/bigquery_magics/core.py
new file mode 100644
index 0000000..6ada795
--- /dev/null
+++ b/bigquery_magics/core.py
@@ -0,0 +1,73 @@
+# Copyright 2024 Google LLC
+
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+
+#     https://www.apache.org/licenses/LICENSE-2.0
+
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import copy
+from google.api_core import client_info
+from google.cloud import bigquery
+import IPython  # type: ignore
+from bigquery_magics import environment
+import bigquery_magics.config
+import bigquery_magics.version
+
+context = bigquery_magics.config.context
+
+
+def _get_user_agent():
+    identities = [
+        f"ipython-{IPython.__version__}",
+        f"bigquery-magics/{bigquery_magics.version.__version__}",
+    ]
+
+    if environment.is_vscode():
+        identities.append("vscode")
+        if environment.is_vscode_google_cloud_code_extension_installed():
+            identities.append(environment.GOOGLE_CLOUD_CODE_EXTENSION_NAME)
+    elif environment.is_jupyter():
+        identities.append("jupyter")
+        if environment.is_jupyter_bigquery_plugin_installed():
+            identities.append(environment.BIGQUERY_JUPYTER_PLUGIN_NAME)
+
+    return " ".join(identities)
+
+
+def create_bq_client(project: str, bigquery_api_endpoint: str, location: str):
+    """Creates a BigQuery client.
+
+    Args:
+        project: Project to use for api calls, None to obtain the project from the context.
+        bigquery_api_endpoint: Bigquery client endpoint.
+        location: Cloud region to use for api calls.
+
+    Returns:
+        google.cloud.bigquery.client.Client: The BigQuery client.
+    """
+    bigquery_client_options = copy.deepcopy(context.bigquery_client_options)
+    if bigquery_api_endpoint:
+        if isinstance(bigquery_client_options, dict):
+            bigquery_client_options["api_endpoint"] = bigquery_api_endpoint
+        else:
+            bigquery_client_options.api_endpoint = bigquery_api_endpoint
+
+    bq_client = bigquery.Client(
+        project=project or context.project,
+        credentials=context.credentials,
+        default_query_job_config=context.default_query_job_config,
+        client_info=client_info.ClientInfo(user_agent=_get_user_agent()),
+        client_options=bigquery_client_options,
+        location=location,
+    )
+    if context._connection:
+        bq_client._connection = context._connection
+
+    return bq_client
diff --git a/bigquery_magics/graph_server.py b/bigquery_magics/graph_server.py
index 2cf58bc..904d728 100644
--- a/bigquery_magics/graph_server.py
+++ b/bigquery_magics/graph_server.py
@@ -19,6 +19,10 @@
 import threading
 from typing import Any, Dict, List
 
+from google.cloud import bigquery
+
+from bigquery_magics import core
+
 
 def execute_node_expansion(params, request):
     return {"error": "Node expansion not yet implemented"}
@@ -251,7 +255,25 @@ 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"]))
+        params = json.loads(data["params"])
+
+        query_results = None
+        if "query_result" in params:
+            query_results = params["query_result"]
+        else:
+            bq_client = core.create_bq_client(
+                params["args"]["project"],
+                params["args"]["bigquery_api_endpoint"],
+                params["args"]["location"],
+            )
+
+            table_ref = bigquery.TableReference.from_api_repr(
+                params["destination_table"]
+            )
+            query_results = json.loads(
+                bq_client.list_rows(table_ref).to_dataframe().to_json()
+            )
+        response = convert_graph_data(query_results=query_results)
         self.do_data_response(response)
 
     def handle_post_node_expansion(self):
diff --git a/tests/unit/bigquery/test_bigquery.py b/tests/unit/bigquery/test_bigquery.py
index d5206c1..1ed5539 100644
--- a/tests/unit/bigquery/test_bigquery.py
+++ b/tests/unit/bigquery/test_bigquery.py
@@ -566,6 +566,7 @@ def test_bigquery_graph_json_json_result(monkeypatch):
     # Set up the context with monkeypatch so that it's reset for subsequent
     # tests.
     monkeypatch.setattr(bigquery_magics.context, "_credentials", mock_credentials)
+    monkeypatch.setattr(bigquery_magics.context, "_project", PROJECT_ID)
 
     # Mock out the BigQuery Storage API.
     bqstorage_mock = mock.create_autospec(bigquery_storage.BigQueryReadClient)
@@ -600,6 +601,9 @@ def test_bigquery_graph_json_json_result(monkeypatch):
         google.cloud.bigquery.job.QueryJob, instance=True
     )
     query_job_mock.to_dataframe.return_value = result
+    query_job_mock.configuration.destination.project = PROJECT_ID
+    query_job_mock.configuration.destination.dataset_id = DATASET_ID
+    query_job_mock.configuration.destination.table_id = TABLE_ID
 
     with run_query_patch as run_query_mock, (
         bqstorage_client_patch
@@ -633,6 +637,7 @@ def test_bigquery_graph_json_result(monkeypatch):
     # Set up the context with monkeypatch so that it's reset for subsequent
     # tests.
     monkeypatch.setattr(bigquery_magics.context, "_credentials", mock_credentials)
+    monkeypatch.setattr(bigquery_magics.context, "_project", PROJECT_ID)
 
     # Mock out the BigQuery Storage API.
     bqstorage_mock = mock.create_autospec(bigquery_storage.BigQueryReadClient)
@@ -664,6 +669,9 @@ def test_bigquery_graph_json_result(monkeypatch):
         google.cloud.bigquery.job.QueryJob, instance=True
     )
     query_job_mock.to_dataframe.return_value = result
+    query_job_mock.configuration.destination.project = PROJECT_ID
+    query_job_mock.configuration.destination.dataset_id = DATASET_ID
+    query_job_mock.configuration.destination.table_id = TABLE_ID
 
     with run_query_patch as run_query_mock, (
         bqstorage_client_patch
@@ -693,6 +701,12 @@ def test_bigquery_graph_json_result(monkeypatch):
             "mUZpbkdyYXBoLlBlcnNvbgB4kQQ=" in html_content
         )  # identifier in 3rd row of query result
 
+        # Verify that args are present in the HTML.
+        assert '\\"args\\": {' in html_content
+        assert '\\"bigquery_api_endpoint\\": null' in html_content
+        assert '\\"project\\": null' in html_content
+        assert '\\"location\\": null' in html_content
+
         # Make sure we can run a second graph query, after the graph server is already running.
         try:
             return_value = ip.run_cell_magic("bigquery", "--graph", sql)
@@ -717,12 +731,198 @@ def test_bigquery_graph_json_result(monkeypatch):
             "mUZpbkdyYXBoLlBlcnNvbgB4kQQ=" in html_content
         )  # identifier in 3rd row of query result
 
+        # Verify that args are present in the HTML.
+        assert '\\"args\\": {' in html_content
+        assert '\\"bigquery_api_endpoint\\": null' in html_content
+        assert '\\"project\\": null' in html_content
+        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
 
 
+@pytest.mark.skipif(
+    graph_visualization is None or bigquery_storage is None,
+    reason="Requires `spanner-graph-notebook` and `google-cloud-bigquery-storage`",
+)
+def test_bigquery_graph_size_exceeds_max(monkeypatch):
+    globalipapp.start_ipython()
+    ip = globalipapp.get_ipython()
+    ip.extension_manager.load_extension("bigquery_magics")
+    mock_credentials = mock.create_autospec(
+        google.auth.credentials.Credentials, instance=True
+    )
+    monkeypatch.setattr(bigquery_magics.context, "_credentials", mock_credentials)
+    monkeypatch.setattr(bigquery_magics.context, "_project", PROJECT_ID)
+
+    # Set threshold to a very small value to trigger the error.
+    monkeypatch.setattr(magics, "MAX_GRAPH_VISUALIZATION_SIZE", 5)
+
+    bqstorage_mock = mock.create_autospec(bigquery_storage.BigQueryReadClient)
+    bqstorage_instance_mock = mock.create_autospec(
+        bigquery_storage.BigQueryReadClient, instance=True
+    )
+    bqstorage_instance_mock._transport = mock.Mock()
+    bqstorage_mock.return_value = bqstorage_instance_mock
+    bqstorage_client_patch = mock.patch(
+        "google.cloud.bigquery_storage.BigQueryReadClient", bqstorage_mock
+    )
+
+    sql = "SELECT graph_json FROM t"
+    result = pandas.DataFrame(['{"id": 1}'], columns=["graph_json"])
+    run_query_patch = mock.patch("bigquery_magics.bigquery._run_query", autospec=True)
+    display_patch = mock.patch("IPython.display.display", autospec=True)
+    query_job_mock = mock.create_autospec(
+        google.cloud.bigquery.job.QueryJob, instance=True
+    )
+    query_job_mock.to_dataframe.return_value = result
+    query_job_mock.configuration.destination.project = PROJECT_ID
+    query_job_mock.configuration.destination.dataset_id = DATASET_ID
+    query_job_mock.configuration.destination.table_id = TABLE_ID
+
+    with run_query_patch as run_query_mock, (
+        bqstorage_client_patch
+    ), display_patch as display_mock:
+        run_query_mock.return_value = query_job_mock
+
+        ip.run_cell_magic("bigquery", "--graph", sql)
+
+        # Should display error message
+        assert display_mock.called
+        html_content = display_mock.call_args[0][0].data
+        assert (
+            "Error:</b> The query result is too large for graph visualization."
+            in html_content
+        )
+
+
+@pytest.mark.skipif(
+    graph_visualization is None or bigquery_storage is None,
+    reason="Requires `spanner-graph-notebook` and `google-cloud-bigquery-storage`",
+)
+def test_bigquery_graph_size_exceeds_query_result_max(monkeypatch):
+    globalipapp.start_ipython()
+    ip = globalipapp.get_ipython()
+    ip.extension_manager.load_extension("bigquery_magics")
+    mock_credentials = mock.create_autospec(
+        google.auth.credentials.Credentials, instance=True
+    )
+    monkeypatch.setattr(bigquery_magics.context, "_credentials", mock_credentials)
+    monkeypatch.setattr(bigquery_magics.context, "_project", PROJECT_ID)
+
+    # Set threshold to a very small value, but larger than the other one.
+    # We want estimated_size > MAX_GRAPH_VISUALIZATION_QUERY_RESULT_SIZE
+    # and estimated_size <= MAX_GRAPH_VISUALIZATION_SIZE
+    monkeypatch.setattr(magics, "MAX_GRAPH_VISUALIZATION_SIZE", 1000000)
+    monkeypatch.setattr(magics, "MAX_GRAPH_VISUALIZATION_QUERY_RESULT_SIZE", 5)
+
+    bqstorage_mock = mock.create_autospec(bigquery_storage.BigQueryReadClient)
+    bqstorage_instance_mock = mock.create_autospec(
+        bigquery_storage.BigQueryReadClient, instance=True
+    )
+    bqstorage_instance_mock._transport = mock.Mock()
+    bqstorage_mock.return_value = bqstorage_instance_mock
+    bqstorage_client_patch = mock.patch(
+        "google.cloud.bigquery_storage.BigQueryReadClient", bqstorage_mock
+    )
+
+    sql = "SELECT graph_json FROM t"
+    result = pandas.DataFrame(['{"id": 1977323800}'], columns=["graph_json"])
+    run_query_patch = mock.patch("bigquery_magics.bigquery._run_query", autospec=True)
+    display_patch = mock.patch("IPython.display.display", autospec=True)
+    query_job_mock = mock.create_autospec(
+        google.cloud.bigquery.job.QueryJob, instance=True
+    )
+    query_job_mock.to_dataframe.return_value = result
+    query_job_mock.configuration.destination.project = PROJECT_ID
+    query_job_mock.configuration.destination.dataset_id = DATASET_ID
+    query_job_mock.configuration.destination.table_id = TABLE_ID
+
+    with run_query_patch as run_query_mock, (
+        bqstorage_client_patch
+    ), display_patch as display_mock:
+        run_query_mock.return_value = query_job_mock
+
+        ip.run_cell_magic("bigquery", "--graph", sql)
+
+        # Should display visualization but without query_result embedded.
+        assert display_mock.called
+        html_content = display_mock.call_args[0][0].data
+        assert "<script>" in html_content
+        assert "1977323800" not in html_content
+
+
+@pytest.mark.skipif(
+    graph_visualization is None or bigquery_storage is None,
+    reason="Requires `spanner-graph-notebook` and `google-cloud-bigquery-storage`",
+)
+def test_bigquery_graph_with_args_serialization(monkeypatch):
+    globalipapp.start_ipython()
+    ip = globalipapp.get_ipython()
+    ip.extension_manager.load_extension("bigquery_magics")
+    mock_credentials = mock.create_autospec(
+        google.auth.credentials.Credentials, instance=True
+    )
+
+    # Set up the context with monkeypatch so that it's reset for subsequent
+    # tests.
+    monkeypatch.setattr(bigquery_magics.context, "_credentials", mock_credentials)
+    monkeypatch.setattr(bigquery_magics.context, "_project", PROJECT_ID)
+
+    # Mock out the BigQuery Storage API.
+    bqstorage_mock = mock.create_autospec(bigquery_storage.BigQueryReadClient)
+    bqstorage_instance_mock = mock.create_autospec(
+        bigquery_storage.BigQueryReadClient, instance=True
+    )
+    bqstorage_instance_mock._transport = mock.Mock()
+    bqstorage_mock.return_value = bqstorage_instance_mock
+    bqstorage_client_patch = mock.patch(
+        "google.cloud.bigquery_storage.BigQueryReadClient", bqstorage_mock
+    )
+    display_patch = mock.patch("IPython.display.display", autospec=True)
+
+    sql = "SELECT graph_json FROM t"
+    graph_json_rows = [
+        """
+        [{"identifier":"mUZpbkdyYXBoLlBlcnNvbgB4kQI=","kind":"node","labels":["Person"],"properties":{"id":1}}]
+        """
+    ]
+    result = pandas.DataFrame(graph_json_rows, columns=["graph_json"])
+    run_query_patch = mock.patch("bigquery_magics.bigquery._run_query", autospec=True)
+    query_job_mock = mock.create_autospec(
+        google.cloud.bigquery.job.QueryJob, instance=True
+    )
+    query_job_mock.to_dataframe.return_value = result
+    query_job_mock.configuration.destination.project = PROJECT_ID
+    query_job_mock.configuration.destination.dataset_id = DATASET_ID
+    query_job_mock.configuration.destination.table_id = TABLE_ID
+
+    with run_query_patch as run_query_mock, (
+        bqstorage_client_patch
+    ), display_patch as display_mock:
+        run_query_mock.return_value = query_job_mock
+
+        endpoint = "https://example.com"
+        project = "my-project"
+        location = "us-west1"
+        magic_args = (
+            f"--graph --bigquery_api_endpoint={endpoint} "
+            f"--project={project} --location={location}"
+        )
+        ip.run_cell_magic("bigquery", magic_args, sql)
+
+        assert display_mock.called
+        html_content = display_mock.call_args_list[0][0][0].data
+
+        # Verify that args are present in the HTML with their assigned values.
+        assert '\\"args\\": {' in html_content
+        assert f'\\"bigquery_api_endpoint\\": \\"{endpoint}\\"' in html_content
+        assert f'\\"project\\": \\"{project}\\"' in html_content
+        assert f'\\"location\\": \\"{location}\\"' in html_content
+
+
 @pytest.mark.skipif(
     graph_visualization is None or bigquery_storage is None,
     reason="Requires `spanner-graph-notebook` and `google-cloud-bigquery-storage`",
@@ -742,6 +942,7 @@ def test_bigquery_graph_colab(monkeypatch):
     # Set up the context with monkeypatch so that it's reset for subsequent
     # tests.
     monkeypatch.setattr(bigquery_magics.context, "_credentials", mock_credentials)
+    monkeypatch.setattr(bigquery_magics.context, "_project", PROJECT_ID)
 
     # Mock out the BigQuery Storage API.
     bqstorage_mock = mock.create_autospec(bigquery_storage.BigQueryReadClient)
@@ -773,6 +974,9 @@ def test_bigquery_graph_colab(monkeypatch):
         google.cloud.bigquery.job.QueryJob, instance=True
     )
     query_job_mock.to_dataframe.return_value = result
+    query_job_mock.configuration.destination.project = PROJECT_ID
+    query_job_mock.configuration.destination.dataset_id = DATASET_ID
+    query_job_mock.configuration.destination.table_id = "test_destination_table"
 
     with run_query_patch as run_query_mock, (
         bqstorage_client_patch
@@ -804,6 +1008,12 @@ def test_bigquery_graph_colab(monkeypatch):
             "mUZpbkdyYXBoLlBlcnNvbgB4kQQ=" in html_content
         )  # identifier in 3rd row of query result
 
+        # Verify that args are present in the HTML.
+        assert '\\"args\\": {' in html_content
+        assert '\\"bigquery_api_endpoint\\": null' in html_content
+        assert '\\"project\\": null' in html_content
+        assert '\\"location\\": null' in html_content
+
         # Make sure we actually used colab path, not GraphServer path.
         assert sys.modules["google.colab"].output.register_callback.called
 
@@ -819,7 +1029,7 @@ def test_bigquery_graph_colab(monkeypatch):
 )
 def test_colab_query_callback():
     result = bigquery_magics.bigquery._colab_query_callback(
-        "query", json.dumps({"result": {}})
+        "query", json.dumps({"query_result": {"result": {}}})
     )
     assert result.data == {
         "response": {
@@ -2839,3 +3049,35 @@ def test_test_bigquery_magic__extension_loaded__is_registered_set_to_true():
     ip.extension_manager.load_extension("bigquery_magics")
 
     assert bigquery_magics.is_registered is True
+
+
+def test__estimate_json_size_empty():
+    df = pandas.DataFrame()
+    assert magics._estimate_json_size(df) == 2
+
+
+def test__estimate_json_size_bool():
+    df = pandas.DataFrame({"a": [True, False]}, dtype="bool")
+    # key_overhead: (len('"a":') + 1) * 2 = 5 * 2 = 10
+    # structural_overhead: (2 * 2) + 2 + (2 - 1) = 7
+    # total_val_len: 4 (true) + 5 (false) = 9
+    # total = 10 + 7 + 9 = 26
+    assert magics._estimate_json_size(df) == 26
+
+
+def test__estimate_json_size_int():
+    df = pandas.DataFrame({"a": [3, 4]})
+    # key_overhead: (len('"a":') + 1) * 2 = 5 * 2 = 10
+    # structural_overhead: (2 * 2) + 2 + (2 - 1) = 7
+    # total_val_len: 4 (true) + 5 (false) = 9
+    # total = 10 + 7 + 9 = 26
+    assert magics._estimate_json_size(df) == 19
+
+
+def test__estimate_json_size_string():
+    df = pandas.DataFrame({"a": ["x", "yz"]})
+    # key_overhead: (len('"a":') + 1) * 2 = 5 * 2 = 10
+    # structural_overhead: (2 * 2) + 2 + (2 - 1) = 7
+    # total_val_len: (1 + 2) + (2 * 2) = 7
+    # total = 10 + 7 + 7 = 24
+    assert magics._estimate_json_size(df) == 24
diff --git a/tests/unit/test_graph_server.py b/tests/unit/test_graph_server.py
index 7c9845b..093361f 100644
--- a/tests/unit/test_graph_server.py
+++ b/tests/unit/test_graph_server.py
@@ -14,7 +14,9 @@
 
 import json
 import unittest
+from unittest import mock
 
+import pandas as pd
 import pytest
 import requests
 
@@ -476,8 +478,10 @@ def test_post_query(self):
         )
 
         data = {
-            "result": {
-                "0": json.dumps(row_alex_owns_account),
+            "query_result": {
+                "result": {
+                    "0": json.dumps(row_alex_owns_account),
+                }
             }
         }
         response = requests.post(route, json={"params": json.dumps(data)})
@@ -494,6 +498,49 @@ def test_post_query(self):
         )
         self.assertIsNone(response_data["schema"])
 
+    @pytest.mark.skipif(
+        graph_visualization is None, reason="Requires `spanner-graph-notebook`"
+    )
+    def test_post_query_from_table(self):
+        self.assertTrue(self.server_thread.is_alive())
+        route = graph_server.graph_server.build_route(
+            graph_server.GraphServer.endpoints["post_query"]
+        )
+
+        params = {
+            "destination_table": {"projectId": "p", "datasetId": "d", "tableId": "t"},
+            "args": {
+                "project": "p",
+                "bigquery_api_endpoint": "e",
+                "location": "l",
+            },
+        }
+
+        with mock.patch("bigquery_magics.core.create_bq_client") as mock_create:
+            mock_client = mock.Mock()
+            mock_create.return_value = mock_client
+            mock_df = pd.DataFrame(
+                [json.dumps(row_alex_owns_account)], columns=["result"]
+            )
+            mock_client.list_rows.return_value.to_dataframe.return_value = mock_df
+
+            response = requests.post(route, json={"params": json.dumps(params)})
+            self.assertEqual(response.status_code, 200)
+
+            mock_create.assert_called_once_with("p", "e", "l")
+            mock_client.list_rows.assert_called_once()
+
+        response_data = response.json()["response"]
+        self.assertEqual(len(response_data["nodes"]), 2)
+        self.assertEqual(len(response_data["edges"]), 1)
+
+        _validate_nodes_and_edges(response.json())
+
+        self.assertEqual(
+            response_data["query_result"], {"result": [row_alex_owns_account_converted]}
+        )
+        self.assertIsNone(response_data["schema"])
+
     @pytest.mark.skipif(
         graph_visualization is None, reason="Requires `spanner-graph-notebook`"
     )

From c81b1aabe1adbb7bd0404e3207a557d7226b3979 Mon Sep 17 00:00:00 2001
From: Eric Feiveson <ericfe+github@google.com>
Date: Tue, 3 Feb 2026 14:28:43 -0800
Subject: [PATCH 2/7] Fix colab path

---
 bigquery_magics/bigquery.py     |  3 +--
 bigquery_magics/graph_server.py | 38 +++++++++++++++++----------------
 tests/unit/test_graph_server.py | 24 ++++++++++-----------
 3 files changed, 33 insertions(+), 32 deletions(-)

diff --git a/bigquery_magics/bigquery.py b/bigquery_magics/bigquery.py
index 411e816..cdc9bf9 100644
--- a/bigquery_magics/bigquery.py
+++ b/bigquery_magics/bigquery.py
@@ -599,9 +599,8 @@ def _handle_result(result, args):
 
 
 def _colab_query_callback(query: str, params: str):
-    parsed_params = json.loads(params)
     return IPython.core.display.JSON(
-        graph_server.convert_graph_data(query_results=parsed_params["query_result"])
+        graph_server.convert_graph_params(json.loads(params))
     )
 
 
diff --git a/bigquery_magics/graph_server.py b/bigquery_magics/graph_server.py
index 904d728..27c47d3 100644
--- a/bigquery_magics/graph_server.py
+++ b/bigquery_magics/graph_server.py
@@ -58,7 +58,7 @@ def _stringify_properties(d: Any) -> Any:
         return _stringify_value(d)
 
 
-def convert_graph_data(query_results: Dict[str, Dict[str, str]]):
+def _convert_graph_data(query_results: Dict[str, Dict[str, str]]):
     """
     Converts graph data to the form expected by the visualization framework.
 
@@ -147,6 +147,24 @@ def convert_graph_data(query_results: Dict[str, Dict[str, str]]):
         return {"error": getattr(e, "message", str(e))}
 
 
+def convert_graph_params(params: Dict[str, Any]):
+    query_results = None
+    if "query_result" in params:
+        query_results = params["query_result"]
+    else:
+        bq_client = core.create_bq_client(
+            params["args"]["project"],
+            params["args"]["bigquery_api_endpoint"],
+            params["args"]["location"],
+        )
+
+        table_ref = bigquery.TableReference.from_api_repr(params["destination_table"])
+        query_results = json.loads(
+            bq_client.list_rows(table_ref).to_dataframe().to_json()
+        )
+    return _convert_graph_data(query_results=query_results)
+
+
 class GraphServer:
     """
     Http server invoked by Javascript to obtain the query results for visualization.
@@ -257,23 +275,7 @@ def handle_post_query(self):
         data = self.parse_post_data()
         params = json.loads(data["params"])
 
-        query_results = None
-        if "query_result" in params:
-            query_results = params["query_result"]
-        else:
-            bq_client = core.create_bq_client(
-                params["args"]["project"],
-                params["args"]["bigquery_api_endpoint"],
-                params["args"]["location"],
-            )
-
-            table_ref = bigquery.TableReference.from_api_repr(
-                params["destination_table"]
-            )
-            query_results = json.loads(
-                bq_client.list_rows(table_ref).to_dataframe().to_json()
-            )
-        response = convert_graph_data(query_results=query_results)
+        response = convert_graph_params(params)
         self.do_data_response(response)
 
     def handle_post_node_expansion(self):
diff --git a/tests/unit/test_graph_server.py b/tests/unit/test_graph_server.py
index 093361f..773b2eb 100644
--- a/tests/unit/test_graph_server.py
+++ b/tests/unit/test_graph_server.py
@@ -211,7 +211,7 @@ def _validate_nodes_and_edges(result):
     graph_visualization is None, reason="Requires `spanner-graph-notebook`"
 )
 def test_convert_one_column_no_rows():
-    result = graph_server.convert_graph_data({"result": {}})
+    result = graph_server._convert_graph_data({"result": {}})
     assert result == {
         "response": {
             "edges": [],
@@ -226,7 +226,7 @@ def test_convert_one_column_no_rows():
     graph_visualization is None, reason="Requires `spanner-graph-notebook`"
 )
 def test_convert_one_column_one_row():
-    result = graph_server.convert_graph_data(
+    result = graph_server._convert_graph_data(
         {
             "result": {
                 "0": json.dumps(row_alex_owns_account),
@@ -249,7 +249,7 @@ def test_convert_one_column_one_row():
     graph_visualization is None, reason="Requires `spanner-graph-notebook`"
 )
 def test_convert_one_column_two_rows_null_json():
-    result = graph_server.convert_graph_data(
+    result = graph_server._convert_graph_data(
         {
             "result": {
                 "0": None,
@@ -276,7 +276,7 @@ def test_convert_one_column_two_rows_null_json():
     graph_visualization is None, reason="Requires `spanner-graph-notebook`"
 )
 def test_convert_one_column_two_rows():
-    result = graph_server.convert_graph_data(
+    result = graph_server._convert_graph_data(
         {
             "result": {
                 "0": json.dumps(row_alex_owns_account_converted),
@@ -300,7 +300,7 @@ def test_convert_one_column_two_rows():
     graph_visualization is None, reason="Requires `spanner-graph-notebook`"
 )
 def test_convert_one_row_two_columns():
-    result = graph_server.convert_graph_data(
+    result = graph_server._convert_graph_data(
         {
             "col1": {
                 "0": json.dumps(row_alex_owns_account_converted),
@@ -329,7 +329,7 @@ def test_convert_one_row_two_columns():
 def test_convert_nongraph_json():
     # If we have valid json that doesn't represent a graph, we don't expect to get nodes and edges,
     # but we should at least have row data, allowing the tabular view to work.
-    result = graph_server.convert_graph_data(
+    result = graph_server._convert_graph_data(
         {
             "result": {
                 "0": json.dumps({"foo": 1, "bar": 2}),
@@ -348,7 +348,7 @@ def test_convert_nongraph_json():
     graph_visualization is None, reason="Requires `spanner-graph-notebook`"
 )
 def test_convert_outer_key_not_string():
-    result = graph_server.convert_graph_data(
+    result = graph_server._convert_graph_data(
         {
             0: {
                 "0": json.dumps({"foo": 1, "bar": 2}),
@@ -362,7 +362,7 @@ def test_convert_outer_key_not_string():
     graph_visualization is None, reason="Requires `spanner-graph-notebook`"
 )
 def test_convert_outer_value_not_dict():
-    result = graph_server.convert_graph_data({"result": 0})
+    result = graph_server._convert_graph_data({"result": 0})
     assert result == {"error": "Expected outer value to be dict, got <class 'int'>"}
 
 
@@ -370,7 +370,7 @@ def test_convert_outer_value_not_dict():
     graph_visualization is None, reason="Requires `spanner-graph-notebook`"
 )
 def test_convert_inner_value_not_string():
-    result = graph_server.convert_graph_data(
+    result = graph_server._convert_graph_data(
         {
             "col1": {
                 "0": json.dumps(row_alex_owns_account),
@@ -398,7 +398,7 @@ def test_convert_inner_value_not_string():
     graph_visualization is None, reason="Requires `spanner-graph-notebook`"
 )
 def test_convert_empty_dict():
-    result = graph_server.convert_graph_data({})
+    result = graph_server._convert_graph_data({})
     assert result == {
         "response": {
             "nodes": [],
@@ -413,7 +413,7 @@ def test_convert_empty_dict():
     graph_visualization is None, reason="Requires `spanner-graph-notebook`"
 )
 def test_convert_wrong_row_index():
-    result0 = graph_server.convert_graph_data(
+    result0 = graph_server._convert_graph_data(
         {
             "result": {
                 "0": json.dumps(row_alex_owns_account),
@@ -422,7 +422,7 @@ def test_convert_wrong_row_index():
     )
 
     # Changing the index should not impact the result.
-    result1 = graph_server.convert_graph_data(
+    result1 = graph_server._convert_graph_data(
         {
             "result": {
                 "1": json.dumps(row_alex_owns_account),

From 0aaa7ad3e553eace633c0537cdf34631b9f6a8d2 Mon Sep 17 00:00:00 2001
From: Eric Feiveson <ericfe+github@google.com>
Date: Thu, 5 Feb 2026 13:41:27 -0800
Subject: [PATCH 3/7] wip

---
 bigquery_magics/bigquery.py | 9 ++++++---
 1 file changed, 6 insertions(+), 3 deletions(-)

diff --git a/bigquery_magics/bigquery.py b/bigquery_magics/bigquery.py
index cdc9bf9..ec14e21 100644
--- a/bigquery_magics/bigquery.py
+++ b/bigquery_magics/bigquery.py
@@ -629,8 +629,8 @@ def _colab_node_expansion_callback(request: dict, params_str: str):
 singleton_server_thread: threading.Thread = None
 
 
-MAX_GRAPH_VISUALIZATION_SIZE = 5000000
-MAX_GRAPH_VISUALIZATION_QUERY_RESULT_SIZE = 100000
+MAX_GRAPH_VISUALIZATION_SIZE = 100000000#2000000
+MAX_GRAPH_VISUALIZATION_QUERY_RESULT_SIZE = 100000000 #100000
 
 
 def _estimate_json_size(df: pandas.DataFrame) -> int:
@@ -712,7 +712,10 @@ def _add_graph_widget(query_result: pandas.DataFrame, query_job: Any, args: Any)
         "location": args.location,
     }
 
-    estimated_size = _estimate_json_size(query_result)
+    estimated_size = query_result.memory_usage(index=True, deep=True).sum()
+    #old_estimated_size = _estimate_json_size(query_result)
+    #actual_size = len(query_result.to_json())
+    #raise ValueError(f'Estimated size changing from {old_estimated_size} to {estimated_size} (actual size: {actual_size})')
     if estimated_size > MAX_GRAPH_VISUALIZATION_SIZE:
         IPython.display.display(
             IPython.core.display.HTML(

From 76f4fe4d598027da7dfffbbb5f5f6244db24a573 Mon Sep 17 00:00:00 2001
From: Eric Feiveson <ericfe+github@google.com>
Date: Thu, 5 Feb 2026 15:03:26 -0800
Subject: [PATCH 4/7] Respond to code review comments

---
 bigquery_magics/bigquery.py     | 49 ++++-----------------------------
 bigquery_magics/core.py         |  4 ++-
 bigquery_magics/graph_server.py |  6 ++--
 3 files changed, 11 insertions(+), 48 deletions(-)

diff --git a/bigquery_magics/bigquery.py b/bigquery_magics/bigquery.py
index ec14e21..98b29d0 100644
--- a/bigquery_magics/bigquery.py
+++ b/bigquery_magics/bigquery.py
@@ -546,7 +546,9 @@ def _query_with_pandas(query: str, params: List[Any], args: Any):
 
 def _create_clients(args: Any) -> Tuple[bigquery.Client, Any]:
     bq_client = core.create_bq_client(
-        args.project, args.bigquery_api_endpoint, args.location
+        project=args.project,
+        bigquery_api_endpoint=args.bigquery_api_endpoint,
+        location=args.location,
     )
 
     # Check and instantiate bq storage client
@@ -629,46 +631,8 @@ def _colab_node_expansion_callback(request: dict, params_str: str):
 singleton_server_thread: threading.Thread = None
 
 
-MAX_GRAPH_VISUALIZATION_SIZE = 100000000#2000000
-MAX_GRAPH_VISUALIZATION_QUERY_RESULT_SIZE = 100000000 #100000
-
-
-def _estimate_json_size(df: pandas.DataFrame) -> int:
-    """Approximates the length of df.to_json(orient='records')
-    without materializing the string.
-    """
-    num_rows, num_cols = df.shape
-    if num_rows == 0:
-        return 2  # "[]"
-
-    # 1. Key overhead: "column_name": (repeated for every row)
-    # Includes quotes, colon, and comma separator per field
-    key_overhead = sum(len(f'"{col}":') + 1 for col in df.columns) * num_rows
-
-    # 2. Row structural overhead: { } per row and [ ] for the list
-    # Plus commas between rows (num_rows - 1)
-    structural_overhead = (2 * num_rows) + 2 + (num_rows - 1)
-
-    # 3. Value lengths
-    total_val_len = 0
-    for col in df.columns:
-        series = df[col]
-
-        if pandas.api.types.is_bool_dtype(series):
-            # true (4) or false (5)
-            total_val_len += series.map({True: 4, False: 5}).sum()
-        elif pandas.api.types.is_numeric_dtype(series):
-            # Numeric values (no quotes). Sample for average length to save memory.
-            sample_size = min(len(series), 1000)
-            avg_len = series.sample(sample_size).astype(str).str.len().mean()
-            total_val_len += avg_len * num_rows
-        else:
-            # Strings/Objects: "value" + quotes (2) + rough escaping factor
-            # .str.len() is relatively memory-efficient
-            val_chars = series.astype(str).str.len().sum()
-            total_val_len += val_chars + (2 * num_rows)
-
-    return int(key_overhead + structural_overhead + total_val_len)
+MAX_GRAPH_VISUALIZATION_SIZE = 2_000_000
+MAX_GRAPH_VISUALIZATION_QUERY_RESULT_SIZE = 100_000
 
 
 def _add_graph_widget(query_result: pandas.DataFrame, query_job: Any, args: Any):
@@ -713,9 +677,6 @@ def _add_graph_widget(query_result: pandas.DataFrame, query_job: Any, args: Any)
     }
 
     estimated_size = query_result.memory_usage(index=True, deep=True).sum()
-    #old_estimated_size = _estimate_json_size(query_result)
-    #actual_size = len(query_result.to_json())
-    #raise ValueError(f'Estimated size changing from {old_estimated_size} to {estimated_size} (actual size: {actual_size})')
     if estimated_size > MAX_GRAPH_VISUALIZATION_SIZE:
         IPython.display.display(
             IPython.core.display.HTML(
diff --git a/bigquery_magics/core.py b/bigquery_magics/core.py
index 6ada795..7bb38d6 100644
--- a/bigquery_magics/core.py
+++ b/bigquery_magics/core.py
@@ -13,9 +13,11 @@
 # limitations under the License.
 
 import copy
+
 from google.api_core import client_info
 from google.cloud import bigquery
 import IPython  # type: ignore
+
 from bigquery_magics import environment
 import bigquery_magics.config
 import bigquery_magics.version
@@ -41,7 +43,7 @@ def _get_user_agent():
     return " ".join(identities)
 
 
-def create_bq_client(project: str, bigquery_api_endpoint: str, location: str):
+def create_bq_client(*, project: str, bigquery_api_endpoint: str, location: str):
     """Creates a BigQuery client.
 
     Args:
diff --git a/bigquery_magics/graph_server.py b/bigquery_magics/graph_server.py
index 27c47d3..65c301c 100644
--- a/bigquery_magics/graph_server.py
+++ b/bigquery_magics/graph_server.py
@@ -153,9 +153,9 @@ def convert_graph_params(params: Dict[str, Any]):
         query_results = params["query_result"]
     else:
         bq_client = core.create_bq_client(
-            params["args"]["project"],
-            params["args"]["bigquery_api_endpoint"],
-            params["args"]["location"],
+            project=params["args"]["project"],
+            bigquery_api_endpoint=params["args"]["bigquery_api_endpoint"],
+            location=params["args"]["location"],
         )
 
         table_ref = bigquery.TableReference.from_api_repr(params["destination_table"])

From faf644ceba5418c693ea897dda6c5f5490f15076 Mon Sep 17 00:00:00 2001
From: Eric Feiveson <ericfe+github@google.com>
Date: Thu, 5 Feb 2026 15:07:11 -0800
Subject: [PATCH 5/7] Change copyright date to 2026

---
 bigquery_magics/core.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/bigquery_magics/core.py b/bigquery_magics/core.py
index 7bb38d6..4b7287e 100644
--- a/bigquery_magics/core.py
+++ b/bigquery_magics/core.py
@@ -1,4 +1,4 @@
-# Copyright 2024 Google LLC
+# Copyright 2026 Google LLC
 
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

From e18fa8f5c6dfce1b1dd557bf0188c514898a776b Mon Sep 17 00:00:00 2001
From: Eric Feiveson <ericfe+github@google.com>
Date: Thu, 5 Feb 2026 15:10:20 -0800
Subject: [PATCH 6/7] Remove tests for removed function.

---
 tests/unit/bigquery/test_bigquery.py | 32 ----------------------------
 1 file changed, 32 deletions(-)

diff --git a/tests/unit/bigquery/test_bigquery.py b/tests/unit/bigquery/test_bigquery.py
index df61ebe..38bed90 100644
--- a/tests/unit/bigquery/test_bigquery.py
+++ b/tests/unit/bigquery/test_bigquery.py
@@ -3046,35 +3046,3 @@ def test_test_bigquery_magic__extension_loaded__is_registered_set_to_true():
     ip.extension_manager.load_extension("bigquery_magics")
 
     assert bigquery_magics.is_registered is True
-
-
-def test__estimate_json_size_empty():
-    df = pandas.DataFrame()
-    assert magics._estimate_json_size(df) == 2
-
-
-def test__estimate_json_size_bool():
-    df = pandas.DataFrame({"a": [True, False]}, dtype="bool")
-    # key_overhead: (len('"a":') + 1) * 2 = 5 * 2 = 10
-    # structural_overhead: (2 * 2) + 2 + (2 - 1) = 7
-    # total_val_len: 4 (true) + 5 (false) = 9
-    # total = 10 + 7 + 9 = 26
-    assert magics._estimate_json_size(df) == 26
-
-
-def test__estimate_json_size_int():
-    df = pandas.DataFrame({"a": [3, 4]})
-    # key_overhead: (len('"a":') + 1) * 2 = 5 * 2 = 10
-    # structural_overhead: (2 * 2) + 2 + (2 - 1) = 7
-    # total_val_len: 4 (true) + 5 (false) = 9
-    # total = 10 + 7 + 9 = 26
-    assert magics._estimate_json_size(df) == 19
-
-
-def test__estimate_json_size_string():
-    df = pandas.DataFrame({"a": ["x", "yz"]})
-    # key_overhead: (len('"a":') + 1) * 2 = 5 * 2 = 10
-    # structural_overhead: (2 * 2) + 2 + (2 - 1) = 7
-    # total_val_len: (1 + 2) + (2 * 2) = 7
-    # total = 10 + 7 + 7 = 24
-    assert magics._estimate_json_size(df) == 24

From e19456957c4d693d052d4d8953e77dc98164d779 Mon Sep 17 00:00:00 2001
From: Eric Feiveson <ericfe+github@google.com>
Date: Thu, 5 Feb 2026 15:16:58 -0800
Subject: [PATCH 7/7] Fix test_graph_server to account for create_bq_client()
 now taking keyword arguments.

---
 tests/unit/test_graph_server.py | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/tests/unit/test_graph_server.py b/tests/unit/test_graph_server.py
index 773b2eb..065ad54 100644
--- a/tests/unit/test_graph_server.py
+++ b/tests/unit/test_graph_server.py
@@ -527,7 +527,9 @@ def test_post_query_from_table(self):
             response = requests.post(route, json={"params": json.dumps(params)})
             self.assertEqual(response.status_code, 200)
 
-            mock_create.assert_called_once_with("p", "e", "l")
+            mock_create.assert_called_once_with(
+                project="p", bigquery_api_endpoint="e", location="l"
+            )
             mock_client.list_rows.assert_called_once()
 
         response_data = response.json()["response"]