feat: add templateParameters field for mssqlsql (#671)

Add `templateParameters to support non-filter parameters and DDL
statements.

Part of https://github.com/googleapis/genai-toolbox/issues/535

internal/tools/mssqlsql/mssqlsql.go

@@ -54,13 +54,14 @@ var _ compatibleSource = &mssql.Source{}
 var compatibleSources = [...]string{cloudsqlmssql.SourceKind, mssql.SourceKind}
 
 type Config struct {
-	Name         string           `yaml:"name" validate:"required"`
+	Name               string           `yaml:"name" validate:"required"`
-	Kind         string           `yaml:"kind" validate:"required"`
+	Kind               string           `yaml:"kind" validate:"required"`
-	Source       string           `yaml:"source" validate:"required"`
+	Source             string           `yaml:"source" validate:"required"`
-	Description  string           `yaml:"description" validate:"required"`
+	Description        string           `yaml:"description" validate:"required"`
-	Statement    string           `yaml:"statement" validate:"required"`
+	Statement          string           `yaml:"statement" validate:"required"`
-	AuthRequired []string         `yaml:"authRequired"`
+	AuthRequired       []string         `yaml:"authRequired"`
-	Parameters   tools.Parameters `yaml:"parameters"`
+	Parameters         tools.Parameters `yaml:"parameters"`
+	TemplateParameters tools.Parameters `yaml:"templateParameters"`
 }
 
 // validate interface
@@ -83,22 +84,26 @@ func (cfg Config) Initialize(srcs map[string]sources.Source) (tools.Tool, error)
 		return nil, fmt.Errorf("invalid source for %q tool: source kind must be one of %q", kind, compatibleSources)
 	}
 
+	allParameters, paramManifest, paramMcpManifest := tools.ProcessParameters(cfg.TemplateParameters, cfg.Parameters)
+
 	mcpManifest := tools.McpManifest{
 		Name:        cfg.Name,
 		Description: cfg.Description,
-		InputSchema: cfg.Parameters.McpManifest(),
+		InputSchema: paramMcpManifest,
 	}
 
 	// finish tool setup
 	t := Tool{
-		Name:         cfg.Name,
+		Name:               cfg.Name,
-		Kind:         kind,
+		Kind:               kind,
-		Parameters:   cfg.Parameters,
+		Parameters:         cfg.Parameters,
-		Statement:    cfg.Statement,
+		TemplateParameters: cfg.TemplateParameters,
-		AuthRequired: cfg.AuthRequired,
+		AllParams:          allParameters,
-		Db:           s.MSSQLDB(),
+		Statement:          cfg.Statement,
-		manifest:     tools.Manifest{Description: cfg.Description, Parameters: cfg.Parameters.Manifest(), AuthRequired: cfg.AuthRequired},
+		AuthRequired:       cfg.AuthRequired,
-		mcpManifest:  mcpManifest,
+		Db:                 s.MSSQLDB(),
+		manifest:           tools.Manifest{Description: cfg.Description, Parameters: paramManifest, AuthRequired: cfg.AuthRequired},
+		mcpManifest:        mcpManifest,
 	}
 	return t, nil
 }
@@ -107,10 +112,12 @@ func (cfg Config) Initialize(srcs map[string]sources.Source) (tools.Tool, error)
 var _ tools.Tool = Tool{}
 
 type Tool struct {
-	Name         string           `yaml:"name"`
+	Name               string           `yaml:"name"`
-	Kind         string           `yaml:"kind"`
+	Kind               string           `yaml:"kind"`
-	AuthRequired []string         `yaml:"authRequired"`
+	AuthRequired       []string         `yaml:"authRequired"`
-	Parameters   tools.Parameters `yaml:"parameters"`
+	Parameters         tools.Parameters `yaml:"parameters"`
+	TemplateParameters tools.Parameters `yaml:"templateParameters"`
+	AllParams          tools.Parameters `yaml:"allParams"`
 
 	Db          *sql.DB
 	Statement   string
@@ -119,18 +126,29 @@ type Tool struct {
 }
 
 func (t Tool) Invoke(ctx context.Context, params tools.ParamValues) ([]any, error) {
-	namedArgs := make([]any, 0, len(params))
+	paramsMap := params.AsMap()
-	paramsMap := params.AsReversedMap()
+	newStatement, err := tools.ResolveTemplateParams(t.TemplateParameters, t.Statement, paramsMap)
+	if err != nil {
+		return nil, fmt.Errorf("unable to extract template params %w", err)
+	}
+
+	newParams, err := tools.GetParams(t.Parameters, paramsMap)
+	if err != nil {
+		return nil, fmt.Errorf("unable to extract standard params %w", err)
+	}
+
+	namedArgs := make([]any, 0, len(newParams))
+	newParamsMap := newParams.AsReversedMap()
 	// To support both named args (e.g @id) and positional args (e.g @p1), check if arg name is contained in the statement.
-	for _, v := range params.AsSlice() {
+	for _, v := range newParams.AsSlice() {
-		paramName := paramsMap[v]
+		paramName := newParamsMap[v]
-		if strings.Contains(t.Statement, "@"+paramName) {
+		if strings.Contains(newStatement, "@"+paramName) {
 			namedArgs = append(namedArgs, sql.Named(paramName, v))
 		} else {
 			namedArgs = append(namedArgs, v)
 		}
 	}
-	rows, err := t.Db.QueryContext(ctx, t.Statement, namedArgs...)
+	rows, err := t.Db.QueryContext(ctx, newStatement, namedArgs...)
 	if err != nil {
 		return nil, fmt.Errorf("unable to execute query: %w", err)
 	}
@@ -173,7 +191,7 @@ func (t Tool) Invoke(ctx context.Context, params tools.ParamValues) ([]any, erro
 }
 
 func (t Tool) ParseParams(data map[string]any, claims map[string]map[string]any) (tools.ParamValues, error) {
-	return tools.ParseParams(t.Parameters, data, claims)
+	return tools.ParseParams(t.AllParams, data, claims)
 }
 
 func (t Tool) Manifest() tools.Manifest {

internal/tools/mssqlsql/mssqlsql_test.go

@@ -90,5 +90,86 @@ func TestParseFromYamlMssql(t *testing.T) {
 			}
 		})
 	}
-
+}
+
+func TestParseFromYamlWithTemplateMssql(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: mssql-sql
+					source: my-instance
+					description: some description
+					statement: |
+						SELECT * FROM SQL_STATEMENT;
+					authRequired:
+						- my-google-auth-service
+						- other-auth-service
+					parameters:
+						- name: country
+						  type: string
+						  description: some description
+						  authServices:
+							- name: my-google-auth-service
+							  field: user_id
+							- name: other-auth-service
+							  field: user_id
+					templateParameters:
+						- name: tableName
+						  type: string
+						  description: The table to select hotels from.
+						- name: fieldArray
+						  type: array
+						  description: The columns to return for the query.
+						  items: 
+								name: column
+								type: string
+								description: A column name that will be returned from the query.
+			`,
+			want: server.ToolConfigs{
+				"example_tool": mssqlsql.Config{
+					Name:         "example_tool",
+					Kind:         "mssql-sql",
+					Source:       "my-instance",
+					Description:  "some description",
+					Statement:    "SELECT * FROM SQL_STATEMENT;\n",
+					AuthRequired: []string{"my-google-auth-service", "other-auth-service"},
+					Parameters: []tools.Parameter{
+						tools.NewStringParameterWithAuth("country", "some description",
+							[]tools.ParamAuthService{{Name: "my-google-auth-service", Field: "user_id"},
+								{Name: "other-auth-service", Field: "user_id"}}),
+					},
+					TemplateParameters: []tools.Parameter{
+						tools.NewStringParameter("tableName", "The table to select hotels from."),
+						tools.NewArrayParameter("fieldArray", "The columns to return for the query.", tools.NewStringParameter("column", "A column name that will be returned from the query.")),
+					},
+				},
+			},
+		},
+	}
+	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/cloudsqlmssql/cloud_sql_mssql_integration_test.go

@@ -118,6 +118,7 @@ func TestCloudSQLMssqlToolEndpoints(t *testing.T) {
 	// create table name with UUID
 	tableNameParam := "param_table_" + strings.ReplaceAll(uuid.New().String(), "-", "")
 	tableNameAuth := "auth_table_" + strings.ReplaceAll(uuid.New().String(), "-", "")
+	tableNameTemplateParam := "template_param_table_" + strings.ReplaceAll(uuid.New().String(), "-", "")
 
 	// set up data for param tool
 	create_statement1, insert_statement1, tool_statement1, params1 := tests.GetMssqlParamToolInfo(tableNameParam)
@@ -132,6 +133,8 @@ func TestCloudSQLMssqlToolEndpoints(t *testing.T) {
 	// Write config into a file and pass it to command
 	toolsFile := tests.GetToolsConfig(sourceConfig, CLOUD_SQL_MSSQL_TOOL_KIND, tool_statement1, tool_statement2)
 	toolsFile = tests.AddMssqlExecuteSqlConfig(t, toolsFile)
+	tmplSelectCombined, tmplSelectFilterCombined := tests.GetMssqlTmplToolStatement()
+	toolsFile = tests.AddTemplateParamConfig(t, toolsFile, CLOUD_SQL_MSSQL_TOOL_KIND, tmplSelectCombined, tmplSelectFilterCombined)
 
 	cmd, cleanup, err := tests.StartCmd(ctx, toolsFile, args...)
 	if err != nil {
@@ -154,6 +157,7 @@ func TestCloudSQLMssqlToolEndpoints(t *testing.T) {
 	tests.RunToolInvokeTest(t, select1Want, invokeParamWant)
 	tests.RunExecuteSqlToolInvokeTest(t, createTableStatement, select1Want)
 	tests.RunMCPToolCallMethod(t, mcpInvokeParamWant, failInvocationWant)
+	tests.RunToolInvokeWithTemplateParameters(t, tableNameTemplateParam)
 }
 
 // Test connection with different IP type

tests/common.go

@@ -297,6 +297,13 @@ func GetMssqlAuthToolInfo(tableName string) (string, string, string, []any) {
 	return create_statement, insert_statement, tool_statement, params
 }
 
+// GetMssqlTmplToolStatement returns statements and param for template parameter test cases for mysql-sql kind
+func GetMssqlTmplToolStatement() (string, string) {
+	tmplSelectCombined := "SELECT * FROM {{.tableName}} WHERE id = @id"
+	tmplSelectFilterCombined := "SELECT * FROM {{.tableName}} WHERE {{.columnFilter}} = @name"
+	return tmplSelectCombined, tmplSelectFilterCombined
+}
+
 // GetMysqlParamToolInfo returns statements and param for my-param-tool mysql-sql kind
 func GetMysqlParamToolInfo(tableName string) (string, string, string, []any) {
 	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

@@ -90,6 +90,7 @@ func TestMssqlToolEndpoints(t *testing.T) {
 	// create table name with UUID
 	tableNameParam := "param_table_" + strings.ReplaceAll(uuid.New().String(), "-", "")
 	tableNameAuth := "auth_table_" + strings.ReplaceAll(uuid.New().String(), "-", "")
+	tableNameTemplateParam := "template_param_table_" + strings.ReplaceAll(uuid.New().String(), "-", "")
 
 	// set up data for param tool
 	create_statement1, insert_statement1, tool_statement1, params1 := tests.GetMssqlParamToolInfo(tableNameParam)
@@ -104,6 +105,8 @@ func TestMssqlToolEndpoints(t *testing.T) {
 	// Write config into a file and pass it to command
 	toolsFile := tests.GetToolsConfig(sourceConfig, MSSQL_TOOL_KIND, tool_statement1, tool_statement2)
 	toolsFile = tests.AddMssqlExecuteSqlConfig(t, toolsFile)
+	tmplSelectCombined, tmplSelectFilterCombined := tests.GetMssqlTmplToolStatement()
+	toolsFile = tests.AddTemplateParamConfig(t, toolsFile, MSSQL_TOOL_KIND, tmplSelectCombined, tmplSelectFilterCombined)
 
 	cmd, cleanup, err := tests.StartCmd(ctx, toolsFile, args...)
 	if err != nil {
@@ -126,4 +129,5 @@ func TestMssqlToolEndpoints(t *testing.T) {
 	tests.RunToolInvokeTest(t, select1Want, invokeParamWant)
 	tests.RunExecuteSqlToolInvokeTest(t, createTableStatement, select1Want)
 	tests.RunMCPToolCallMethod(t, mcpInvokeParamWant, failInvocationWant)
+	tests.RunToolInvokeWithTemplateParameters(t, tableNameTemplateParam)
 }