feat: add mysql-execute-sql tool (#577)

This tool can be used across mysql sources.

`mysql-execute-sql` config is as below:

```
tools:
  mysql_execute_sql_tool:
    kind: "mysql-execute-sql"
    source: my-mysql-source
    description: Use this tool to execute sql.
```

The `mysql-execute-sql` tool takes one parameter. Example request as
follow:
```
curl -X POST -H "Content-Type: application/json" -d '{"sql": "SELECT 1"}' http://127.0.0.1:5000/api/tool/mysql_execute_sql_tool/invoke
```
This commit is contained in:
trehanshakuntG
2025-05-20 02:17:39 +05:30
committed by GitHub
parent e89abac29f
commit 8590061ae4
8 changed files with 445 additions and 4 deletions

View File

@@ -0,0 +1,37 @@
---
title: "mysql-execute-sql"
type: docs
weight: 1
description: >
A "mysql-execute-sql" tool executes a SQL statement against a MySQL
database.
---
## About
A `mysql-execute-sql` tool executes 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-execute-sql` takes one input parameter `sql` and run the sql
statement against the `source`.
## Example
```yaml
tools:
execute_sql_tool:
kind: mysql-execute-sql
source: my-mysql-instance
description: Use this tool to execute sql statement.
```
## Reference
| **field** | **type** | **required** | **description** |
|-------------|:------------------------------------------:|:------------:|--------------------------------------------------------------------------------------------------|
| kind | string | true | Must be "mysql-execute-sql". |
| 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. |

View File

@@ -45,6 +45,7 @@ import (
"github.com/googleapis/genai-toolbox/internal/tools/dgraph"
httptool "github.com/googleapis/genai-toolbox/internal/tools/http"
"github.com/googleapis/genai-toolbox/internal/tools/mssqlsql"
"github.com/googleapis/genai-toolbox/internal/tools/mysqlexecutesql"
"github.com/googleapis/genai-toolbox/internal/tools/mysqlsql"
neo4jtool "github.com/googleapis/genai-toolbox/internal/tools/neo4j"
"github.com/googleapis/genai-toolbox/internal/tools/postgresexecutesql"
@@ -412,6 +413,12 @@ func (c *ToolConfigs) UnmarshalYAML(ctx context.Context, unmarshal func(interfac
return fmt.Errorf("unable to parse as %q: %w", kind, err)
}
(*c)[name] = actual
case mysqlexecutesql.ToolKind:
actual := mysqlexecutesql.Config{Name: name}
if err := dec.DecodeContext(ctx, &actual); err != nil {
return fmt.Errorf("unable to parse as %q: %w", kind, err)
}
(*c)[name] = actual
case couchbasetool.ToolKind:
actual := couchbasetool.Config{Name: name}
if err := dec.DecodeContext(ctx, &actual); err != nil {

View File

@@ -0,0 +1,172 @@
// 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 mysqlexecutesql
import (
"context"
"database/sql"
"fmt"
"github.com/googleapis/genai-toolbox/internal/sources"
"github.com/googleapis/genai-toolbox/internal/sources/cloudsqlmysql"
"github.com/googleapis/genai-toolbox/internal/sources/mysql"
"github.com/googleapis/genai-toolbox/internal/tools"
)
const ToolKind string = "mysql-execute-sql"
type compatibleSource interface {
MySQLPool() *sql.DB
}
// validate compatible sources are still compatible
var _ compatibleSource = &cloudsqlmysql.Source{}
var _ compatibleSource = &mysql.Source{}
var compatibleSources = [...]string{cloudsqlmysql.SourceKind, mysql.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 ToolKind
}
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", ToolKind, compatibleSources)
}
sqlParameter := tools.NewStringParameter("sql", "The sql to execute.")
parameters := tools.Parameters{sqlParameter}
mcpManifest := tools.McpManifest{
Name: cfg.Name,
Description: cfg.Description,
InputSchema: parameters.McpManifest(),
}
// finish tool setup
t := Tool{
Name: cfg.Name,
Kind: ToolKind,
Parameters: parameters,
AuthRequired: cfg.AuthRequired,
Pool: s.MySQLPool(),
manifest: tools.Manifest{Description: cfg.Description, Parameters: parameters.Manifest(), AuthRequired: cfg.AuthRequired},
mcpManifest: mcpManifest,
}
return t, nil
}
// validate interface
var _ tools.Tool = Tool{}
type Tool struct {
Name string `yaml:"name"`
Kind string `yaml:"kind"`
AuthRequired []string `yaml:"authRequired"`
Parameters tools.Parameters `yaml:"parameters"`
Pool *sql.DB
manifest tools.Manifest
mcpManifest tools.McpManifest
}
func (t Tool) Invoke(ctx context.Context, params tools.ParamValues) ([]any, error) {
sliceParams := params.AsSlice()
sql, ok := sliceParams[0].(string)
if !ok {
return nil, fmt.Errorf("unable to get cast %s", sliceParams[0])
}
results, err := t.Pool.QueryContext(ctx, sql)
if err != nil {
return nil, fmt.Errorf("unable to execute query: %w", err)
}
cols, err := results.Columns()
if err != nil {
return nil, fmt.Errorf("unable to retrieve rows column name: %w", err)
}
// create an array of values for each column, which can be re-used to scan each row
rawValues := make([]any, len(cols))
values := make([]any, len(cols))
for i := range rawValues {
values[i] = &rawValues[i]
}
var out []any
for results.Next() {
err := results.Scan(values...)
if err != nil {
return nil, fmt.Errorf("unable to parse row: %w", err)
}
vMap := make(map[string]any)
for i, name := range cols {
b, ok := rawValues[i].([]byte)
if ok {
vMap[name] = string(b)
} else {
vMap[name] = rawValues[i]
}
}
out = append(out, vMap)
}
err = results.Close()
if err != nil {
return nil, fmt.Errorf("unable to close rows: %w", err)
}
if err := results.Err(); err != nil {
return nil, fmt.Errorf("errors encountered by results.Scan: %w", err)
}
return out, nil
}
func (t Tool) ParseParams(data map[string]any, claims map[string]map[string]any) (tools.ParamValues, error) {
return tools.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)
}

View File

@@ -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 mysqlexecutesql_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/mysqlexecutesql"
)
func TestParseFromYamlExecuteSql(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-execute-sql
source: my-instance
description: some description
authRequired:
- my-google-auth-service
- other-auth-service
`,
want: server.ToolConfigs{
"example_tool": mysqlexecutesql.Config{
Name: "example_tool",
Kind: mysqlexecutesql.ToolKind,
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)
}
})
}
}

View File

@@ -110,8 +110,8 @@ func TestCloudSQLMysqlToolEndpoints(t *testing.T) {
}
// create table name with UUID
tableNameParam := "param_table_" + strings.Replace(uuid.New().String(), "-", "", -1)
tableNameAuth := "auth_table_" + strings.Replace(uuid.New().String(), "-", "", -1)
tableNameParam := "param_table_" + strings.ReplaceAll(uuid.New().String(), "-", "")
tableNameAuth := "auth_table_" + strings.ReplaceAll(uuid.New().String(), "-", "")
// set up data for param tool
create_statement1, insert_statement1, tool_statement1, params1 := tests.GetMysqlParamToolInfo(tableNameParam)
@@ -125,6 +125,7 @@ func TestCloudSQLMysqlToolEndpoints(t *testing.T) {
// Write config into a file and pass it to command
toolsFile := tests.GetToolsConfig(sourceConfig, CLOUD_SQL_MYSQL_TOOL_KIND, tool_statement1, tool_statement2)
toolsFile = tests.AddMySqlExecuteSqlConfig(t, toolsFile)
cmd, cleanup, err := tests.StartCmd(ctx, toolsFile, args...)
if err != nil {
@@ -145,6 +146,7 @@ func TestCloudSQLMysqlToolEndpoints(t *testing.T) {
select1Want, failInvocationWant := tests.GetMysqlWants()
invokeParamWant, mcpInvokeParamWant := tests.GetNonSpannerInvokeParamWant()
tests.RunToolInvokeTest(t, select1Want, invokeParamWant)
tests.RunMySqlExecuteSqlToolInvokeTest(t, select1Want)
tests.RunMCPToolCallMethod(t, mcpInvokeParamWant, failInvocationWant)
}

View File

@@ -127,6 +127,29 @@ func AddPgExecuteSqlConfig(t *testing.T, config map[string]any) map[string]any {
return config
}
// AddMySqlExecuteSqlConfig gets the tools config for `mysql-execute-sql`
func AddMySqlExecuteSqlConfig(t *testing.T, config map[string]any) map[string]any {
tools, ok := config["tools"].(map[string]any)
if !ok {
t.Fatalf("unable to get tools from config")
}
tools["my-exec-sql-tool"] = map[string]any{
"kind": "mysql-execute-sql",
"source": "my-instance",
"description": "Tool to execute sql",
}
tools["my-auth-exec-sql-tool"] = map[string]any{
"kind": "mysql-execute-sql",
"source": "my-instance",
"description": "Tool to execute sql",
"authRequired": []string{
"my-google-auth",
},
}
config["tools"] = tools
return config
}
// GetHTTPToolsConfig returns a mock HTTP tool's config file
func GetHTTPToolsConfig(sourceConfig map[string]any, toolKind string) map[string]any {
// Write config into a file and pass it to command

View File

@@ -87,8 +87,8 @@ func TestMysqlToolEndpoints(t *testing.T) {
}
// create table name with UUID
tableNameParam := "param_table_" + strings.Replace(uuid.New().String(), "-", "", -1)
tableNameAuth := "auth_table_" + strings.Replace(uuid.New().String(), "-", "", -1)
tableNameParam := "param_table_" + strings.ReplaceAll(uuid.New().String(), "-", "")
tableNameAuth := "auth_table_" + strings.ReplaceAll(uuid.New().String(), "-", "")
// set up data for param tool
create_statement1, insert_statement1, tool_statement1, params1 := tests.GetMysqlParamToolInfo(tableNameParam)
@@ -102,6 +102,7 @@ func TestMysqlToolEndpoints(t *testing.T) {
// Write config into a file and pass it to command
toolsFile := tests.GetToolsConfig(sourceConfig, MYSQL_TOOL_KIND, tool_statement1, tool_statement2)
toolsFile = tests.AddMySqlExecuteSqlConfig(t, toolsFile)
cmd, cleanup, err := tests.StartCmd(ctx, toolsFile, args...)
if err != nil {
@@ -122,5 +123,6 @@ func TestMysqlToolEndpoints(t *testing.T) {
select1Want, failInvocationWant := tests.GetMysqlWants()
invokeParamWant, mcpInvokeParamWant := tests.GetNonSpannerInvokeParamWant()
tests.RunToolInvokeTest(t, select1Want, invokeParamWant)
tests.RunMySqlExecuteSqlToolInvokeTest(t, select1Want)
tests.RunMCPToolCallMethod(t, mcpInvokeParamWant, failInvocationWant)
}

View File

@@ -423,6 +423,128 @@ func RunPgExecuteSqlToolInvokeTest(t *testing.T, select_1_want string) {
}
}
func RunMySqlExecuteSqlToolInvokeTest(t *testing.T, select_1_want string) {
// Get ID token
idToken, err := GetGoogleIdToken(ClientId)
if err != nil {
t.Fatalf("error getting Google ID token: %s", err)
}
// Test tool invoke endpoint
invokeTcs := []struct {
name string
api string
requestHeader map[string]string
requestBody io.Reader
want string
isErr bool
}{
{
name: "invoke my-exec-sql-tool",
api: "http://127.0.0.1:5000/api/tool/my-exec-sql-tool/invoke",
requestHeader: map[string]string{},
requestBody: bytes.NewBuffer([]byte(`{"sql":"SELECT 1"}`)),
want: select_1_want,
isErr: false,
},
{
name: "invoke my-exec-sql-tool create table",
api: "http://127.0.0.1:5000/api/tool/my-exec-sql-tool/invoke",
requestHeader: map[string]string{},
requestBody: bytes.NewBuffer([]byte(`{"sql":"CREATE TABLE t (id SERIAL PRIMARY KEY, name TEXT)"}`)),
want: "null",
isErr: false,
},
{
name: "invoke my-exec-sql-tool select table",
api: "http://127.0.0.1:5000/api/tool/my-exec-sql-tool/invoke",
requestHeader: map[string]string{},
requestBody: bytes.NewBuffer([]byte(`{"sql":"SELECT * FROM t"}`)),
want: "null",
isErr: false,
},
{
name: "invoke my-exec-sql-tool drop table",
api: "http://127.0.0.1:5000/api/tool/my-exec-sql-tool/invoke",
requestHeader: map[string]string{},
requestBody: bytes.NewBuffer([]byte(`{"sql":"DROP TABLE t"}`)),
want: "null",
isErr: false,
},
{
name: "invoke my-exec-sql-tool without body",
api: "http://127.0.0.1:5000/api/tool/my-exec-sql-tool/invoke",
requestHeader: map[string]string{},
requestBody: bytes.NewBuffer([]byte(`{}`)),
isErr: true,
},
{
name: "Invoke my-auth-exec-sql-tool with auth token",
api: "http://127.0.0.1:5000/api/tool/my-auth-exec-sql-tool/invoke",
requestHeader: map[string]string{"my-google-auth_token": idToken},
requestBody: bytes.NewBuffer([]byte(`{"sql":"SELECT 1"}`)),
isErr: false,
want: select_1_want,
},
{
name: "Invoke my-auth-exec-sql-tool with invalid auth token",
api: "http://127.0.0.1:5000/api/tool/my-auth-exec-sql-tool/invoke",
requestHeader: map[string]string{"my-google-auth_token": "INVALID_TOKEN"},
requestBody: bytes.NewBuffer([]byte(`{"sql":"SELECT 1"}`)),
isErr: true,
},
{
name: "Invoke my-auth-exec-sql-tool without auth token",
api: "http://127.0.0.1:5000/api/tool/my-auth-exec-sql-tool/invoke",
requestHeader: map[string]string{},
requestBody: bytes.NewBuffer([]byte(`{"sql":"SELECT 1"}`)),
isErr: true,
},
}
for _, tc := range invokeTcs {
t.Run(tc.name, func(t *testing.T) {
// Send Tool invocation request
req, err := http.NewRequest(http.MethodPost, tc.api, tc.requestBody)
if err != nil {
t.Fatalf("unable to create request: %s", err)
}
req.Header.Add("Content-type", "application/json")
for k, v := range tc.requestHeader {
req.Header.Add(k, v)
}
resp, err := http.DefaultClient.Do(req)
if err != nil {
t.Fatalf("unable to send request: %s", err)
}
defer resp.Body.Close()
if resp.StatusCode != http.StatusOK {
if tc.isErr {
return
}
bodyBytes, _ := io.ReadAll(resp.Body)
t.Fatalf("response status code is not 200, got %d: %s", resp.StatusCode, string(bodyBytes))
}
// Check response body
var body map[string]interface{}
err = json.NewDecoder(resp.Body).Decode(&body)
if err != nil {
t.Fatalf("error parsing response body")
}
got, ok := body["result"].(string)
if !ok {
t.Fatalf("unable to find result in response body")
}
if got != tc.want {
t.Fatalf("unexpected value: got %q, want %q", got, tc.want)
}
})
}
}
// RunMCPToolCallMethod runs the tool/call for mcp endpoint
func RunMCPToolCallMethod(t *testing.T, invoke_param_want, fail_invocation_want string) {
// Test tool invoke endpoint