aiml
SearchMy learning
Menu

Lesson 4 of 6Module 2

Query BigQuery with BigQueryToolset

Query docs analytics through ADK BigQueryToolset with SELECT-only execution, a per-query bytes cap, and offline tests for refused jobs.

Prerequisites

An ADK agent with tools, as in Building Agents with Google ADK 2.0, and the gcloud and bq command-line tools.

By the end, you can

  • Connect an ADK agent to BigQuery with BigQueryToolset and Application Default Credentials, and expose only the tools it needs.
  • Limit what the agent's SQL can do with BigQueryToolConfig: WriteMode.BLOCKED, maximum_bytes_billed, compute_project_id and job labels.
  • Read the errors ADK and BigQuery return for a refused write and an oversized scan, and pin both in offline tests.
Runtime
Python 3.12+, a Google Cloud project with BigQuery and Application Default Credentials, a Gemini API key
Dependencies
google-adk[gcp] 2.12.0 with google-cloud-bigquery 3.46.2

Lab code, tests and fixtures (ZIP, 22 KB)Quiz check your understanding

Progress is saved in this browser.

The docs team keeps page views and site searches in BigQuery. Give its assistant three tools to inspect schemas and query that data, then ask which pages were busiest and which searches returned no results. BigQueryToolset lets the model write SQL. The application permits only SELECT statements, selects the query project and sets a 10 MiB billed-bytes limit on each query job.

When you need this. When an agent answers questions from data that already lives in a warehouse, such as a docs assistant reading traffic or a release bot checking build durations, we let it write SQL against BigQuery instead of exporting the data into a prompt. The toolset fits when the questions are open-ended; for a fixed set of reviewed queries, the preceding MCP Toolbox approach keeps SQL in reviewed tool definitions instead of asking the model to write it.

The tables live in your Google Cloud project. Query results are sent to the model through the tool response.

From a SQL proposal to a checked query job
data flowreturnfailureADK agent processRunnerrun_asyncdocs_analystAgentinvokeGeminigemini-3.8-flashrequestfunction callInMemorySessionServicesession eventseventshistoryBigQueryToolsetexecute_sqlWriteMode.BLOCKEDexecute_sqltool resultSELECT?dry-run resultDELETE refusedno execution jobnoBigQuerydocs tables or public tablecap 10 MiB per query jobyes: query jobrows / errordata flowreturnfailureRunnerrun_asyncInMemorySessionServicedocs_analystAgentinvokeGeminigemini-3.8-flashrequestcallBigQueryToolsetexecute_sqlWriteMode.BLOCKEDexecute_sqltool resultSELECT?dry-run resultDELETErefusedno jobnoBigQuerydocs or public10 MiB capyesrows / error
Runner invokes the docs_analyst Agent and keeps its events in InMemorySessionService. Gemini proposes execute_sql; BigQueryToolset runs the installed ADK tool with Application Default Credentials. ADK checks the statement type returned by a BigQuery dry run, refuses DELETE, and submits SELECT with maximum_bytes_billed. Query rows or an error return through the tool response for the next model turn. The traffic queries read aiml_docs_analytics; the capped query targeted public stackoverflow.posts_questions.

Create the lesson dataset

Download the complete lab ZIP, extract it, and open a terminal in adk-bigquery-toolset/. The commands below use that directory and its pinned dependencies. The [gcp] extra of google-adk installs google-cloud-bigquery, which the toolset needs. Sign in with gcloud auth application-default login, and set GOOGLE_CLOUD_PROJECT to a project where you can create a dataset and run query jobs. The model runs on the Gemini API, so GOOGLE_API_KEY must be set as well.

Prepare the BigQuery clientTerminal · bash
python3.12 -m venv .venv
source .venv/bin/activate
python -m pip install -r requirements-lock.txt
gcloud auth application-default login
export GOOGLE_CLOUD_PROJECT=your-project-id

setup.sh creates the aiml_docs_analytics dataset in the US multi-region with two tables loaded from the CSV files in data/: page_views (one row per page per day) and search_queries (weekly searches and how many results each returned). Load jobs are free, and the whole dataset is 1.5 KB. Every resource carries the label aiml-lesson, so you can find and delete it later.

setup.shSource · bashDownload
#!/usr/bin/env bash
# Create the lesson dataset (two tables, 1.5 KB of CSV) in your project. Load jobs are free.
set -euo pipefail
PROJECT="${GOOGLE_CLOUD_PROJECT:?set GOOGLE_CLOUD_PROJECT}"
LABEL="aiml-lesson:query-bigquery-with-bigquerytoolset"
bq --project_id="${PROJECT}" mk --dataset --location=US --label="${LABEL}" "${PROJECT}:aiml_docs_analytics"
bq --project_id="${PROJECT}" load --source_format=CSV --skip_leading_rows=1 \
  aiml_docs_analytics.page_views data/page_views.csv day:DATE,page:STRING,views:INT64,avg_seconds:INT64
bq --project_id="${PROJECT}" load --source_format=CSV --skip_leading_rows=1 \
  aiml_docs_analytics.search_queries data/search_queries.csv week:DATE,query:STRING,results:INT64,searches:INT64
for table in page_views search_queries; do
  bq --project_id="${PROJECT}" update --set_label="${LABEL}" "${PROJECT}:aiml_docs_analytics.${table}"
done
bq --project_id="${PROJECT}" ls aiml_docs_analytics

Run the reviewed setup script in your selected project:

Create the lesson datasetTerminal · bash
./setup.sh

Configure the toolset and its limits

agent.pySource · pythonDownload
"""The docs assistant answers questions about docs traffic from BigQuery through BigQueryToolset."""
import asyncio
import os
import sys

import google.auth
from google.adk import Agent, Runner
from google.adk.integrations.bigquery import BigQueryCredentialsConfig, BigQueryToolset
from google.adk.integrations.bigquery.config import BigQueryToolConfig, WriteMode
from google.adk.sessions import InMemorySessionService
from google.genai import types

PROJECT = os.environ["GOOGLE_CLOUD_PROJECT"]
DATASET = "aiml_docs_analytics"
MODEL = os.environ.get("DOCS_MODEL", "gemini-3.8-flash")
READ_TOOLS = ["list_table_ids", "get_table_info", "execute_sql"]


def tool_config(**overrides) -> BigQueryToolConfig:
    """The limits ADK enforces on every query, whatever the model writes."""
    settings = dict(
        write_mode=WriteMode.BLOCKED,          # dry-run first; only SELECT statements run
        maximum_bytes_billed=10 * 1024 * 1024,  # 10 MiB, the minimum accepted by this ADK config
        max_query_result_rows=50,
        compute_project_id=PROJECT,             # jobs run and bill only in this project
        default_project_id=PROJECT,
        default_dataset_id=DATASET,
        location="US",
        job_labels={"aiml-lesson": "query-bigquery-with-bigquerytoolset"},
    )
    settings.update(overrides)
    return BigQueryToolConfig(**settings)


def bigquery(config: BigQueryToolConfig | None = None, tool_filter: list[str] | None = READ_TOOLS,
             credentials=None) -> BigQueryToolset:
    # Application Default Credentials: `gcloud auth application-default login` locally,
    # the attached service account on Cloud Run or Agent Runtime.
    if credentials is None:
        credentials, _ = google.auth.default(scopes=["https://www.googleapis.com/auth/bigquery"])
    return BigQueryToolset(
        credentials_config=BigQueryCredentialsConfig(credentials=credentials),
        bigquery_tool_config=config or tool_config(),
        tool_filter=tool_filter,
    )


def build_agent(model=None, toolset: BigQueryToolset | None = None) -> Agent:
    return Agent(
        name="docs_analyst",
        model=model or MODEL,
        instruction=f"""You answer questions about how developers use the docuctl docs.
The data is in the BigQuery dataset {DATASET}: page views per page per day, and weekly site searches
with how many results each returned. Look at the table schema before you write SQL, and select
only the columns you need. Report the numbers the query returned; do not estimate.""",
        tools=[toolset or bigquery()],
        generate_content_config=types.GenerateContentConfig(
            temperature=0.2, thinking_config=types.ThinkingConfig(thinking_level="low")),
    )


async def ask(text: str, agent: Agent | None = None) -> list:
    sessions = InMemorySessionService()
    runner = Runner(app_name="docs", agent=agent or build_agent(), session_service=sessions)
    try:
        session = await sessions.create_session(app_name="docs", user_id="dev")
        message = types.Content(role="user", parts=[types.Part(text=text)])
        return [e async for e in runner.run_async(user_id="dev", session_id=session.id, new_message=message)]
    finally:
        await runner.close()


def show(events: list) -> str:
    lines = []
    for event in events:
        for call in event.get_function_calls():
            lines.append(f"call   {call.name} {call.args}")
        for result in event.get_function_responses():
            text = str(result.response)
            lines.append(f"result {result.name} {text if len(text) <= 400 else text[:400] + ' ...'}")
        if event.is_final_response() and event.author != "user" and event.content:
            lines.append("answer " + "".join(p.text or "" for p in event.content.parts if not p.thought).strip())
    return "\n".join(lines)


async def list_tools(tool_filter: list[str] | None = READ_TOOLS, credentials=None) -> list[str]:
    return [tool.name for tool in await bigquery(tool_filter=tool_filter, credentials=credentials).get_tools()]


if __name__ == "__main__":
    if sys.argv[1:] == ["--tools"]:
        print(asyncio.run(list_tools()))
    else:
        print(show(asyncio.run(ask(" ".join(sys.argv[1:])))))

BigQueryCredentialsConfig(credentials=...) takes any google.auth credentials. Here they are Application Default Credentials: your user account on a laptop, the attached service account on Cloud Run or Agent Runtime. The tools act with exactly the permissions of those credentials, which is why the limits below matter.

tool_filter picks 3 of the toolset’s 11 tools: get_table_info, list_table_ids, execute_sql. The others include forecast and detect_anomalies, which run BigQuery ML functions, and ask_data_insights, which calls a separate service. The example exposes only the three tools it uses. This filter controls the offered declarations; IAM still controls database access.

BigQueryToolConfig supplies client-side query controls:

  • write_mode=WriteMode.BLOCKED is the default, and the code sets it anyway so that a future default change cannot widen access. Before running anything, execute_sql sends the SQL as a dry run and runs it only if BigQuery reports the statement type SELECT. PROTECTED also allows writes to temporary tables in a BigQuery session; ALLOWED runs anything the credentials can do.
  • maximum_bytes_billed is set on every query job. BigQuery checks it before reading data and fails the job without billing it. ADK refuses values below 10,485,760 bytes (10 MiB), the smallest amount an on-demand query bills.
  • compute_project_id makes the tool refuse to run a job in any other project, so this tool refuses a different job project. It can still read tables in other projects when the credentials permit. default_project_id and default_dataset_id make project_id and dataset_id optional for the model.
  • job_labels tags every job with aiml-lesson. The labels show up in billing exports and in the job list, so you can find the activity’s jobs when inspecting usage.

These limits live in the client, inside your process. In production, also give the agent’s service account only roles/bigquery.jobUser on the project and roles/bigquery.dataViewer on the datasets it reads. IAM roles are enforced by BigQuery itself, so they still hold when the client code is wrong. The recorded run used user ADC; a separately restricted service account was not tested.

A write in blocked mode. The dry run identifies DELETE and execute_sql returns Read-only mode only supports SELECT statements. before submitting an execution job. Keep WriteMode.BLOCKED for this read tool; offer writes separately after review.

A cap below 10 MiB. BigQueryToolConfig raises a validation error at startup. Supply at least 10485760 bytes and test the configuration before starting the agent.

A job in another project. With compute_project_id configured, the tool refuses it before calling BigQuery. Set that field to the project intended to pay, alongside the allowed IAM permissions.

Ask the agent about docs traffic

Terminal commandsTerminal · bash
python agent.py --tools
python agent.py "Which three docs pages had the most views in the week of October 12, and which site searches returned no results?"

Inspect the tables and schemas the agent requests before its aggregate queries. Compare the busiest pages with data/page_views.csv and the zero-result searches with data/search_queries.csv. A plausible summary can still omit a date filter or join the wrong rows. The fixture searches include a removed command, docuctl release revert, and a missing topic, canary deploy; these suggest different documentation fixes.

A scan over the bytes cap. BigQuery returns bytesBilledLimitExceeded; the saved job has no billed-byte count. Choose partitioned or summary tables and a cap suited to their scans. This cap applies to one query job, not to the model’s whole run or its Gemini charges.

Try asking the agent to delete a row only in this disposable lesson dataset. The blocked tool should return the read-only error. A generated DELETE suggestion is text; it does not mean a write executed. The offline test exercises the refused-write and over-cap paths without sending a query to BigQuery.

How query settings reach BigQuery

BigQueryToolset wraps plain Python functions as BigQueryTools. Each tool gets the credentials and the tool config at call time, and its declaration hides those arguments from the model. This is what the model saw for execute_sql; project_id is optional because default_project_id is set:

Inspect the saved test fixture

Recorded example. The archived “Query BigQuery with BigQueryToolset” demonstration retains the earlier shipctl scenario; the current lab uses docuctl and does not claim a new model result.

recorded-runs.json: execute_sql parameters as the model saw themRecorded output · jsonDownload
{
  "properties": {
    "project_id": {
      "title": "Project Id",
      "type": "string",
      "description": "Defaults to 'vertex-ai-experminent' when not provided."
    },
    "query": {
      "title": "Query",
      "type": "string"
    },
    "dry_run": {
      "default": false,
      "title": "Dry Run",
      "type": "boolean"
    }
  },
  "required": [
    "query"
  ],
  "title": "execute_sqlParams",
  "type": "object"
}

For an allowed execution, execute_sql first requests a BigQuery dry run and then submits the query. The dry run is a separate request: BigQuery parses the SQL and returns its statement type and the bytes it would scan, without running it or billing anything. In blocked mode, anything other than SELECT stops there. If the statement is allowed and the caller did not request dry-run-only behavior, the execution request carries maximum_bytes_billed, the labels and max_query_result_rows (50 here). If the result fills that limit, the tool adds result_is_likely_truncated so the model knows it has seen only part of the rows.

The execute_sql wrapper catches ordinary exceptions. Errors from BigQuery, such as a syntax error, a missing table or the bytes cap, come back as {"status": "ERROR", "error_details": ...}, and the model reads them like any other result. That keeps the agent running, but it also means the model decides what to do next, so the application must handle repeated requests.

Test the agent offline

test_agent.py replays the recorded model turns and replaces the BigQuery client with a stub, so nothing reaches Google Cloud. The real ADK code still runs: the tool filter, the dry-run check, the job settings and the project guardrail.

Terminal commandsTerminal · bash
python -m unittest -v test_agent.py
Read the offline test
test_agent.pySource · pythonDownload
"""Replay the recorded Gemini turns through the real BigQueryToolset, offline.

Run: python -m unittest -v test_agent.py
The model is replaced by its recorded turns and the BigQuery client by a stub, so nothing reaches
Google Cloud. ADK's own code decides what runs: the dry-run check, the job settings and the guardrails.
"""
import asyncio
import json
import logging
import os
import unittest
import warnings
from pathlib import Path
from types import SimpleNamespace
from unittest import mock

warnings.filterwarnings("ignore")
logging.disable(logging.CRITICAL)

from google.adk.integrations.bigquery import client as bq_client  # noqa: E402
from google.adk.models.base_llm import BaseLlm  # noqa: E402
from google.adk.models.llm_response import LlmResponse  # noqa: E402
from google.genai import types  # noqa: E402
from google.oauth2.credentials import Credentials  # noqa: E402
from pydantic import ValidationError  # noqa: E402

# Fixture project only: cloud clients and model calls are replaced in these tests.
os.environ.setdefault("GOOGLE_CLOUD_PROJECT", "vertex-ai-experminent")
import agent  # noqa: E402

RECORD = json.loads(Path(__file__).with_name("recorded-runs.json").read_text(encoding="utf-8"))
OFFLINE = Credentials(token="offline-test")  # a valid-looking token; the stub never sends it


def setUpModule():
    warnings.simplefilter("ignore")  # unittest resets filters; ADK flags its BigQuery tools as experimental


class ReplayLlm(BaseLlm):
    turns: list

    async def generate_content_async(self, llm_request, stream=False):
        yield LlmResponse(content=types.Content.model_validate(self.turns.pop(0)))


class StubBigQuery:
    """Answers the client calls ADK makes; records every job ADK would have started."""

    def __init__(self, rows):
        self.rows, self.dry_runs, self.executed = list(rows), [], []

    def query(self, query, project=None, job_config=None):
        self.dry_runs.append(query)
        kind = query.lstrip().split()[0].upper()
        return SimpleNamespace(statement_type="SELECT" if kind in ("SELECT", "WITH") else kind,
                               destination=None, to_api_repr=lambda: {"query": query})

    def query_and_wait(self, query, job_config=None, project=None, max_results=None):
        self.executed.append({"query": query, "maximum_bytes_billed": job_config.maximum_bytes_billed,
                              "labels": dict(job_config.labels)})
        return [SimpleNamespace(items=lambda r=row: r.items()) for row in (self.rows.pop(0) if self.rows else [])]

    def list_tables(self, dataset):
        return [SimpleNamespace(table_id=t) for t in ("page_views", "search_queries")]

    def get_table(self, ref):
        return SimpleNamespace(to_api_repr=lambda: {"id": str(ref)})


def errors(key: str) -> list[dict]:
    return [r["response"] for r in RECORD["runs"][key]["tool_results"]
            if isinstance(r["response"], dict) and r["response"].get("status") == "ERROR"]


def replay(key: str) -> StubBigQuery:
    run = RECORD["runs"][key]
    rows = [r["response"].get("rows", []) for r in run["tool_results"]
            if r["tool"] == "execute_sql" and r["response"].get("status") == "SUCCESS"]
    stub = StubBigQuery(rows)
    model = ReplayLlm(model="replay", turns=list(run["model_turns"]))
    with mock.patch.object(bq_client, "get_bigquery_client", return_value=stub):
        asyncio.run(agent.ask("replayed question", agent.build_agent(model, agent.bigquery(credentials=OFFLINE))))
    return stub


class BigQueryToolsetTests(unittest.TestCase):
    def test_tool_filter_offers_only_the_read_tools(self):
        offered = asyncio.run(agent.list_tools(credentials=OFFLINE))
        self.assertEqual(sorted(offered), sorted(RECORD["tools"]))
        hidden = set(RECORD["all_tools"]) - set(offered)
        self.assertTrue({"forecast", "detect_anomalies", "ask_data_insights"} <= hidden)

    def test_every_query_runs_with_the_cap_and_the_lesson_label(self):
        stub = replay("traffic")
        self.assertTrue(stub.executed)
        for job in stub.executed:
            self.assertEqual(job["maximum_bytes_billed"], 10 * 1024 * 1024)
            self.assertEqual(job["labels"]["aiml-lesson"], "query-bigquery-with-bigquerytoolset")

    def test_blocked_mode_rejects_a_delete_before_it_runs(self):
        stub = replay("write")
        self.assertTrue(any(q.lstrip().upper().startswith("DELETE") for q in stub.dry_runs))
        self.assertFalse(any(job["query"].lstrip().upper().startswith("DELETE") for job in stub.executed))
        self.assertEqual(errors("write")[0]["error_details"], "Read-only mode only supports SELECT statements.")

    def test_a_query_over_the_bytes_cap_fails_without_billing(self):
        self.assertIn("Query exceeded limit for bytes billed: 10485760.", errors("bytes")[0]["error_details"])
        capped = [j for j in RECORD["jobs"] if j["error"] == "bytesBilledLimitExceeded"]
        self.assertTrue(capped and all(j["total_bytes_billed"] in (None, 0) for j in capped))

    def test_a_cap_below_10_mib_is_refused_at_startup(self):
        with self.assertRaises(ValidationError):
            agent.tool_config(maximum_bytes_billed=1_000_000)
        self.assertIn("must be set >=10485760", RECORD["cap_below_minimum"])

    def test_compute_project_id_keeps_jobs_in_the_lesson_project(self):
        call = {"role": "model", "parts": [{"function_call": {"name": "execute_sql", "args": {
            "project_id": "bigquery-public-data", "query": "SELECT 1"}}}]}
        done = {"role": "model", "parts": [{"text": "done"}]}
        stub = StubBigQuery([])
        model = ReplayLlm(model="replay", turns=[call, done])
        with mock.patch.object(bq_client, "get_bigquery_client", return_value=stub):
            events = asyncio.run(agent.ask("q", agent.build_agent(model, agent.bigquery(credentials=OFFLINE))))
        self.assertIn("is restricted to execute queries only in the project", agent.show(events))
        self.assertEqual(stub.dry_runs, [])


if __name__ == "__main__":
    unittest.main()

Cost and cleanup

The 10 MiB setting limits each query job. A conversation can submit several jobs, and Gemini requests are charged separately. Review query jobs with the aiml-lesson label in the selected project and compare their processed bytes, billed bytes and errors. Those counters describe jobs, not your total invoice. Check the current pricing page before running a larger activity.

Delete the dataset and both tables when you are done:

cleanup.shSource · bashDownload
#!/usr/bin/env bash
# Delete the lesson dataset and both tables, then show that it is gone.
set -euo pipefail
PROJECT="${GOOGLE_CLOUD_PROJECT:?set GOOGLE_CLOUD_PROJECT}"
bq --project_id="${PROJECT}" rm -r -f -d "${PROJECT}:aiml_docs_analytics"
bq --project_id="${PROJECT}" show "${PROJECT}:aiml_docs_analytics" 2>&1 | grep -o "Not found: Dataset" || exit 1
Terminal commandsTerminal · bash
./cleanup.sh

bq rm -r deletes the tables with the dataset; the last line fails the script unless BigQuery reports the dataset as not found.

Check your understanding

Does maximum_bytes_billed=10485760 cap the cost of the entire agent conversation?

Choose one answer, then check your reasoning.

Read the explanation

A per-query bytes cap and run-level call limits solve different problems. Review both for an application.

What’s next

References

Official API documentation and the installed ADK interfaces were checked October 22, 2026. Exact package versions are pinned in the lab lockfile.

Written by Arjun Prabhulal.

Module 2

Connect databases, search and other agents