feat(elasticsearch)!: add vector search support and remove query passing through param

.ci/integration.cloudbuild.yaml

@@ -1169,13 +1169,7 @@ steps:
     env:
       - "GOPATH=/gopath"
       - "SERVICE_ACCOUNT_EMAIL=$SERVICE_ACCOUNT_EMAIL"
-    secretEnv:
-      [
-        "CLIENT_ID",
-        "ELASTICSEARCH_USER",
-        "ELASTICSEARCH_PASS",
-        "ELASTICSEARCH_HOST",
-      ]
+    secretEnv: ["CLIENT_ID", "API_KEY"]
     volumes:
       - name: "go"
         path: "/gopath"
@@ -1494,12 +1488,6 @@ availableSecrets:
       env: YUGABYTEDB_USER
     - versionName: projects/$PROJECT_ID/secrets/yugabytedb_pass/versions/latest
       env: YUGABYTEDB_PASS
-    - versionName: projects/$PROJECT_ID/secrets/elastic_search_host/versions/latest
-      env: ELASTICSEARCH_HOST
-    - versionName: projects/$PROJECT_ID/secrets/elastic_search_user/versions/latest
-      env: ELASTICSEARCH_USER
-    - versionName: projects/$PROJECT_ID/secrets/elastic_search_pass/versions/latest
-      env: ELASTICSEARCH_PASS
     - versionName: projects/$PROJECT_ID/secrets/snowflake_account/versions/latest
       env: SNOWFLAKE_ACCOUNT
     - versionName: projects/$PROJECT_ID/secrets/snowflake_user/versions/latest
@@ -1593,4 +1581,4 @@ substitutions:
   _MARIADB_PORT: "3307"
   _MARIADB_DATABASE: test_database
   _SNOWFLAKE_DATABASE: "test"
-  _SNOWFLAKE_SCHEMA: "PUBLIC"
+  _SNOWFLAKE_SCHEMA: "PUBLIC"

internal/tools/elasticsearch/elasticsearchesql/elasticsearchesql.go

@@ -49,7 +49,7 @@ type Config struct {
 	Source       string                 `yaml:"source" validate:"required"`
 	Description  string                 `yaml:"description" validate:"required"`
 	AuthRequired []string               `yaml:"authRequired" validate:"required"`
-	Query        string                 `yaml:"query"`
+	Query        string                 `yaml:"query" validate:"required"`
 	Format       string                 `yaml:"format"`
 	Timeout      int                    `yaml:"timeout"`
 	Parameters   parameters.Parameters  `yaml:"parameters"`
@@ -109,25 +109,21 @@ func (t Tool) Invoke(ctx context.Context, resourceMgr tools.SourceProvider, para
 	}
 
 	query := t.Query
-	sqlParams := make([]map[string]any, 0, len(params))
 	paramMap := params.AsMap()
-	// If a query is provided in the params and not already set in the tool, use it.
-	if queryVal, ok := paramMap["query"]; ok {
-		if str, ok := queryVal.(string); ok && t.Query == "" {
-			query = str
-		}
-
-		// Drop the query param if not a string or if the tool already has a query.
-		delete(paramMap, "query")
-	}
 
+	var paramsList []map[string]any
 	for _, param := range t.Parameters {
 		if param.GetType() == "array" {
 			return nil, util.NewAgentError("array parameters are not supported yet", nil)
 		}
-		sqlParams = append(sqlParams, map[string]any{param.GetName(): paramMap[param.GetName()]})
+
+		// ES|QL requires an array of single-key objects for named parameters
+		if val, ok := paramMap[param.GetName()]; ok {
+			paramsList = append(paramsList, map[string]any{param.GetName(): val})
+		}
 	}
-	resp, err := source.RunSQL(ctx, t.Format, query, sqlParams)
+
+	resp, err := source.RunSQL(ctx, t.Format, query, paramsList)
 	if err != nil {
 		return nil, util.ProcessGeneralError(err)
 	}

tests/elasticsearch/elasticsearch_integration_test.go

@@ -20,14 +20,15 @@ import (
 	"encoding/json"
 	"fmt"
 	"net/http"
-	"os"
 	"regexp"
 	"strings"
 	"testing"
 	"time"
 
 	"github.com/elastic/go-elasticsearch/v9"
 	"github.com/elastic/go-elasticsearch/v9/esapi"
+	"github.com/testcontainers/testcontainers-go"
+	"github.com/testcontainers/testcontainers-go/wait"
 
 	"github.com/googleapis/mcp-toolbox/internal/testutils"
 	"github.com/googleapis/mcp-toolbox/tests"
@@ -36,15 +37,57 @@ import (
 var (
 	ElasticsearchSourceType = "elasticsearch"
 	ElasticsearchToolType   = "elasticsearch-esql"
-	EsAddress               = os.Getenv("ELASTICSEARCH_HOST")
-	EsUser                  = os.Getenv("ELASTICSEARCH_USER")
-	EsPass                  = os.Getenv("ELASTICSEARCH_PASS")
+	EsAddress               = ""
+	EsUser                  = "elastic"
+	EsPass                  = "test-password"
 )
 
-func getElasticsearchVars(t *testing.T) map[string]any {
-	if EsAddress == "" {
-		t.Fatal("'ELASTICSEARCH_HOST' not set")
+func setupElasticsearchContainer(ctx context.Context, t *testing.T) (string, func()) {
+	t.Helper()
+
+	req := testcontainers.ContainerRequest{
+		Image:        "docker.elastic.co/elasticsearch/elasticsearch:9.3.2",
+		ExposedPorts: []string{"9200/tcp"},
+		Env: map[string]string{
+			"discovery.type":         "single-node",
+			"xpack.security.enabled": "false",
+		},
+		WaitingFor: wait.ForAll(
+			wait.ForHTTP("/"),
+			wait.ForExposedPort(),
+		),
 	}
+
+	container, err := testcontainers.GenericContainer(ctx, testcontainers.GenericContainerRequest{
+		ContainerRequest: req,
+		Started:          true,
+	})
+	if err != nil {
+		t.Fatalf("failed to start elasticsearch container: %s", err)
+	}
+
+	cleanup := func() {
+		if err := container.Terminate(ctx); err != nil {
+			t.Fatalf("failed to terminate container: %s", err)
+		}
+	}
+
+	host, err := container.Host(ctx)
+	if err != nil {
+		cleanup()
+		t.Fatalf("failed to get container host: %s", err)
+	}
+
+	mappedPort, err := container.MappedPort(ctx, "9200")
+	if err != nil {
+		cleanup()
+		t.Fatalf("failed to get container mapped port: %s", err)
+	}
+
+	return fmt.Sprintf("http://%s:%s", host, mappedPort.Port()), cleanup
+}
+
+func getElasticsearchVars(t *testing.T) map[string]any {
 	return map[string]any{
 		"type":      ElasticsearchSourceType,
 		"addresses": []string{EsAddress},
@@ -64,9 +107,13 @@ type ElasticsearchWants struct {
 }
 
 func TestElasticsearchToolEndpoints(t *testing.T) {
-	ctx, cancel := context.WithTimeout(context.Background(), time.Minute)
+	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Minute)
 	defer cancel()
 
+	var containerCleanup func()
+	EsAddress, containerCleanup = setupElasticsearchContainer(ctx, t)
+	defer containerCleanup()
+
 	args := []string{"--enable-api"}
 
 	sourceConfig := getElasticsearchVars(t)
@@ -77,6 +124,10 @@ func TestElasticsearchToolEndpoints(t *testing.T) {
 
 	toolsConfig := getElasticsearchToolsConfig(sourceConfig, ElasticsearchToolType, paramToolStatement, idParamToolStatement, nameParamToolStatement, arrayParamToolStatement, authToolStatement)
 
+	searchStmt := fmt.Sprintf(`FROM %s | WHERE KNN(embedding, ?query) | LIMIT 1 | KEEP id, name`, index)
+	insertStmt := fmt.Sprintf("FROM %s | WHERE name == ?content | EVAL dummy = ?text_to_embed | LIMIT 0", index)
+	toolsConfig = tests.AddSemanticSearchConfig(t, toolsConfig, ElasticsearchToolType, insertStmt, searchStmt)
+
 	cmd, cleanup, err := tests.StartCmd(ctx, toolsConfig, args...)
 	if err != nil {
 		t.Fatalf("failed to start cmd: %v", err)
@@ -100,28 +151,65 @@ func TestElasticsearchToolEndpoints(t *testing.T) {
 		t.Fatalf("error creating the Elasticsearch client: %s", err)
 	}
 
-	// Delete index if already exists
+	// Delete indices if already exists
 	defer func() {
 		_, err = esapi.IndicesDeleteRequest{
 			Index: []string{index},
 		}.Do(ctx, esClient)
 		if err != nil {
-			t.Fatalf("error deleting index: %s", err)
+			t.Errorf("error deleting indices: %s", err)
 		}
 	}()
 
 	alice := fmt.Sprintf(`{
-									"id": 1,
-									"name": "Alice",
-									"email": "%s"
-								}`, tests.ServiceAccountEmail)
+                  "id": 1,
+                  "name": "Alice",
+                  "email": "%s"
+                }`, tests.ServiceAccountEmail)
+
+	// Create index with mapping for vector search
+	mapping := `{
+    "mappings": {
+      "properties": {
+        "embedding": {
+          "type": "dense_vector",
+          "dims": 768,
+          "index": true,
+          "similarity": "cosine"
+        }
+      }
+    }
+  }`
+	res, err := esapi.IndicesCreateRequest{
+		Index: index,
+		Body:  strings.NewReader(mapping),
+	}.Do(ctx, esClient)
+	if err != nil {
+		t.Fatalf("error creating index: %s", err)
+	}
+	if res.IsError() {
+		t.Logf("Create index response error (might be ignored): %s", res.String())
+	}
+
+	vectorSize := 768
+	var sb strings.Builder
+	sb.WriteString("[")
+	for i := 0; i < vectorSize; i++ {
+		sb.WriteString("0.1")
+		if i < vectorSize-1 {
+			sb.WriteString(", ")
+		}
+	}
+	sb.WriteString("]")
+	semanticDoc := fmt.Sprintf(`{"id": 5, "name": "Semantic", "embedding": %s}`, sb.String())
 
 	// Index sample documents
 	sampleDocs := []string{
 		alice,
 		`{"id": 2, "name": "Jane", "email": "janedoe@gmail.com"}`,
 		`{"id": 3, "name": "Sid"}`,
 		`{"id": 4, "name": "null"}`,
+		semanticDoc,
 	}
 	for _, doc := range sampleDocs {
 		res, err := esapi.IndexRequest{
@@ -150,27 +238,28 @@ func TestElasticsearchToolEndpoints(t *testing.T) {
 	)
 	tests.RunMCPToolCallMethod(t, wants.McpMyFailTool, wants.McpSelect1, tests.WithMcpMyToolId3NameAliceWant(wants.McpMyToolId3NameAlice))
 
+	// Semantic search tests
+	semanticSearchWant := `[{"id":5,"name":"Semantic"}]`
+	tests.RunSemanticSearchToolInvokeTest(t, "[]", "[]", semanticSearchWant)
 	runExecuteEsqlTest(t, index)
-
 }
 
 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.
+	paramToolStatement := fmt.Sprintf(`FROM %s | WHERE id == ?id OR name == ?name | SORT id ASC | KEEP id, name, name.keyword, email, email.keyword`, index)
+	idParamToolStatement := fmt.Sprintf(`FROM %s | WHERE id == ?id | KEEP id, name, name.keyword, email, email.keyword`, index)
+	nameParamToolStatement := fmt.Sprintf(`FROM %s | WHERE name == ?name | KEEP id, name, name.keyword, email, email.keyword`, index)
 	authToolStatement := fmt.Sprintf(`FROM %s | WHERE email == ?email | KEEP name`, index)
-	return paramToolStatement, idParamToolStatement, nameParamToolStatement, arrayParamToolStatement, authToolStatement
+	return paramToolStatement, idParamToolStatement, nameParamToolStatement, "", 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)
+	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"},{"email":null,"email.keyword":null,"id":5,"name":"Semantic","name.keyword":"Semantic"}]`, 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}"}]}`
+	mcpMyFailToolWant := `{"content":[{"type":"text","text":"{\"error\":{\"root_cause\":[{\"type\":\"parsing_exception\",\"reason\":\"line 1:1: mismatched input 'SELEC' expecting {, 'row', 'from', 'ts', 'set', 'show'}\"}],\"type\":\"parsing_exception\",\"reason\":\"line 1:1: mismatched input 'SELEC' expecting {, 'row', 'from', 'ts', 'set', '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)
+	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\"},{\"email\":null,\"email.keyword\":null,\"id\":5,\"name\":\"Semantic\",\"name.keyword\":\"Semantic\"}]"}]}}`, tests.ServiceAccountEmail)
 
 	return ElasticsearchWants{
 		Select1:               select1Want,
@@ -199,7 +288,7 @@ func getElasticsearchToolsConfig(sourceConfig map[string]any, toolType, paramToo
 				"type":        toolType,
 				"source":      "my-instance",
 				"description": "Simple tool to test end to end functionality.",
-				"query":       "FROM test-index | SORT id ASC",
+				"query":       "FROM test-index | SORT id ASC | KEEP id, name, name.keyword, email, email.keyword",
 			},
 			"my-tool": map[string]any{
 				"type":        toolType,
@@ -246,34 +335,6 @@ func getElasticsearchToolsConfig(sourceConfig map[string]any, toolType, paramToo
 					},
 				},
 			},
-			"my-array-tool": map[string]any{
-				"type":        toolType,
-				"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{
 				"type":        toolType,
 				"source":      "my-instance",
@@ -298,7 +359,7 @@ func getElasticsearchToolsConfig(sourceConfig map[string]any, toolType, paramToo
 				"type":        toolType,
 				"source":      "my-instance",
 				"description": "Tool to test auth required invocation.",
-				"query":       "FROM test-index | SORT id ASC",
+				"query":       "FROM test-index | SORT id ASC | KEEP id, name, name.keyword, email, email.keyword",
 				"authRequired": []string{
 					"my-google-auth",
 				},
@@ -339,7 +400,7 @@ func runExecuteEsqlTest(t *testing.T, index string) {
 		if !ok {
 			t.Fatalf("unable to find result in response body")
 		}
-		want := `[{"id":1},{"id":2},{"id":3},{"id":4}]`
+		want := `[{"id":1},{"id":2},{"id":3},{"id":4},{"id":5}]`
 		if got != want {
 			t.Fatalf("unexpected value: got %q, want %q", got, want)
 		}

tests/embedding.go

@@ -52,11 +52,16 @@ func AddSemanticSearchConfig(t *testing.T, config map[string]any, toolKind, inse
 		t.Fatalf("unable to get tools from config")
 	}
 
+	queryKey := "statement"
+	if toolKind == "elasticsearch-esql" {
+		queryKey = "query"
+	}
+
 	tools["insert_docs"] = map[string]any{
 		"kind":        toolKind,
 		"source":      "my-instance",
 		"description": "Stores content and its vector embedding into the documents table.",
-		"statement":   insertStmt,
+		queryKey:      insertStmt,
 		"parameters": []any{
 			map[string]any{
 				"name":        "content",
@@ -77,7 +82,7 @@ func AddSemanticSearchConfig(t *testing.T, config map[string]any, toolKind, inse
 		"kind":        toolKind,
 		"source":      "my-instance",
 		"description": "Finds the most semantically similar document to the query vector.",
-		"statement":   searchStmt,
+		queryKey:      searchStmt,
 		"parameters": []any{
 			map[string]any{
 				"name":        "query",