From a57ebdf4f0c5d9db11ffb7f3b31524ef5d8e7340 Mon Sep 17 00:00:00 2001 From: Karol Droste Date: Tue, 23 Dec 2025 16:59:01 +0000 Subject: [PATCH 1/2] Add example for calling BigQuery via MCP --- examples/bigquery_mcp/main.go | 116 ++++++++++++++++++++++++++++++++++ 1 file changed, 116 insertions(+) create mode 100644 examples/bigquery_mcp/main.go diff --git a/examples/bigquery_mcp/main.go b/examples/bigquery_mcp/main.go new file mode 100644 index 00000000..df5c31cd --- /dev/null +++ b/examples/bigquery_mcp/main.go @@ -0,0 +1,116 @@ +// Copyright 2025 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 +// +// http://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. + +// Package provides an example ADK agent that uses BigQuery via MCP. +package main + +import ( + "context" + "log" + "net/http" + "os" + "os/signal" + + "github.com/modelcontextprotocol/go-sdk/mcp" + "golang.org/x/oauth2" + "google.golang.org/genai" + + "google.golang.org/adk/agent" + "google.golang.org/adk/agent/llmagent" + "google.golang.org/adk/cmd/launcher" + "google.golang.org/adk/cmd/launcher/full" + "google.golang.org/adk/model/gemini" + "google.golang.org/adk/tool" + "google.golang.org/adk/tool/mcptoolset" +) + +// In order to run this example you need to: +// - set the environment variable `GCP_PROJECT_ID` +// - set the environment variable `GOOGLE_API_KEY` +// - set the environment variable `TEST_TOKEN` (you can use command like "export TEST_TOKEN=$(gcloud auth print-access-token)") +// - ensure you have enabled "BigQuery API" (bigquery.googleapis.com) for project indicated in `GCP_PROJECT_ID` +// - ensure you have access to the project indidated in `GCP_PROJECT_ID` +// You can try prompt like: +// `select server date using googlesql from project ` + value of `GCP_PROJECT_ID` + +// TransportWithHeaders adds "X-Goog-User-Project" header to BigQuery MCP calls +type TransportWithHeaders struct { + parent http.RoundTripper + project string // GCP Project ID +} + +func (t *TransportWithHeaders) RoundTrip(req *http.Request) (*http.Response, error) { + if t.project != "" { + req.Header.Add("X-Goog-User-Project", t.project) + } + return t.parent.RoundTrip(req) +} + +func main() { + project := os.Getenv("GCP_PROJECT_ID") + + // Create context that cancels on interrupt signal (Ctrl+C) + ctx, stop := signal.NotifyContext(context.Background(), os.Interrupt) + defer stop() + + model, err := gemini.NewModel(ctx, "gemini-2.5-flash", &genai.ClientConfig{ + APIKey: os.Getenv("GOOGLE_API_KEY"), + }) + if err != nil { + log.Fatalf("Failed to create model: %v", err) + } + + oauthClient := oauth2.NewClient(ctx, oauth2.StaticTokenSource( + &oauth2.Token{AccessToken: os.Getenv("TEST_TOKEN")}, + )) + + transport := &mcp.StreamableClientTransport{ + Endpoint: "https://bigquery.googleapis.com/mcp", + + HTTPClient: &http.Client{ + Transport: &TransportWithHeaders{parent: oauthClient.Transport, project: project}, + }, + } + mcpToolSet, err := mcptoolset.New(mcptoolset.Config{ + Transport: transport, + // WARNING: we need to filter out "get_table_info" tool because of errors is causes ("reference to undefined schema at $defs.RangePartitioning.properties.range") + // we need now just 'execute_sql' tool + ToolFilter: tool.StringPredicate([]string{"execute_sql"}), + }) + if err != nil { + log.Fatalf("Failed to create MCP tool set: %v", err) + } + + // Create LLMAgent with MCP tool set + a, err := llmagent.New(llmagent.Config{ + Name: "helper_agent", + Model: model, + Description: "Helper agent.", + Instruction: "You are a helpful assistant that helps users with various tasks.", + Toolsets: []tool.Toolset{ + mcpToolSet, + }, + }) + if err != nil { + log.Fatalf("Failed to create agent: %v", err) + } + + config := &launcher.Config{ + AgentLoader: agent.NewSingleLoader(a), + } + l := full.NewLauncher() + if err = l.Execute(ctx, config, os.Args[1:]); err != nil { + log.Fatalf("Run failed: %v\n\n%s", err, l.CommandLineSyntax()) + } +} From 4c76152b329e5f6c1340f8ab536d22845f343f73 Mon Sep 17 00:00:00 2001 From: Karol Droste Date: Tue, 23 Dec 2025 17:28:54 +0000 Subject: [PATCH 2/2] Small fixes --- examples/bigquery_mcp/main.go | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/examples/bigquery_mcp/main.go b/examples/bigquery_mcp/main.go index df5c31cd..78f8b245 100644 --- a/examples/bigquery_mcp/main.go +++ b/examples/bigquery_mcp/main.go @@ -52,6 +52,7 @@ type TransportWithHeaders struct { func (t *TransportWithHeaders) RoundTrip(req *http.Request) (*http.Response, error) { if t.project != "" { + req = req.Clone(req.Context()) req.Header.Add("X-Goog-User-Project", t.project) } return t.parent.RoundTrip(req) @@ -59,20 +60,22 @@ func (t *TransportWithHeaders) RoundTrip(req *http.Request) (*http.Response, err func main() { project := os.Getenv("GCP_PROJECT_ID") + apiKey := os.Getenv("GOOGLE_API_KEY") + token := os.Getenv("TEST_TOKEN") // Create context that cancels on interrupt signal (Ctrl+C) ctx, stop := signal.NotifyContext(context.Background(), os.Interrupt) defer stop() model, err := gemini.NewModel(ctx, "gemini-2.5-flash", &genai.ClientConfig{ - APIKey: os.Getenv("GOOGLE_API_KEY"), + APIKey: apiKey, }) if err != nil { log.Fatalf("Failed to create model: %v", err) } oauthClient := oauth2.NewClient(ctx, oauth2.StaticTokenSource( - &oauth2.Token{AccessToken: os.Getenv("TEST_TOKEN")}, + &oauth2.Token{AccessToken: token}, )) transport := &mcp.StreamableClientTransport{