feat(tools/mysql-get-query-plan): tool impl + docs + tests (#2123)

## Description

Tool mysql-get-query-plan implementation, along with tests and docs.
Tool used to get information about how MySQL executes a SQL statement
(EXPLAIN).

## PR Checklist

- [x] Make sure you reviewed

[CONTRIBUTING.md](https://github.com/googleapis/genai-toolbox/blob/main/CONTRIBUTING.md)
- [x] Make sure to open an issue as a

[bug/issue](https://github.com/googleapis/genai-toolbox/issues/new/choose)
  before writing your code! That way we can discuss the change, evaluate
  designs, and agree on the general idea
- [x] Ensure the tests and linter pass
- [x] Code coverage does not decrease (if any source code was changed)
- [x] Appropriate docs were updated (if necessary)
- [x] Make sure to add `!` if this involve a breaking change

🛠️ Fixes #1692

---------

Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>
Co-authored-by: Averi Kitsch <akitsch@google.com>
Co-authored-by: Yuan Teoh <45984206+Yuan325@users.noreply.github.com>

cmd/root.go

@@ -169,6 +169,7 @@ import (
 	_ "github.com/googleapis/genai-toolbox/internal/tools/mssql/mssqllisttables"
 	_ "github.com/googleapis/genai-toolbox/internal/tools/mssql/mssqlsql"
 	_ "github.com/googleapis/genai-toolbox/internal/tools/mysql/mysqlexecutesql"
+	_ "github.com/googleapis/genai-toolbox/internal/tools/mysql/mysqlgetqueryplan"
 	_ "github.com/googleapis/genai-toolbox/internal/tools/mysql/mysqllistactivequeries"
 	_ "github.com/googleapis/genai-toolbox/internal/tools/mysql/mysqllisttablefragmentation"
 	_ "github.com/googleapis/genai-toolbox/internal/tools/mysql/mysqllisttables"

docs/en/resources/sources/cloud-sql-mysql.md

@@ -31,6 +31,9 @@ to a database by following these instructions][csql-mysql-quickstart].
 - [`mysql-list-active-queries`](../tools/mysql/mysql-list-active-queries.md)
   List active queries in Cloud SQL for MySQL.
 
+- [`mysql-get-query-plan`](../tools/mysql/mysql-get-query-plan.md)
+  Provide information about how MySQL executes a SQL statement (EXPLAIN).
+
 - [`mysql-list-tables`](../tools/mysql/mysql-list-tables.md)
   List tables in a Cloud SQL for MySQL database.
 

docs/en/resources/sources/mysql.md

@@ -25,6 +25,9 @@ reliability, performance, and ease of use.
 - [`mysql-list-active-queries`](../tools/mysql/mysql-list-active-queries.md)
   List active queries in MySQL.
 
+- [`mysql-get-query-plan`](../tools/mysql/mysql-get-query-plan.md)
+  Provide information about how MySQL executes a SQL statement (EXPLAIN).
+
 - [`mysql-list-tables`](../tools/mysql/mysql-list-tables.md)
   List tables in a MySQL database.
 

docs/en/resources/tools/mysql/mysql-get-query-plan.md

@@ -0,0 +1,39 @@
+---
+title: "mysql-get-query-plan"
+type: docs
+weight: 1
+description: >
+  A "mysql-get-query-plan" tool gets the execution plan for a SQL statement against a MySQL
+  database.
+aliases:
+- /resources/tools/mysql-get-query-plan
+---
+
+## About
+
+A `mysql-get-query-plan` tool gets the execution plan for a SQL statement against a MySQL
+database. It's compatible with any of the following sources:
+
+- [cloud-sql-mysql](../../sources/cloud-sql-mysql.md)
+- [mysql](../../sources/mysql.md)
+
+`mysql-get-query-plan` takes one input parameter `sql_statement` and gets the execution plan for the SQL
+statement against the `source`.
+
+## Example
+
+```yaml
+tools:
+ get_query_plan_tool:
+    kind: mysql-get-query-plan
+    source: my-mysql-instance
+    description: Use this tool to get the execution plan for a sql statement.
+```
+
+## Reference
+
+| **field**   |                  **type**                  | **required** | **description**                                                                                  |
+|-------------|:------------------------------------------:|:------------:|--------------------------------------------------------------------------------------------------|
+| kind        |                   string                   |     true     | Must be "mysql-get-query-plan".                                                                     |
+| source      |                   string                   |     true     | Name of the source the SQL should execute on.                                                    |
+| description |                   string                   |     true     | Description of the tool that is passed to the LLM.                                               |

internal/prebuiltconfigs/tools/cloud-sql-mysql.yaml

@@ -32,16 +32,9 @@ tools:
     source: cloud-sql-mysql-source
     description: Lists top N (default 10) ongoing queries from processlist and innodb_trx, ordered by execution time in descending order. Returns detailed information of those queries in json format, including process id, query, transaction duration, transaction wait duration, process time, transaction state, process state, username with host, transaction rows locked, transaction rows modified, and db schema.
   get_query_plan:
-    kind: mysql-sql
+    kind: mysql-get-query-plan
     source: cloud-sql-mysql-source
     description: "Provide information about how MySQL executes a SQL statement. Common use cases include: 1) analyze query plan to improve its performance, and 2) determine effectiveness of existing indexes and evalueate new ones."
-    statement: |
-      EXPLAIN FORMAT=JSON {{.sql_statement}};
-    templateParameters:
-      - name: sql_statement
-        type: string
-        description: "the SQL statement to explain"
-        required: true
   list_tables:
     kind: mysql-list-tables
     source: cloud-sql-mysql-source

internal/prebuiltconfigs/tools/mysql.yaml

@@ -36,16 +36,9 @@ tools:
     source: mysql-source
     description: Lists top N (default 10) ongoing queries from processlist and innodb_trx, ordered by execution time in descending order. Returns detailed information of those queries in json format, including process id, query, transaction duration, transaction wait duration, process time, transaction state, process state, username with host, transaction rows locked, transaction rows modified, and db schema.
   get_query_plan:
-    kind: mysql-sql
+    kind: mysql-get-query-plan
     source: mysql-source
     description: "Provide information about how MySQL executes a SQL statement. Common use cases include: 1) analyze query plan to improve its performance, and 2) determine effectiveness of existing indexes and evalueate new ones."
-    statement: |
-      EXPLAIN FORMAT=JSON {{.sql_statement}};
-    templateParameters:
-      - name: sql_statement
-        type: string
-        description: "the SQL statement to explain"
-        required: true
   list_tables:
     kind: mysql-list-tables
     source: mysql-source

internal/tools/mysql/mysqlgetqueryplan/mysqlgetqueryplan.go

@@ -0,0 +1,184 @@
+// 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 mysqlgetqueryplan
+
+import (
+	"context"
+	"database/sql"
+	"encoding/json"
+	"fmt"
+
+	yaml "github.com/goccy/go-yaml"
+	"github.com/googleapis/genai-toolbox/internal/sources"
+	"github.com/googleapis/genai-toolbox/internal/sources/cloudsqlmysql"
+	"github.com/googleapis/genai-toolbox/internal/sources/mindsdb"
+	"github.com/googleapis/genai-toolbox/internal/sources/mysql"
+	"github.com/googleapis/genai-toolbox/internal/tools"
+	"github.com/googleapis/genai-toolbox/internal/util"
+	"github.com/googleapis/genai-toolbox/internal/util/parameters"
+)
+
+const kind string = "mysql-get-query-plan"
+
+func init() {
+	if !tools.Register(kind, newConfig) {
+		panic(fmt.Sprintf("tool kind %q already registered", kind))
+	}
+}
+
+func newConfig(ctx context.Context, name string, decoder *yaml.Decoder) (tools.ToolConfig, error) {
+	actual := Config{Name: name}
+	if err := decoder.DecodeContext(ctx, &actual); err != nil {
+		return nil, err
+	}
+	return actual, nil
+}
+
+type compatibleSource interface {
+	MySQLPool() *sql.DB
+}
+
+// validate compatible sources are still compatible
+var _ compatibleSource = &cloudsqlmysql.Source{}
+var _ compatibleSource = &mysql.Source{}
+var _ compatibleSource = &mindsdb.Source{}
+
+var compatibleSources = [...]string{cloudsqlmysql.SourceKind, mysql.SourceKind, mindsdb.SourceKind}
+
+type Config struct {
+	Name         string   `yaml:"name" validate:"required"`
+	Kind         string   `yaml:"kind" validate:"required"`
+	Source       string   `yaml:"source" validate:"required"`
+	Description  string   `yaml:"description" validate:"required"`
+	AuthRequired []string `yaml:"authRequired"`
+}
+
+// validate interface
+var _ tools.ToolConfig = Config{}
+
+func (cfg Config) ToolConfigKind() string {
+	return kind
+}
+
+func (cfg Config) Initialize(srcs map[string]sources.Source) (tools.Tool, error) {
+	// verify source exists
+	rawS, ok := srcs[cfg.Source]
+	if !ok {
+		return nil, fmt.Errorf("no source named %q configured", cfg.Source)
+	}
+
+	// verify the source is compatible
+	s, ok := rawS.(compatibleSource)
+	if !ok {
+		return nil, fmt.Errorf("invalid source for %q tool: source kind must be one of %q", kind, compatibleSources)
+	}
+
+	sqlParameter := parameters.NewStringParameter("sql_statement", "The sql statement to explain.")
+	params := parameters.Parameters{sqlParameter}
+
+	mcpManifest := tools.GetMcpManifest(cfg.Name, cfg.Description, cfg.AuthRequired, params, nil)
+
+	// finish tool setup
+	t := Tool{
+		Config:      cfg,
+		Parameters:  params,
+		Pool:        s.MySQLPool(),
+		manifest:    tools.Manifest{Description: cfg.Description, Parameters: params.Manifest(), AuthRequired: cfg.AuthRequired},
+		mcpManifest: mcpManifest,
+	}
+	return t, nil
+}
+
+// validate interface
+var _ tools.Tool = Tool{}
+
+type Tool struct {
+	Config
+	Parameters parameters.Parameters `yaml:"parameters"`
+
+	Pool        *sql.DB
+	manifest    tools.Manifest
+	mcpManifest tools.McpManifest
+}
+
+func (t Tool) Invoke(ctx context.Context, resourceMgr tools.SourceProvider, params parameters.ParamValues, accessToken tools.AccessToken) (any, error) {
+	paramsMap := params.AsMap()
+	sql, ok := paramsMap["sql_statement"].(string)
+	if !ok {
+		return nil, fmt.Errorf("unable to get cast %s", paramsMap["sql_statement"])
+	}
+
+	// Log the query executed for debugging.
+	logger, err := util.LoggerFromContext(ctx)
+	if err != nil {
+		return nil, fmt.Errorf("error getting logger: %s", err)
+	}
+	logger.DebugContext(ctx, fmt.Sprintf("executing `%s` tool query: %s", kind, sql))
+
+	query := fmt.Sprintf("EXPLAIN FORMAT=JSON %s", sql)
+	results, err := t.Pool.QueryContext(ctx, query)
+	if err != nil {
+		return nil, fmt.Errorf("unable to execute query: %w", err)
+	}
+	defer results.Close()
+
+	var plan string
+	if results.Next() {
+		if err := results.Scan(&plan); err != nil {
+			return nil, fmt.Errorf("unable to parse row: %w", err)
+		}
+	} else {
+		return nil, fmt.Errorf("no query plan returned")
+	}
+
+	if err := results.Err(); err != nil {
+		return nil, fmt.Errorf("errors encountered during row iteration: %w", err)
+	}
+
+	var out any
+	if err := json.Unmarshal([]byte(plan), &out); err != nil {
+		return nil, fmt.Errorf("failed to unmarshal query plan json: %w", err)
+	}
+
+	return out, nil
+}
+
+func (t Tool) ParseParams(data map[string]any, claims map[string]map[string]any) (parameters.ParamValues, error) {
+	return parameters.ParseParams(t.Parameters, data, claims)
+}
+
+func (t Tool) Manifest() tools.Manifest {
+	return t.manifest
+}
+
+func (t Tool) McpManifest() tools.McpManifest {
+	return t.mcpManifest
+}
+
+func (t Tool) Authorized(verifiedAuthServices []string) bool {
+	return tools.IsAuthorized(t.AuthRequired, verifiedAuthServices)
+}
+
+func (t Tool) RequiresClientAuthorization(resourceMgr tools.SourceProvider) bool {
+	return false
+}
+
+func (t Tool) ToConfig() tools.ToolConfig {
+	return t.Config
+}
+
+func (t Tool) GetAuthTokenHeaderName() string {
+	return "Authorization"
+}

internal/tools/mysql/mysqlgetqueryplan/mysqlgetqueryplan_test.go

@@ -0,0 +1,76 @@
+// 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 mysqlgetqueryplan_test
+
+import (
+	"testing"
+
+	yaml "github.com/goccy/go-yaml"
+	"github.com/google/go-cmp/cmp"
+	"github.com/googleapis/genai-toolbox/internal/server"
+	"github.com/googleapis/genai-toolbox/internal/testutils"
+	"github.com/googleapis/genai-toolbox/internal/tools/mysql/mysqlgetqueryplan"
+)
+
+func TestParseFromYamlGetQueryPlan(t *testing.T) {
+	ctx, err := testutils.ContextWithNewLogger()
+	if err != nil {
+		t.Fatalf("unexpected error: %s", err)
+	}
+	tcs := []struct {
+		desc string
+		in   string
+		want server.ToolConfigs
+	}{
+		{
+			desc: "basic example",
+			in: `
+			tools:
+				example_tool:
+					kind: mysql-get-query-plan
+					source: my-instance
+					description: some description
+					authRequired:
+						- my-google-auth-service
+						- other-auth-service
+			`,
+			want: server.ToolConfigs{
+				"example_tool": mysqlgetqueryplan.Config{
+					Name:         "example_tool",
+					Kind:         "mysql-get-query-plan",
+					Source:       "my-instance",
+					Description:  "some description",
+					AuthRequired: []string{"my-google-auth-service", "other-auth-service"},
+				},
+			},
+		},
+	}
+	for _, tc := range tcs {
+		t.Run(tc.desc, func(t *testing.T) {
+			got := struct {
+				Tools server.ToolConfigs `yaml:"tools"`
+			}{}
+			// Parse contents
+			err := yaml.UnmarshalContext(ctx, testutils.FormatYaml(tc.in), &got)
+			if err != nil {
+				t.Fatalf("unable to unmarshal: %s", err)
+			}
+			if diff := cmp.Diff(tc.want, got.Tools); diff != "" {
+				t.Fatalf("incorrect parse: diff %v", diff)
+			}
+		})
+	}
+
+}

tests/cloudsqlmysql/cloud_sql_mysql_integration_test.go

@@ -163,6 +163,7 @@ func TestCloudSQLMySQLToolEndpoints(t *testing.T) {
 	const expectedOwner = "'toolbox-identity'@'%'"
 	tests.RunMySQLListTablesTest(t, CloudSQLMySQLDatabase, tableNameParam, tableNameAuth, expectedOwner)
 	tests.RunMySQLListActiveQueriesTest(t, ctx, pool)
+	tests.RunMySQLGetQueryPlanTest(t, ctx, pool, CloudSQLMySQLDatabase, tableNameParam)
 }
 
 // Test connection with different IP type

tests/common.go

@@ -448,6 +448,11 @@ func AddMySQLPrebuiltToolConfig(t *testing.T, config map[string]any) map[string]
 		"source":      "my-instance",
 		"description": "Lists table fragmentation in the database.",
 	}
+	tools["get_query_plan"] = map[string]any{
+		"kind":        "mysql-get-query-plan",
+		"source":      "my-instance",
+		"description": "Gets the query plan for a SQL statement.",
+	}
 	config["tools"] = tools
 	return config
 }

tests/mysql/mysql_integration_test.go

@@ -143,4 +143,5 @@ func TestMySQLToolEndpoints(t *testing.T) {
 	tests.RunMySQLListActiveQueriesTest(t, ctx, pool)
 	tests.RunMySQLListTablesMissingUniqueIndexes(t, ctx, pool, MySQLDatabase)
 	tests.RunMySQLListTableFragmentationTest(t, MySQLDatabase, tableNameParam, tableNameAuth)
+	tests.RunMySQLGetQueryPlanTest(t, ctx, pool, MySQLDatabase, tableNameParam)
 }

tests/tool.go

@@ -3377,6 +3377,81 @@ func RunMySQLListTableFragmentationTest(t *testing.T, databaseName, tableNamePar
 	}
 }
 
+func RunMySQLGetQueryPlanTest(t *testing.T, ctx context.Context, pool *sql.DB, databaseName, tableNameParam string) {
+	// Create a simple query to explain
+	query := fmt.Sprintf("SELECT * FROM %s", tableNameParam)
+
+	invokeTcs := []struct {
+		name           string
+		requestBody    io.Reader
+		wantStatusCode int
+		checkResult    func(t *testing.T, result any)
+	}{
+		{
+			name:           "invoke get_query_plan with valid query",
+			requestBody:    bytes.NewBufferString(fmt.Sprintf(`{"sql_statement": "%s"}`, query)),
+			wantStatusCode: http.StatusOK,
+			checkResult: func(t *testing.T, result any) {
+				resultMap, ok := result.(map[string]any)
+				if !ok {
+					t.Fatalf("result should be a map, got %T", result)
+				}
+				if _, ok := resultMap["query_block"]; !ok {
+					t.Errorf("result should contain 'query_block', got %v", resultMap)
+				}
+			},
+		},
+		{
+			name:           "invoke get_query_plan with invalid query",
+			requestBody:    bytes.NewBufferString(`{"sql_statement": "SELECT * FROM non_existent_table"}`),
+			wantStatusCode: http.StatusBadRequest,
+			checkResult:    nil,
+		},
+	}
+
+	for _, tc := range invokeTcs {
+		t.Run(tc.name, func(t *testing.T) {
+			const api = "http://127.0.0.1:5000/api/tool/get_query_plan/invoke"
+			resp, respBytes := RunRequest(t, http.MethodPost, api, tc.requestBody, nil)
+			if resp.StatusCode != tc.wantStatusCode {
+				t.Fatalf("wrong status code: got %d, want %d, body: %s", resp.StatusCode, tc.wantStatusCode, string(respBytes))
+			}
+			if tc.wantStatusCode != http.StatusOK {
+				return
+			}
+
+			var bodyWrapper map[string]json.RawMessage
+
+			if err := json.Unmarshal(respBytes, &bodyWrapper); err != nil {
+				t.Fatalf("error parsing response wrapper: %s, body: %s", err, string(respBytes))
+			}
+
+			resultJSON, ok := bodyWrapper["result"]
+			if !ok {
+				t.Fatal("unable to find 'result' in response body")
+			}
+
+			var resultString string
+			if err := json.Unmarshal(resultJSON, &resultString); err != nil {
+				if string(resultJSON) == "null" {
+					resultString = "null"
+				} else {
+					t.Fatalf("'result' is not a JSON-encoded string: %s", err)
+				}
+			}
+
+			var got map[string]any
+			if err := json.Unmarshal([]byte(resultString), &got); err != nil {
+				t.Fatalf("failed to unmarshal actual result string: %v", err)
+			}
+
+			if tc.checkResult != nil {
+				tc.checkResult(t, got)
+			}
+		})
+	}
+}
+
 // RunMSSQLListTablesTest run tests againsts the mssql-list-tables tools.
 func RunMSSQLListTablesTest(t *testing.T, tableNameParam, tableNameAuth string) {
 	// TableNameParam columns to construct want.