Add example for calling BigQuery via MCP #438
Conversation
Summary of ChangesHello @kdroste-google, I'm Gemini Code Assist1! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed! This pull request introduces a new Go example that illustrates how to build an ADK LLM agent capable of interacting with Google BigQuery through the Model Context Protocol (MCP). The example provides a practical demonstration of setting up the necessary authentication, handling project-specific headers for BigQuery requests, and configuring the agent's toolset to perform SQL queries, offering a clear blueprint for similar integrations. Highlights
🧠 New Feature in Public Preview: You can now enable Memory to help Gemini Code Assist learn from your team's feedback. This makes future code reviews more consistent and personalized to your project's style. Click here to enable Memory in your admin console. Using Gemini Code AssistThe full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips. Invoking Gemini You can request assistance from Gemini at any point by creating a comment using either
Customization To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a Limitations & Feedback Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for Github and other Google products, sign up here. You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension. Footnotes
|
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request adds a new example demonstrating how to call BigQuery via MCP using the ADK. The example is well-structured and provides clear instructions for running it. My review includes a few suggestions to improve code quality, robustness, and adherence to Go conventions. Specifically, I've recommended adding validation for required environment variables, improving the http.RoundTripper implementation to avoid side effects, and correcting the package comment to follow Go's documentation standards.
| // 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. |
There was a problem hiding this comment.
According to Go documentation conventions, the package comment for a main package should start with // Command <command-name> ... or // Package main .... The current comment is missing the package/command name.
| // Package provides an example ADK agent that uses BigQuery via MCP. | |
| // Command bigquery_mcp provides an example ADK agent that uses BigQuery via MCP. |
| 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) | ||
| } |
There was a problem hiding this comment.
The RoundTrip method modifies the http.Request object in place by adding a header. The http.RoundTripper interface documentation states that implementations should not modify the request. To adhere to this contract and avoid potential side effects, it's safer to create a clone of the request before adding the header.
| 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 (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) | |
| } |
| 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")}, | ||
| )) |
There was a problem hiding this comment.
The program relies on several environment variables (GCP_PROJECT_ID, GOOGLE_API_KEY, TEST_TOKEN) but doesn't check for their presence and reads them in different places. If any of these are missing, the program will fail later with potentially confusing error messages. It's better to read and validate all required environment variables at the beginning of the main function and exit with a clear error message if they are not set. This also improves code readability by grouping configuration loading.
project := os.Getenv("GCP_PROJECT_ID")
if project == "" {
log.Fatal("GCP_PROJECT_ID environment variable must be set")
}
apiKey := os.Getenv("GOOGLE_API_KEY")
if apiKey == "" {
log.Fatal("GOOGLE_API_KEY environment variable must be set")
}
testToken := os.Getenv("TEST_TOKEN")
if testToken == "" {
log.Fatal("TEST_TOKEN environment variable must be set")
}
// 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: apiKey,
})
if err != nil {
log.Fatalf("Failed to create model: %v", err)
}
oauthClient := oauth2.NewClient(ctx, oauth2.StaticTokenSource(
&oauth2.Token{AccessToken: testToken},
))
|
/gemini review |
There was a problem hiding this comment.
Code Review
This pull request adds a new example demonstrating how to call the BigQuery API using MCP. The code is clear and serves as a good reference. I have provided a few suggestions to improve the example's robustness and fix minor typos in comments. Specifically, adding checks for required environment variables will improve the user experience by providing clearer error messages.
| // - 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` |
There was a problem hiding this comment.
There is a typo in the comment. indidated should be indicated.
| // - ensure you have access to the project indidated in `GCP_PROJECT_ID` | |
| // - ensure you have access to the project indicated in `GCP_PROJECT_ID` |
| project := os.Getenv("GCP_PROJECT_ID") | ||
| apiKey := os.Getenv("GOOGLE_API_KEY") | ||
| token := os.Getenv("TEST_TOKEN") |
There was a problem hiding this comment.
The application expects several environment variables to be set, but it doesn't check for their existence. If they are missing, the program will fail with a less-than-clear error message later on. It's better to check for these variables at startup and provide a clear error message to the user if they are not set.
project := os.Getenv("GCP_PROJECT_ID")
if project == "" {
log.Fatal("GCP_PROJECT_ID environment variable must be set")
}
apiKey := os.Getenv("GOOGLE_API_KEY")
if apiKey == "" {
log.Fatal("GOOGLE_API_KEY environment variable must be set")
}
token := os.Getenv("TEST_TOKEN")
if token == "" {
log.Fatal("TEST_TOKEN environment variable must be set")
}| } | ||
| 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") |
There was a problem hiding this comment.
There is a grammatical error in the comment. is causes should be it causes.
// WARNING: we need to filter out "get_table_info" tool because of errors it causes ("reference to undefined schema at $defs.RangePartitioning.properties.range")1 participant
This example shows how to call BigQuery API using MCP.