From 093aedb9e9ef0f394a9a0708e22594186f06aa7f Mon Sep 17 00:00:00 2001 From: Yuan Teoh Date: Wed, 29 Jan 2025 13:50:54 -0800 Subject: [PATCH 1/2] fix(mssql): fix mssql tool kind to mssql-sql --- internal/server/config.go | 6 +++--- .../tools/{mssql/mssql.go => mssqlsql/mssqlsql.go} | 4 ++-- .../{mssql/mssql_test.go => mssqlsql/mssqlsql_test.go} | 10 +++++----- tests/cloud_sql_mssql_integration_test.go | 8 ++++---- 4 files changed, 14 insertions(+), 14 deletions(-) rename internal/tools/{mssql/mssql.go => mssqlsql/mssqlsql.go} (98%) rename internal/tools/{mssql/mssql_test.go => mssqlsql/mssqlsql_test.go} (92%) diff --git a/internal/server/config.go b/internal/server/config.go index 16caa940b..0a1e44bdd 100644 --- a/internal/server/config.go +++ b/internal/server/config.go @@ -29,7 +29,7 @@ import ( postgressrc "github.com/googleapis/genai-toolbox/internal/sources/postgres" spannersrc "github.com/googleapis/genai-toolbox/internal/sources/spanner" "github.com/googleapis/genai-toolbox/internal/tools" - "github.com/googleapis/genai-toolbox/internal/tools/mssql" + "github.com/googleapis/genai-toolbox/internal/tools/mssqlsql" "github.com/googleapis/genai-toolbox/internal/tools/mysql" neo4jtool "github.com/googleapis/genai-toolbox/internal/tools/neo4j" "github.com/googleapis/genai-toolbox/internal/tools/postgressql" @@ -272,8 +272,8 @@ func (c *ToolConfigs) UnmarshalYAML(unmarshal func(interface{}) error) error { return fmt.Errorf("unable to parse as %q: %w", k.Kind, err) } (*c)[name] = actual - case mssql.ToolKind: - actual := mssql.Config{Name: name} + case mssqlsql.ToolKind: + actual := mssqlsql.Config{Name: name} if err := u.Unmarshal(&actual); err != nil { return fmt.Errorf("unable to parse as %q: %w", k.Kind, err) } diff --git a/internal/tools/mssql/mssql.go b/internal/tools/mssqlsql/mssqlsql.go similarity index 98% rename from internal/tools/mssql/mssql.go rename to internal/tools/mssqlsql/mssqlsql.go index 2997b361e..0944a41b0 100644 --- a/internal/tools/mssql/mssql.go +++ b/internal/tools/mssqlsql/mssqlsql.go @@ -12,7 +12,7 @@ // See the License for the specific language governing permissions and // limitations under the License. -package mssql +package mssqlsql import ( "database/sql" @@ -24,7 +24,7 @@ import ( "github.com/googleapis/genai-toolbox/internal/tools" ) -const ToolKind string = "mssql" +const ToolKind string = "mssql-sql" type compatibleSource interface { MSSQLDB() *sql.DB diff --git a/internal/tools/mssql/mssql_test.go b/internal/tools/mssqlsql/mssqlsql_test.go similarity index 92% rename from internal/tools/mssql/mssql_test.go rename to internal/tools/mssqlsql/mssqlsql_test.go index d4230b394..e101d5890 100644 --- a/internal/tools/mssql/mssql_test.go +++ b/internal/tools/mssqlsql/mssqlsql_test.go @@ -12,7 +12,7 @@ // See the License for the specific language governing permissions and // limitations under the License. -package mssql_test +package mssqlsql_test import ( "testing" @@ -21,7 +21,7 @@ import ( "github.com/googleapis/genai-toolbox/internal/server" "github.com/googleapis/genai-toolbox/internal/testutils" "github.com/googleapis/genai-toolbox/internal/tools" - "github.com/googleapis/genai-toolbox/internal/tools/mssql" + "github.com/googleapis/genai-toolbox/internal/tools/mssqlsql" "gopkg.in/yaml.v3" ) @@ -36,7 +36,7 @@ func TestParseFromYamlMssql(t *testing.T) { in: ` tools: example_tool: - kind: mssql + kind: mssql-sql source: my-instance description: some description statement: | @@ -55,9 +55,9 @@ func TestParseFromYamlMssql(t *testing.T) { field: user_id `, want: server.ToolConfigs{ - "example_tool": mssql.Config{ + "example_tool": mssqlsql.Config{ Name: "example_tool", - Kind: mssql.ToolKind, + Kind: mssqlsql.ToolKind, Source: "my-instance", Description: "some description", Statement: "SELECT * FROM SQL_STATEMENT;\n", diff --git a/tests/cloud_sql_mssql_integration_test.go b/tests/cloud_sql_mssql_integration_test.go index 2cff4ef44..cea51e023 100644 --- a/tests/cloud_sql_mssql_integration_test.go +++ b/tests/cloud_sql_mssql_integration_test.go @@ -133,7 +133,7 @@ func TestCloudSQLMssql(t *testing.T) { }, "tools": map[string]any{ "my-simple-tool": map[string]any{ - "kind": "mssql", + "kind": "mssql-sql", "source": "my-instance", "description": "Simple tool to test end to end functionality.", "statement": "SELECT 1;", @@ -300,7 +300,7 @@ func TestToolInvocationWithParams(t *testing.T) { defer teardownTest(t) // call generic invocation test helper - RunToolInvocationWithParamsTest(t, sourceConfig, "mssql", tableName) + RunToolInvocationWithParamsTest(t, sourceConfig, "mssql-sql", tableName) } // Set up auth test database table @@ -362,7 +362,7 @@ func TestCloudSQLMssqlGoogleAuthenticatedParameter(t *testing.T) { defer teardownTest(t) // call generic auth test helper - RunGoogleAuthenticatedParameterTest(t, sourceConfig, "mssql", tableName) + RunGoogleAuthenticatedParameterTest(t, sourceConfig, "mssql-sql", tableName) } @@ -371,6 +371,6 @@ func TestCloudSQLMssqlAuthRequiredToolInvocation(t *testing.T) { sourceConfig := requireCloudSQLMssqlVars(t) // call generic auth test helper - RunAuthRequiredToolInvocationTest(t, sourceConfig, "mssql") + RunAuthRequiredToolInvocationTest(t, sourceConfig, "mssql-sql") } From 0871c573350b4977ca8bce1af905525e97d2fccc Mon Sep 17 00:00:00 2001 From: Yuan Teoh Date: Wed, 29 Jan 2025 13:58:07 -0800 Subject: [PATCH 2/2] update auth test --- tests/auth_test.go | 6 +++--- tests/common_test.go | 4 ++-- 2 files changed, 5 insertions(+), 5 deletions(-) diff --git a/tests/auth_test.go b/tests/auth_test.go index 2c4ce7a3a..9c26a8a7a 100644 --- a/tests/auth_test.go +++ b/tests/auth_test.go @@ -63,7 +63,7 @@ func RunGoogleAuthenticatedParameterTest(t *testing.T, sourceConfig map[string]a switch { case strings.EqualFold(toolKind, "postgres-sql"): statement = fmt.Sprintf("SELECT * FROM %s WHERE email = $1;", tableName) - case strings.EqualFold(toolKind, "mssql"): + case strings.EqualFold(toolKind, "mssql-sql"): statement = fmt.Sprintf("SELECT * FROM %s WHERE email = @email;", tableName) default: t.Fatalf("invalid tool kind: %s", toolKind) @@ -132,7 +132,7 @@ func RunGoogleAuthenticatedParameterTest(t *testing.T, sourceConfig map[string]a // Tools using database/sql interface only outputs `int64` instead of `int32` var wantString string switch toolKind { - case "mssql": + case "mssql-sql": wantString = fmt.Sprintf("Stub tool call for \"my-auth-tool\"! Parameters parsed: [{\"email\" \"%s\"}] \n Output: [%%!s(int64=1) Alice %s]", SERVICE_ACCOUNT_EMAIL, SERVICE_ACCOUNT_EMAIL) default: wantString = fmt.Sprintf("Stub tool call for \"my-auth-tool\"! Parameters parsed: [{\"email\" \"%s\"}] \n Output: [%%!s(int32=1) Alice %s]", SERVICE_ACCOUNT_EMAIL, SERVICE_ACCOUNT_EMAIL) @@ -216,7 +216,7 @@ func RunAuthRequiredToolInvocationTest(t *testing.T, sourceConfig map[string]any // Tools using database/sql interface only outputs `int64` instead of `int32` var wantString string switch toolKind { - case "mssql": + case "mssql-sql": wantString = "Stub tool call for \"my-auth-tool\"! Parameters parsed: [] \n Output: [%!s(int64=1)]" default: wantString = "Stub tool call for \"my-auth-tool\"! Parameters parsed: [] \n Output: [%!s(int32=1)]" diff --git a/tests/common_test.go b/tests/common_test.go index ed32264fc..f768481ce 100644 --- a/tests/common_test.go +++ b/tests/common_test.go @@ -215,7 +215,7 @@ func RunToolInvocationWithParamsTest(t *testing.T, sourceConfig map[string]any, switch toolKind { case "postgres-sql": statement = fmt.Sprintf("SELECT * FROM %s WHERE id = $1 OR name = $2;", tableName) - case "mssql": + case "mssql-sql": statement = fmt.Sprintf("SELECT * FROM %s WHERE id = @id OR name = @p2;", tableName) default: t.Fatalf("invalid tool kind: %s", toolKind) @@ -224,7 +224,7 @@ func RunToolInvocationWithParamsTest(t *testing.T, sourceConfig map[string]any, // Tools using database/sql interface only outputs `int64` instead of `int32` var wantString string switch toolKind { - case "mssql": + case "mssql-sql": wantString = "Stub tool call for \"my-tool\"! Parameters parsed: [{\"id\" '\\x03'} {\"name\" \"Alice\"}] \n Output: [%!s(int64=1) Alice][%!s(int64=3) Sid]" default: wantString = "Stub tool call for \"my-tool\"! Parameters parsed: [{\"id\" '\\x03'} {\"name\" \"Alice\"}] \n Output: [%!s(int32=1) Alice][%!s(int32=3) Sid]"