MCP Toolbox for Databases

// 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 elasticsearch import ( "context" "fmt" "os" "regexp" "strings" "testing" "time" "github.com/elastic/go-elasticsearch/v9" "github.com/elastic/go-elasticsearch/v9/esapi" "github.com/googleapis/genai-toolbox/internal/testutils" "github.com/googleapis/genai-toolbox/tests" ) var ( ElasticsearchSourceKind = "elasticsearch" ElasticsearchToolKind = "elasticsearch-esql" EsAddress = os.Getenv("ELASTICSEARCH_HOST") EsUser = os.Getenv("ELASTICSEARCH_USER") EsPass = os.Getenv("ELASTICSEARCH_PASS") ) func getElasticsearchVars(t *testing.T) map[string]any { if EsAddress == "" { t.Fatal("'ELASTICSEARCH_HOST' not set") } return map[string]any{ "kind": ElasticsearchSourceKind, "addresses": []string{EsAddress}, "username": EsUser, "password": EsPass, } } type ElasticsearchWants struct { Select1 string MyToolId3NameAlice string MyToolById4 string Null string McpMyFailTool string McpMyToolId3NameAlice string McpSelect1 string } func TestElasticsearchToolEndpoints(t *testing.T) { ctx, cancel := context.WithTimeout(context.Background(), time.Minute) defer cancel() var args []string sourceConfig := getElasticsearchVars(t) index := "test-index" paramToolStatement, idParamToolStatement, nameParamToolStatement, arrayParamToolStatement, authToolStatement := getElasticsearchQueries(index) toolsConfig := getElasticsearchToolsConfig(sourceConfig, ElasticsearchToolKind, paramToolStatement, idParamToolStatement, nameParamToolStatement, arrayParamToolStatement, authToolStatement) cmd, cleanup, err := tests.StartCmd(ctx, toolsConfig, args...) if err != nil { t.Fatalf("failed to start cmd: %v", err) } defer cleanup() waitCtx, cancel := context.WithTimeout(ctx, 10*time.Second) defer cancel() out, err := testutils.WaitForString(waitCtx, regexp.MustCompile(`Server ready to serve`), cmd.Out) if err != nil { t.Logf("toolbox command logs: \n%s", out) t.Fatalf("toolbox didn't start successfully: %s", err) } esClient, err := elasticsearch.NewBaseClient(elasticsearch.Config{ Addresses: []string{EsAddress}, Username: EsUser, Password: EsPass, }) if err != nil { t.Fatalf("error creating the Elasticsearch client: %s", err) } // Delete index if already exists defer func() { _, err = esapi.IndicesDeleteRequest{ Index: []string{index}, }.Do(ctx, esClient) if err != nil { t.Fatalf("error deleting index: %s", err) } }() alice := fmt.Sprintf(`{ "id": 1, "name": "Alice", "email": "%s" }`, tests.ServiceAccountEmail) // Index sample documents sampleDocs := []string{ alice, `{"id": 2, "name": "Jane", "email": "janedoe@gmail.com"}`, `{"id": 3, "name": "Sid"}`, `{"id": 4, "name": "null"}`, } for _, doc := range sampleDocs { res, err := esapi.IndexRequest{ Index: "test-index", Body: strings.NewReader(doc), Refresh: "true", }.Do(ctx, esClient) if res.IsError() { t.Fatalf("error indexing document: %s", res.String()) } if err != nil { t.Fatalf("error indexing document: %s", err) } } // Get configs for tests wants := getElasticsearchWants() tests.RunToolGetTest(t) tests.RunToolInvokeTest(t, wants.Select1, tests.DisableArrayTest(), tests.WithMyToolId3NameAliceWant(wants.MyToolId3NameAlice), tests.WithMyToolById4Want(wants.MyToolById4), tests.WithNullWant(wants.Null), ) tests.RunMCPToolCallMethod(t, wants.McpMyFailTool, wants.McpSelect1, tests.WithMcpMyToolId3NameAliceWant(wants.McpMyToolId3NameAlice)) } func getElasticsearchQueries(index string) (string, string, string, string, string) { paramToolStatement := fmt.Sprintf(`FROM %s | WHERE id == ?id OR name == ?name | SORT id ASC`, index) idParamToolStatement := fmt.Sprintf(`FROM %s | WHERE id == ?id`, index) nameParamToolStatement := fmt.Sprintf(`FROM %s | WHERE name == ?name`, index) arrayParamToolStatement := fmt.Sprintf(`FROM %s | WHERE first_name == ?first_name_array`, index) // Not supported yet. authToolStatement := fmt.Sprintf(`FROM %s | WHERE email == ?email | KEEP name`, index) return paramToolStatement, idParamToolStatement, nameParamToolStatement, arrayParamToolStatement, authToolStatement } func getElasticsearchWants() ElasticsearchWants { select1Want := fmt.Sprintf(`[{"email":"%[1]s","email.keyword":"%[1]s","id":1,"name":"Alice","name.keyword":"Alice"},{"email":"janedoe@gmail.com","email.keyword":"janedoe@gmail.com","id":2,"name":"Jane","name.keyword":"Jane"},{"email":null,"email.keyword":null,"id":3,"name":"Sid","name.keyword":"Sid"},{"email":null,"email.keyword":null,"id":4,"name":"null","name.keyword":"null"}]`, tests.ServiceAccountEmail) myToolId3NameAliceWant := fmt.Sprintf(`[{"email":"%[1]s","email.keyword":"%[1]s","id":1,"name":"Alice","name.keyword":"Alice"},{"email":null,"email.keyword":null,"id":3,"name":"Sid","name.keyword":"Sid"}]`, tests.ServiceAccountEmail) myToolById4Want := `[{"email":null,"email.keyword":null,"id":4,"name":"null","name.keyword":"null"}]` nullWant := `{"error":{"root_cause":[{"type":"verification_exception","reason":"Found 1 problem\nline 1:25: first argument of [name == ?name] is [text] so second argument must also be [text] but was [null]"}],"type":"verification_exception","reason":"Found 1 problem\nline 1:25: first argument of [name == ?name] is [text] so second argument must also be [text] but was [null]"},"status":400}` mcpMyFailToolWant := `{"content":[{"type":"text","text":"{\"error\":{\"root_cause\":[{\"type\":\"parsing_exception\",\"reason\":\"line 1:1: mismatched input 'SELEC' expecting {, 'row', 'from', 'show'}\"}],\"type\":\"parsing_exception\",\"reason\":\"line 1:1: mismatched input 'SELEC' expecting {, 'row', 'from', 'show'}\",\"caused_by\":{\"type\":\"input_mismatch_exception\",\"reason\":null}},\"status\":400}"}]}` mcpMyToolId3NameAliceWant := fmt.Sprintf(`{"jsonrpc":"2.0","id":"my-tool","result":{"content":[{"type":"text","text":"[{\"email\":\"%[1]s\",\"email.keyword\":\"%[1]s\",\"id\":1,\"name\":\"Alice\",\"name.keyword\":\"Alice\"},{\"email\":null,\"email.keyword\":null,\"id\":3,\"name\":\"Sid\",\"name.keyword\":\"Sid\"}]"}]}}`, tests.ServiceAccountEmail) mcpSelect1Want := fmt.Sprintf(`{"jsonrpc":"2.0","id":"invoke my-auth-required-tool","result":{"content":[{"type":"text","text":"[{\"email\":\"%[1]s\",\"email.keyword\":\"%[1]s\",\"id\":1,\"name\":\"Alice\",\"name.keyword\":\"Alice\"},{\"email\":\"janedoe@gmail.com\",\"email.keyword\":\"janedoe@gmail.com\",\"id\":2,\"name\":\"Jane\",\"name.keyword\":\"Jane\"},{\"email\":null,\"email.keyword\":null,\"id\":3,\"name\":\"Sid\",\"name.keyword\":\"Sid\"},{\"email\":null,\"email.keyword\":null,\"id\":4,\"name\":\"null\",\"name.keyword\":\"null\"}]"}]}}`, tests.ServiceAccountEmail) return ElasticsearchWants{ Select1: select1Want, MyToolId3NameAlice: myToolId3NameAliceWant, MyToolById4: myToolById4Want, Null: nullWant, McpMyFailTool: mcpMyFailToolWant, McpMyToolId3NameAlice: mcpMyToolId3NameAliceWant, McpSelect1: mcpSelect1Want, } } func getElasticsearchToolsConfig(sourceConfig map[string]any, toolKind, paramToolStatement, idParamToolStmt, nameParamToolStmt, arrayToolStatement, authToolStatement string) map[string]any { toolsFile := map[string]any{ "sources": map[string]any{ "my-instance": sourceConfig, }, "authServices": map[string]any{ "my-google-auth": map[string]any{ "kind": "google", "clientId": tests.ClientId, }, }, "tools": map[string]any{ "my-simple-tool": map[string]any{ "kind": toolKind, "source": "my-instance", "description": "Simple tool to test end to end functionality.", "query": "FROM test-index | SORT id ASC", }, "my-tool": map[string]any{ "kind": toolKind, "source": "my-instance", "description": "Tool to test invocation with params.", "query": paramToolStatement, "parameters": []any{ map[string]any{ "name": "id", "type": "integer", "description": "user ID", }, map[string]any{ "name": "name", "type": "string", "description": "user name", }, }, }, "my-tool-by-id": map[string]any{ "kind": toolKind, "source": "my-instance", "description": "Tool to test invocation with params.", "query": idParamToolStmt, "parameters": []any{ map[string]any{ "name": "id", "type": "integer", "description": "user ID", }, }, }, "my-tool-by-name": map[string]any{ "kind": toolKind, "source": "my-instance", "description": "Tool to test invocation with params.", "query": nameParamToolStmt, "parameters": []any{ map[string]any{ "name": "name", "type": "string", "description": "user name", "required": false, }, }, }, "my-array-tool": map[string]any{ "kind": toolKind, "source": "my-instance", "description": "Tool to test invocation with array params.", "query": arrayToolStatement, "parameters": []any{ map[string]any{ "name": "idArray", "type": "array", "description": "ID array", "items": map[string]any{ "name": "id", "type": "integer", "description": "ID", }, }, map[string]any{ "name": "nameArray", "type": "array", "description": "user name array", "items": map[string]any{ "name": "name", "type": "string", "description": "user name", }, }, }, }, "my-auth-tool": map[string]any{ "kind": toolKind, "source": "my-instance", "description": "Tool to test authenticated parameters.", // statement to auto-fill authenticated parameter "query": authToolStatement, "parameters": []map[string]any{ { "name": "email", "type": "string", "description": "user email", "authServices": []map[string]string{ { "name": "my-google-auth", "field": "email", }, }, }, }, }, "my-auth-required-tool": map[string]any{ "kind": toolKind, "source": "my-instance", "description": "Tool to test auth required invocation.", "query": "FROM test-index | SORT id ASC", "authRequired": []string{ "my-google-auth", }, }, "my-fail-tool": map[string]any{ "kind": toolKind, "source": "my-instance", "description": "Tool to test statement with incorrect syntax.", "query": "SELEC 1;", }, }, } return toolsFile }