Skip to content

Commit b81fc6a

Browse files
Yuan325Yuan325
authored
feat: add templateParameters field for mssqlsql (#671)
Add `templateParameters to support non-filter parameters and DDL statements. Part of #535
1 parent 29aa0a7 commit b81fc6a

5 files changed

Lines changed: 141 additions & 27 deletions

File tree

‎internal/tools/mssqlsql/mssqlsql.go‎

Lines changed: 45 additions & 27 deletions
Original file line numberDiff line numberDiff line change
@@ -54,13 +54,14 @@ var _ compatibleSource = &mssql.Source{}
5454
var compatibleSources = [...]string{cloudsqlmssql.SourceKind, mssql.SourceKind}
5555

5656
type Config struct {
57-
Name string `yaml:"name" validate:"required"`
58-
Kind string `yaml:"kind" validate:"required"`
59-
Source string `yaml:"source" validate:"required"`
60-
Description string `yaml:"description" validate:"required"`
61-
Statement string `yaml:"statement" validate:"required"`
62-
AuthRequired []string `yaml:"authRequired"`
63-
Parameters tools.Parameters `yaml:"parameters"`
57+
Name string `yaml:"name" validate:"required"`
58+
Kind string `yaml:"kind" validate:"required"`
59+
Source string `yaml:"source" validate:"required"`
60+
Description string `yaml:"description" validate:"required"`
61+
Statement string `yaml:"statement" validate:"required"`
62+
AuthRequired []string `yaml:"authRequired"`
63+
Parameters tools.Parameters `yaml:"parameters"`
64+
TemplateParameters tools.Parameters `yaml:"templateParameters"`
6465
}
6566

6667
// validate interface
@@ -83,22 +84,26 @@ func (cfg Config) Initialize(srcs map[string]sources.Source) (tools.Tool, error)
8384
return nil, fmt.Errorf("invalid source for %q tool: source kind must be one of %q", kind, compatibleSources)
8485
}
8586

87+
allParameters, paramManifest, paramMcpManifest := tools.ProcessParameters(cfg.TemplateParameters, cfg.Parameters)
88+
8689
mcpManifest := tools.McpManifest{
8790
Name: cfg.Name,
8891
Description: cfg.Description,
89-
InputSchema: cfg.Parameters.McpManifest(),
92+
InputSchema: paramMcpManifest,
9093
}
9194

9295
// finish tool setup
9396
t := Tool{
94-
Name: cfg.Name,
95-
Kind: kind,
96-
Parameters: cfg.Parameters,
97-
Statement: cfg.Statement,
98-
AuthRequired: cfg.AuthRequired,
99-
Db: s.MSSQLDB(),
100-
manifest: tools.Manifest{Description: cfg.Description, Parameters: cfg.Parameters.Manifest(), AuthRequired: cfg.AuthRequired},
101-
mcpManifest: mcpManifest,
97+
Name: cfg.Name,
98+
Kind: kind,
99+
Parameters: cfg.Parameters,
100+
TemplateParameters: cfg.TemplateParameters,
101+
AllParams: allParameters,
102+
Statement: cfg.Statement,
103+
AuthRequired: cfg.AuthRequired,
104+
Db: s.MSSQLDB(),
105+
manifest: tools.Manifest{Description: cfg.Description, Parameters: paramManifest, AuthRequired: cfg.AuthRequired},
106+
mcpManifest: mcpManifest,
102107
}
103108
return t, nil
104109
}
@@ -107,10 +112,12 @@ func (cfg Config) Initialize(srcs map[string]sources.Source) (tools.Tool, error)
107112
var _ tools.Tool = Tool{}
108113

109114
type Tool struct {
110-
Name string `yaml:"name"`
111-
Kind string `yaml:"kind"`
112-
AuthRequired []string `yaml:"authRequired"`
113-
Parameters tools.Parameters `yaml:"parameters"`
115+
Name string `yaml:"name"`
116+
Kind string `yaml:"kind"`
117+
AuthRequired []string `yaml:"authRequired"`
118+
Parameters tools.Parameters `yaml:"parameters"`
119+
TemplateParameters tools.Parameters `yaml:"templateParameters"`
120+
AllParams tools.Parameters `yaml:"allParams"`
114121

115122
Db *sql.DB
116123
Statement string
@@ -119,18 +126,29 @@ type Tool struct {
119126
}
120127

121128
func (t Tool) Invoke(ctx context.Context, params tools.ParamValues) ([]any, error) {
122-
namedArgs := make([]any, 0, len(params))
123-
paramsMap := params.AsReversedMap()
129+
paramsMap := params.AsMap()
130+
newStatement, err := tools.ResolveTemplateParams(t.TemplateParameters, t.Statement, paramsMap)
131+
if err != nil {
132+
return nil, fmt.Errorf("unable to extract template params %w", err)
133+
}
134+
135+
newParams, err := tools.GetParams(t.Parameters, paramsMap)
136+
if err != nil {
137+
return nil, fmt.Errorf("unable to extract standard params %w", err)
138+
}
139+
140+
namedArgs := make([]any, 0, len(newParams))
141+
newParamsMap := newParams.AsReversedMap()
124142
// To support both named args (e.g @id) and positional args (e.g @p1), check if arg name is contained in the statement.
125-
for _, v := range params.AsSlice() {
126-
paramName := paramsMap[v]
127-
if strings.Contains(t.Statement, "@"+paramName) {
143+
for _, v := range newParams.AsSlice() {
144+
paramName := newParamsMap[v]
145+
if strings.Contains(newStatement, "@"+paramName) {
128146
namedArgs = append(namedArgs, sql.Named(paramName, v))
129147
} else {
130148
namedArgs = append(namedArgs, v)
131149
}
132150
}
133-
rows, err := t.Db.QueryContext(ctx, t.Statement, namedArgs...)
151+
rows, err := t.Db.QueryContext(ctx, newStatement, namedArgs...)
134152
if err != nil {
135153
return nil, fmt.Errorf("unable to execute query: %w", err)
136154
}
@@ -173,7 +191,7 @@ func (t Tool) Invoke(ctx context.Context, params tools.ParamValues) ([]any, erro
173191
}
174192

175193
func (t Tool) ParseParams(data map[string]any, claims map[string]map[string]any) (tools.ParamValues, error) {
176-
return tools.ParseParams(t.Parameters, data, claims)
194+
return tools.ParseParams(t.AllParams, data, claims)
177195
}
178196

179197
func (t Tool) Manifest() tools.Manifest {

‎internal/tools/mssqlsql/mssqlsql_test.go‎

Lines changed: 81 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -90,5 +90,86 @@ func TestParseFromYamlMssql(t *testing.T) {
9090
}
9191
})
9292
}
93+
}
9394

95+
func TestParseFromYamlWithTemplateMssql(t *testing.T) {
96+
ctx, err := testutils.ContextWithNewLogger()
97+
if err != nil {
98+
t.Fatalf("unexpected error: %s", err)
99+
}
100+
tcs := []struct {
101+
desc string
102+
in string
103+
want server.ToolConfigs
104+
}{
105+
{
106+
desc: "basic example",
107+
in: `
108+
tools:
109+
example_tool:
110+
kind: mssql-sql
111+
source: my-instance
112+
description: some description
113+
statement: |
114+
SELECT * FROM SQL_STATEMENT;
115+
authRequired:
116+
- my-google-auth-service
117+
- other-auth-service
118+
parameters:
119+
- name: country
120+
type: string
121+
description: some description
122+
authServices:
123+
- name: my-google-auth-service
124+
field: user_id
125+
- name: other-auth-service
126+
field: user_id
127+
templateParameters:
128+
- name: tableName
129+
type: string
130+
description: The table to select hotels from.
131+
- name: fieldArray
132+
type: array
133+
description: The columns to return for the query.
134+
items:
135+
name: column
136+
type: string
137+
description: A column name that will be returned from the query.
138+
`,
139+
want: server.ToolConfigs{
140+
"example_tool": mssqlsql.Config{
141+
Name: "example_tool",
142+
Kind: "mssql-sql",
143+
Source: "my-instance",
144+
Description: "some description",
145+
Statement: "SELECT * FROM SQL_STATEMENT;\n",
146+
AuthRequired: []string{"my-google-auth-service", "other-auth-service"},
147+
Parameters: []tools.Parameter{
148+
tools.NewStringParameterWithAuth("country", "some description",
149+
[]tools.ParamAuthService{{Name: "my-google-auth-service", Field: "user_id"},
150+
{Name: "other-auth-service", Field: "user_id"}}),
151+
},
152+
TemplateParameters: []tools.Parameter{
153+
tools.NewStringParameter("tableName", "The table to select hotels from."),
154+
tools.NewArrayParameter("fieldArray", "The columns to return for the query.", tools.NewStringParameter("column", "A column name that will be returned from the query.")),
155+
},
156+
},
157+
},
158+
},
159+
}
160+
for _, tc := range tcs {
161+
t.Run(tc.desc, func(t *testing.T) {
162+
got := struct {
163+
Tools server.ToolConfigs `yaml:"tools"`
164+
}{}
165+
// Parse contents
166+
err := yaml.UnmarshalContext(ctx, testutils.FormatYaml(tc.in), &got)
167+
if err != nil {
168+
t.Fatalf("unable to unmarshal: %s", err)
169+
}
170+
if diff := cmp.Diff(tc.want, got.Tools); diff != "" {
171+
t.Fatalf("incorrect parse: diff %v", diff)
172+
}
173+
})
174+
}
94175
}

‎tests/cloudsqlmssql/cloud_sql_mssql_integration_test.go‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -118,6 +118,7 @@ func TestCloudSQLMssqlToolEndpoints(t *testing.T) {
118118
// create table name with UUID
119119
tableNameParam := "param_table_" + strings.ReplaceAll(uuid.New().String(), "-", "")
120120
tableNameAuth := "auth_table_" + strings.ReplaceAll(uuid.New().String(), "-", "")
121+
tableNameTemplateParam := "template_param_table_" + strings.ReplaceAll(uuid.New().String(), "-", "")
121122

122123
// set up data for param tool
123124
create_statement1, insert_statement1, tool_statement1, params1 := tests.GetMssqlParamToolInfo(tableNameParam)
@@ -132,6 +133,8 @@ func TestCloudSQLMssqlToolEndpoints(t *testing.T) {
132133
// Write config into a file and pass it to command
133134
toolsFile := tests.GetToolsConfig(sourceConfig, CLOUD_SQL_MSSQL_TOOL_KIND, tool_statement1, tool_statement2)
134135
toolsFile = tests.AddMssqlExecuteSqlConfig(t, toolsFile)
136+
tmplSelectCombined, tmplSelectFilterCombined := tests.GetMssqlTmplToolStatement()
137+
toolsFile = tests.AddTemplateParamConfig(t, toolsFile, CLOUD_SQL_MSSQL_TOOL_KIND, tmplSelectCombined, tmplSelectFilterCombined)
135138

136139
cmd, cleanup, err := tests.StartCmd(ctx, toolsFile, args...)
137140
if err != nil {
@@ -154,6 +157,7 @@ func TestCloudSQLMssqlToolEndpoints(t *testing.T) {
154157
tests.RunToolInvokeTest(t, select1Want, invokeParamWant)
155158
tests.RunExecuteSqlToolInvokeTest(t, createTableStatement, select1Want)
156159
tests.RunMCPToolCallMethod(t, mcpInvokeParamWant, failInvocationWant)
160+
tests.RunToolInvokeWithTemplateParameters(t, tableNameTemplateParam)
157161
}
158162

159163
// Test connection with different IP type

‎tests/common.go‎

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -297,6 +297,13 @@ func GetMssqlAuthToolInfo(tableName string) (string, string, string, []any) {
297297
return create_statement, insert_statement, tool_statement, params
298298
}
299299

300+
// GetMssqlTmplToolStatement returns statements and param for template parameter test cases for mysql-sql kind
301+
func GetMssqlTmplToolStatement() (string, string) {
302+
tmplSelectCombined := "SELECT * FROM {{.tableName}} WHERE id = @id"
303+
tmplSelectFilterCombined := "SELECT * FROM {{.tableName}} WHERE {{.columnFilter}} = @name"
304+
return tmplSelectCombined, tmplSelectFilterCombined
305+
}
306+
300307
// GetMysqlParamToolInfo returns statements and param for my-param-tool mysql-sql kind
301308
func GetMysqlParamToolInfo(tableName string) (string, string, string, []any) {
302309
create_statement := fmt.Sprintf("CREATE TABLE %s (id INT NOT NULL AUTO_INCREMENT PRIMARY KEY, name VARCHAR(255));", tableName)

‎tests/mssql/mssql_integration_test.go‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -90,6 +90,7 @@ func TestMssqlToolEndpoints(t *testing.T) {
9090
// create table name with UUID
9191
tableNameParam := "param_table_" + strings.ReplaceAll(uuid.New().String(), "-", "")
9292
tableNameAuth := "auth_table_" + strings.ReplaceAll(uuid.New().String(), "-", "")
93+
tableNameTemplateParam := "template_param_table_" + strings.ReplaceAll(uuid.New().String(), "-", "")
9394

9495
// set up data for param tool
9596
create_statement1, insert_statement1, tool_statement1, params1 := tests.GetMssqlParamToolInfo(tableNameParam)
@@ -104,6 +105,8 @@ func TestMssqlToolEndpoints(t *testing.T) {
104105
// Write config into a file and pass it to command
105106
toolsFile := tests.GetToolsConfig(sourceConfig, MSSQL_TOOL_KIND, tool_statement1, tool_statement2)
106107
toolsFile = tests.AddMssqlExecuteSqlConfig(t, toolsFile)
108+
tmplSelectCombined, tmplSelectFilterCombined := tests.GetMssqlTmplToolStatement()
109+
toolsFile = tests.AddTemplateParamConfig(t, toolsFile, MSSQL_TOOL_KIND, tmplSelectCombined, tmplSelectFilterCombined)
107110

108111
cmd, cleanup, err := tests.StartCmd(ctx, toolsFile, args...)
109112
if err != nil {
@@ -126,4 +129,5 @@ func TestMssqlToolEndpoints(t *testing.T) {
126129
tests.RunToolInvokeTest(t, select1Want, invokeParamWant)
127130
tests.RunExecuteSqlToolInvokeTest(t, createTableStatement, select1Want)
128131
tests.RunMCPToolCallMethod(t, mcpInvokeParamWant, failInvocationWant)
132+
tests.RunToolInvokeWithTemplateParameters(t, tableNameTemplateParam)
129133
}

0 commit comments

Comments
 (0)