From 0641da0353857317113b2169e547ca69603ddfde Mon Sep 17 00:00:00 2001 From: gRedHeadphone Date: Fri, 19 Dec 2025 06:32:16 +0530 Subject: [PATCH] feat(tools/mysql-get-query-plan): tool impl + docs + tests (#2123) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## 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 Co-authored-by: Yuan Teoh <45984206+Yuan325@users.noreply.github.com> --- cmd/root.go | 1 + docs/en/resources/sources/cloud-sql-mysql.md | 3 + docs/en/resources/sources/mysql.md | 3 + .../tools/mysql/mysql-get-query-plan.md | 39 ++++ .../tools/cloud-sql-mysql.yaml | 9 +- internal/prebuiltconfigs/tools/mysql.yaml | 9 +- .../mysqlgetqueryplan/mysqlgetqueryplan.go | 184 ++++++++++++++++++ .../mysqlgetqueryplan_test.go | 76 ++++++++ .../cloud_sql_mysql_integration_test.go | 1 + tests/common.go | 5 + tests/mysql/mysql_integration_test.go | 1 + tests/tool.go | 75 +++++++ 12 files changed, 390 insertions(+), 16 deletions(-) create mode 100644 docs/en/resources/tools/mysql/mysql-get-query-plan.md create mode 100644 internal/tools/mysql/mysqlgetqueryplan/mysqlgetqueryplan.go create mode 100644 internal/tools/mysql/mysqlgetqueryplan/mysqlgetqueryplan_test.go diff --git a/cmd/root.go b/cmd/root.go index af4efcbbcb..53021dadbb 100644 --- a/cmd/root.go +++ b/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" diff --git a/docs/en/resources/sources/cloud-sql-mysql.md b/docs/en/resources/sources/cloud-sql-mysql.md index 188bcbce26..e9f89f22a9 100644 --- a/docs/en/resources/sources/cloud-sql-mysql.md +++ b/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. diff --git a/docs/en/resources/sources/mysql.md b/docs/en/resources/sources/mysql.md index 44d46195ac..95f2b96d7c 100644 --- a/docs/en/resources/sources/mysql.md +++ b/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. diff --git a/docs/en/resources/tools/mysql/mysql-get-query-plan.md b/docs/en/resources/tools/mysql/mysql-get-query-plan.md new file mode 100644 index 0000000000..d77b81e097 --- /dev/null +++ b/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. | diff --git a/internal/prebuiltconfigs/tools/cloud-sql-mysql.yaml b/internal/prebuiltconfigs/tools/cloud-sql-mysql.yaml index 0a6008eadc..63a73730b7 100644 --- a/internal/prebuiltconfigs/tools/cloud-sql-mysql.yaml +++ b/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 diff --git a/internal/prebuiltconfigs/tools/mysql.yaml b/internal/prebuiltconfigs/tools/mysql.yaml index 9f85de3642..d3068550eb 100644 --- a/internal/prebuiltconfigs/tools/mysql.yaml +++ b/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 diff --git a/internal/tools/mysql/mysqlgetqueryplan/mysqlgetqueryplan.go b/internal/tools/mysql/mysqlgetqueryplan/mysqlgetqueryplan.go new file mode 100644 index 0000000000..34e148b6cc --- /dev/null +++ b/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" +} diff --git a/internal/tools/mysql/mysqlgetqueryplan/mysqlgetqueryplan_test.go b/internal/tools/mysql/mysqlgetqueryplan/mysqlgetqueryplan_test.go new file mode 100644 index 0000000000..b06248dbaf --- /dev/null +++ b/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) + } + }) + } + +} diff --git a/tests/cloudsqlmysql/cloud_sql_mysql_integration_test.go b/tests/cloudsqlmysql/cloud_sql_mysql_integration_test.go index 192c779ea9..55b3035868 100644 --- a/tests/cloudsqlmysql/cloud_sql_mysql_integration_test.go +++ b/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 diff --git a/tests/common.go b/tests/common.go index e2887c5ed9..5ada5a6b32 100644 --- a/tests/common.go +++ b/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 } diff --git a/tests/mysql/mysql_integration_test.go b/tests/mysql/mysql_integration_test.go index 4cb81197be..113767fd1d 100644 --- a/tests/mysql/mysql_integration_test.go +++ b/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) } diff --git a/tests/tool.go b/tests/tool.go index e5ea67a2c3..65a358ca5d 100644 --- a/tests/tool.go +++ b/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.