mirror of
https://github.com/googleapis/genai-toolbox.git
synced 2026-01-10 07:58:12 -05:00
feat(sources/elasticsearch): add Elasticsearch source and tools (#1109)
Add support for Elasticsearch with the following tools: * search * esql * get_mappings * list_indices This PR fixes #859 --------- Co-authored-by: duwenxin <duwenxin@google.com> Co-authored-by: Wenxin Du <117315983+duwenxin99@users.noreply.github.com>
This commit is contained in:
Laurent Saint-Félix
committed by
GitHub
GitHub
parent
d7f68ebb1a
commit
5367285e91
@@ -705,7 +705,26 @@ steps:
|
|||||||
- |
|
- |
|
||||||
./yugabytedb.test -test.v
|
./yugabytedb.test -test.v
|
||||||
|
|
||||||
|
- id: "elasticsearch"
|
||||||
|
name: golang:1
|
||||||
|
waitFor: ["compile-test-binary"]
|
||||||
|
entrypoint: /bin/bash
|
||||||
|
env:
|
||||||
|
- "GOPATH=/gopath"
|
||||||
|
- "SERVICE_ACCOUNT_EMAIL=$SERVICE_ACCOUNT_EMAIL"
|
||||||
|
secretEnv: ["CLIENT_ID", "ELASTICSEARCH_USER", "ELASTICSEARCH_PASS", "ELASTICSEARCH_HOST"]
|
||||||
|
volumes:
|
||||||
|
- name: "go"
|
||||||
|
path: "/gopath"
|
||||||
|
args:
|
||||||
|
- -c
|
||||||
|
- |
|
||||||
|
.ci/test_with_coverage.sh \
|
||||||
|
"Elasticsearch" \
|
||||||
|
elasticsearch \
|
||||||
|
elasticsearch
|
||||||
|
|
||||||
|
|
||||||
- id: "cassandra"
|
- id: "cassandra"
|
||||||
name: golang:1
|
name: golang:1
|
||||||
waitFor: ["compile-test-binary"]
|
waitFor: ["compile-test-binary"]
|
||||||
@@ -764,7 +783,7 @@ steps:
|
|||||||
.ci/test_with_coverage.sh \
|
.ci/test_with_coverage.sh \
|
||||||
"Serverless Spark" \
|
"Serverless Spark" \
|
||||||
serverlessspark
|
serverlessspark
|
||||||
|
|
||||||
availableSecrets:
|
availableSecrets:
|
||||||
secretManager:
|
secretManager:
|
||||||
- versionName: projects/$PROJECT_ID/secrets/cloud_sql_pg_user/versions/latest
|
- versionName: projects/$PROJECT_ID/secrets/cloud_sql_pg_user/versions/latest
|
||||||
@@ -855,6 +874,12 @@ availableSecrets:
|
|||||||
env: YUGABYTEDB_USER
|
env: YUGABYTEDB_USER
|
||||||
- versionName: projects/$PROJECT_ID/secrets/yugabytedb_pass/versions/latest
|
- versionName: projects/$PROJECT_ID/secrets/yugabytedb_pass/versions/latest
|
||||||
env: YUGABYTEDB_PASS
|
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/cassandra_user/versions/latest
|
- versionName: projects/$PROJECT_ID/secrets/cassandra_user/versions/latest
|
||||||
env: CASSANDRA_USER
|
env: CASSANDRA_USER
|
||||||
- versionName: projects/$PROJECT_ID/secrets/cassandra_pass/versions/latest
|
- versionName: projects/$PROJECT_ID/secrets/cassandra_pass/versions/latest
|
||||||
|
|||||||
@@ -85,6 +85,7 @@ import (
|
|||||||
_ "github.com/googleapis/genai-toolbox/internal/tools/dataplex/dataplexsearchaspecttypes"
|
_ "github.com/googleapis/genai-toolbox/internal/tools/dataplex/dataplexsearchaspecttypes"
|
||||||
_ "github.com/googleapis/genai-toolbox/internal/tools/dataplex/dataplexsearchentries"
|
_ "github.com/googleapis/genai-toolbox/internal/tools/dataplex/dataplexsearchentries"
|
||||||
_ "github.com/googleapis/genai-toolbox/internal/tools/dgraph"
|
_ "github.com/googleapis/genai-toolbox/internal/tools/dgraph"
|
||||||
|
_ "github.com/googleapis/genai-toolbox/internal/tools/elasticsearch/elasticsearchesql"
|
||||||
_ "github.com/googleapis/genai-toolbox/internal/tools/firebird/firebirdexecutesql"
|
_ "github.com/googleapis/genai-toolbox/internal/tools/firebird/firebirdexecutesql"
|
||||||
_ "github.com/googleapis/genai-toolbox/internal/tools/firebird/firebirdsql"
|
_ "github.com/googleapis/genai-toolbox/internal/tools/firebird/firebirdsql"
|
||||||
_ "github.com/googleapis/genai-toolbox/internal/tools/firestore/firestoreadddocuments"
|
_ "github.com/googleapis/genai-toolbox/internal/tools/firestore/firestoreadddocuments"
|
||||||
@@ -195,6 +196,7 @@ import (
|
|||||||
_ "github.com/googleapis/genai-toolbox/internal/sources/couchbase"
|
_ "github.com/googleapis/genai-toolbox/internal/sources/couchbase"
|
||||||
_ "github.com/googleapis/genai-toolbox/internal/sources/dataplex"
|
_ "github.com/googleapis/genai-toolbox/internal/sources/dataplex"
|
||||||
_ "github.com/googleapis/genai-toolbox/internal/sources/dgraph"
|
_ "github.com/googleapis/genai-toolbox/internal/sources/dgraph"
|
||||||
|
_ "github.com/googleapis/genai-toolbox/internal/sources/elasticsearch"
|
||||||
_ "github.com/googleapis/genai-toolbox/internal/sources/firebird"
|
_ "github.com/googleapis/genai-toolbox/internal/sources/firebird"
|
||||||
_ "github.com/googleapis/genai-toolbox/internal/sources/firestore"
|
_ "github.com/googleapis/genai-toolbox/internal/sources/firestore"
|
||||||
_ "github.com/googleapis/genai-toolbox/internal/sources/http"
|
_ "github.com/googleapis/genai-toolbox/internal/sources/http"
|
||||||
|
|||||||
68
docs/en/resources/sources/elasticsearch.md
Normal file
68
docs/en/resources/sources/elasticsearch.md
Normal file
@@ -0,0 +1,68 @@
|
|||||||
|
---
|
||||||
|
title: "Elasticsearch"
|
||||||
|
type: docs
|
||||||
|
weight: 1
|
||||||
|
description: >
|
||||||
|
Elasticsearch is a distributed, free and open search and analytics engine
|
||||||
|
for all types of data, including textual, numerical, geospatial, structured,
|
||||||
|
and unstructured.
|
||||||
|
---
|
||||||
|
|
||||||
|
# Elasticsearch Source
|
||||||
|
|
||||||
|
[Elasticsearch][elasticsearch-docs] is a distributed, free and open search and analytics engine
|
||||||
|
for all types of data, including textual, numerical, geospatial, structured,
|
||||||
|
and unstructured.
|
||||||
|
|
||||||
|
If you are new to Elasticsearch, you can learn how to
|
||||||
|
[set up a cluster and start indexing data][elasticsearch-quickstart].
|
||||||
|
|
||||||
|
Elasticsearch uses [ES|QL][elasticsearch-esql] for querying data. ES|QL
|
||||||
|
is a powerful query language that allows you to search and aggregate data in
|
||||||
|
Elasticsearch.
|
||||||
|
|
||||||
|
See the [official documentation](https://www.elastic.co/guide/en/elasticsearch/reference/current/index.html) for more information.
|
||||||
|
|
||||||
|
[elasticsearch-docs]: https://www.elastic.co/guide/en/elasticsearch/reference/current/index.html
|
||||||
|
[elasticsearch-quickstart]: https://www.elastic.co/guide/en/elasticsearch/reference/current/getting-started.html
|
||||||
|
[elasticsearch-esql]: https://www.elastic.co/guide/en/elasticsearch/reference/current/esql.html
|
||||||
|
|
||||||
|
## Available Tools
|
||||||
|
|
||||||
|
- [`elasticsearch-esql`](../tools/elasticsearch/elasticsearch-esql.md)
|
||||||
|
Execute ES|QL queries.
|
||||||
|
|
||||||
|
## Requirements
|
||||||
|
|
||||||
|
### API Key
|
||||||
|
|
||||||
|
Toolbox uses an [API key][api-key] to authorize and authenticate when
|
||||||
|
interacting with [Elasticsearch][elasticsearch-docs].
|
||||||
|
|
||||||
|
In addition to [setting the API key for your server][set-api-key], you need to
|
||||||
|
ensure the API key has the correct permissions for the queries you intend to
|
||||||
|
run. See [API key management][api-key-management] for more information on
|
||||||
|
applying permissions to an API key.
|
||||||
|
|
||||||
|
[api-key]: https://www.elastic.co/guide/en/elasticsearch/reference/current/security-api-create-api-key.html
|
||||||
|
[set-api-key]: https://www.elastic.co/guide/en/elasticsearch/reference/current/security-api-create-api-key.html
|
||||||
|
[api-key-management]: https://www.elastic.co/guide/en/elasticsearch/reference/current/security-api-get-api-key.html
|
||||||
|
|
||||||
|
## Example
|
||||||
|
|
||||||
|
```yaml
|
||||||
|
sources:
|
||||||
|
my-elasticsearch-source:
|
||||||
|
kind: "elasticsearch"
|
||||||
|
addresses:
|
||||||
|
- "http://localhost:9200"
|
||||||
|
apikey: "my-api-key"
|
||||||
|
```
|
||||||
|
|
||||||
|
## Reference
|
||||||
|
|
||||||
|
| **field** | **type** | **required** | **description** |
|
||||||
|
|-----------|:--------:|:------------:|-------------------------------------------------------------------------------|
|
||||||
|
| kind | string | true | Must be "elasticsearch". |
|
||||||
|
| addresses | []string | true | List of Elasticsearch hosts to connect to. |
|
||||||
|
| apikey | string | true | The API key to use for authentication. |
|
||||||
7
docs/en/resources/tools/elasticsearch/_index.md
Normal file
7
docs/en/resources/tools/elasticsearch/_index.md
Normal file
@@ -0,0 +1,7 @@
|
|||||||
|
---
|
||||||
|
title: "Elasticsearch"
|
||||||
|
type: docs
|
||||||
|
weight: 1
|
||||||
|
description: >
|
||||||
|
Tools that work with Elasticsearch Sources.
|
||||||
|
---
|
||||||
45
docs/en/resources/tools/elasticsearch/elasticsearch-esql.md
Normal file
45
docs/en/resources/tools/elasticsearch/elasticsearch-esql.md
Normal file
@@ -0,0 +1,45 @@
|
|||||||
|
---
|
||||||
|
title: "elasticsearch-esql"
|
||||||
|
type: docs
|
||||||
|
weight: 2
|
||||||
|
description: >
|
||||||
|
Execute ES|QL queries.
|
||||||
|
---
|
||||||
|
|
||||||
|
# elasticsearch-esql
|
||||||
|
|
||||||
|
Execute ES|QL queries.
|
||||||
|
|
||||||
|
This tool allows you to execute ES|QL queries against your Elasticsearch
|
||||||
|
cluster. You can use this to perform complex searches and aggregations.
|
||||||
|
|
||||||
|
See the [official documentation](https://www.elastic.co/docs/reference/query-languages/esql/esql-getting-started) for more information.
|
||||||
|
|
||||||
|
## Example
|
||||||
|
|
||||||
|
```yaml
|
||||||
|
tools:
|
||||||
|
query_my_index:
|
||||||
|
kind: elasticsearch-esql
|
||||||
|
source: elasticsearch-source
|
||||||
|
description: Use this tool to execute ES|QL queries.
|
||||||
|
query: |
|
||||||
|
FROM my-index
|
||||||
|
| KEEP *
|
||||||
|
| LIMIT ?limit
|
||||||
|
parameters:
|
||||||
|
- name: limit
|
||||||
|
type: integer
|
||||||
|
description: Limit the number of results.
|
||||||
|
required: true
|
||||||
|
```
|
||||||
|
|
||||||
|
## Parameters
|
||||||
|
|
||||||
|
| **name** | **type** | **required** | **description** |
|
||||||
|
|------------|:--------:|:------------:|-----------------------------------------------------------------------------------------------------------------------------------------------------|
|
||||||
|
| query | string | false | The ES\|QL query to run. Can also be passed by parameters. |
|
||||||
|
| format | string | false | The format of the query. Default is json. Valid values are csv, json, tsv, txt, yaml, cbor, smile, or arrow. |
|
||||||
|
| timeout | integer | false | The timeout for the query in seconds. Default is 60 (1 minute). |
|
||||||
|
| parameters | [parameters](../#specifying-parameters) | false | List of [parameters](../#specifying-parameters) that will be used with the ES\|QL query.<br/>Only supports “string”, “integer”, “float”, “boolean”. |
|
||||||
|
|
||||||
2
go.mod
2
go.mod
@@ -21,6 +21,8 @@ require (
|
|||||||
github.com/cenkalti/backoff/v5 v5.0.3
|
github.com/cenkalti/backoff/v5 v5.0.3
|
||||||
github.com/couchbase/gocb/v2 v2.11.1
|
github.com/couchbase/gocb/v2 v2.11.1
|
||||||
github.com/couchbase/tools-common/http v1.0.9
|
github.com/couchbase/tools-common/http v1.0.9
|
||||||
|
github.com/elastic/elastic-transport-go/v8 v8.7.0
|
||||||
|
github.com/elastic/go-elasticsearch/v8 v8.19.0
|
||||||
github.com/fsnotify/fsnotify v1.9.0
|
github.com/fsnotify/fsnotify v1.9.0
|
||||||
github.com/go-chi/chi/v5 v5.2.3
|
github.com/go-chi/chi/v5 v5.2.3
|
||||||
github.com/go-chi/httplog/v2 v2.1.1
|
github.com/go-chi/httplog/v2 v2.1.1
|
||||||
|
|||||||
4
go.sum
4
go.sum
@@ -820,6 +820,10 @@ github.com/docopt/docopt-go v0.0.0-20180111231733-ee0de3bc6815/go.mod h1:WwZ+bS3
|
|||||||
github.com/dustin/go-humanize v1.0.0/go.mod h1:HtrtbFcZ19U5GC7JDqmcUSB87Iq5E25KnS6fMYU6eOk=
|
github.com/dustin/go-humanize v1.0.0/go.mod h1:HtrtbFcZ19U5GC7JDqmcUSB87Iq5E25KnS6fMYU6eOk=
|
||||||
github.com/dustin/go-humanize v1.0.1 h1:GzkhY7T5VNhEkwH0PVJgjz+fX1rhBrR7pRT3mDkpeCY=
|
github.com/dustin/go-humanize v1.0.1 h1:GzkhY7T5VNhEkwH0PVJgjz+fX1rhBrR7pRT3mDkpeCY=
|
||||||
github.com/dustin/go-humanize v1.0.1/go.mod h1:Mu1zIs6XwVuF/gI1OepvI0qD18qycQx+mFykh5fBlto=
|
github.com/dustin/go-humanize v1.0.1/go.mod h1:Mu1zIs6XwVuF/gI1OepvI0qD18qycQx+mFykh5fBlto=
|
||||||
|
github.com/elastic/elastic-transport-go/v8 v8.7.0 h1:OgTneVuXP2uip4BA658Xi6Hfw+PeIOod2rY3GVMGoVE=
|
||||||
|
github.com/elastic/elastic-transport-go/v8 v8.7.0/go.mod h1:YLHer5cj0csTzNFXoNQ8qhtGY1GTvSqPnKWKaqQE3Hk=
|
||||||
|
github.com/elastic/go-elasticsearch/v8 v8.19.0 h1:VmfBLNRORY7RZL+9hTxBD97ehl9H8Nxf2QigDh6HuMU=
|
||||||
|
github.com/elastic/go-elasticsearch/v8 v8.19.0/go.mod h1:F3j9e+BubmKvzvLjNui/1++nJuJxbkhHefbaT0kFKGY=
|
||||||
github.com/envoyproxy/go-control-plane v0.9.0/go.mod h1:YTl/9mNaCwkRvm6d1a2C3ymFceY/DCBVvsKhRF0iEA4=
|
github.com/envoyproxy/go-control-plane v0.9.0/go.mod h1:YTl/9mNaCwkRvm6d1a2C3ymFceY/DCBVvsKhRF0iEA4=
|
||||||
github.com/envoyproxy/go-control-plane v0.9.1-0.20191026205805-5f8ba28d4473/go.mod h1:YTl/9mNaCwkRvm6d1a2C3ymFceY/DCBVvsKhRF0iEA4=
|
github.com/envoyproxy/go-control-plane v0.9.1-0.20191026205805-5f8ba28d4473/go.mod h1:YTl/9mNaCwkRvm6d1a2C3ymFceY/DCBVvsKhRF0iEA4=
|
||||||
github.com/envoyproxy/go-control-plane v0.9.4/go.mod h1:6rpuAdCZL397s3pYoYcLgu1mIlRU8Am5FuJP05cCM98=
|
github.com/envoyproxy/go-control-plane v0.9.4/go.mod h1:6rpuAdCZL397s3pYoYcLgu1mIlRU8Am5FuJP05cCM98=
|
||||||
|
|||||||
@@ -36,6 +36,7 @@ var expectedToolSources = []string{
|
|||||||
"cloud-sql-postgres-observability",
|
"cloud-sql-postgres-observability",
|
||||||
"cloud-sql-postgres",
|
"cloud-sql-postgres",
|
||||||
"dataplex",
|
"dataplex",
|
||||||
|
"elasticsearch",
|
||||||
"firestore",
|
"firestore",
|
||||||
"looker-conversational-analytics",
|
"looker-conversational-analytics",
|
||||||
"looker",
|
"looker",
|
||||||
|
|||||||
33
internal/prebuiltconfigs/tools/elasticsearch.yaml
Normal file
33
internal/prebuiltconfigs/tools/elasticsearch.yaml
Normal file
@@ -0,0 +1,33 @@
|
|||||||
|
# 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.
|
||||||
|
|
||||||
|
sources:
|
||||||
|
elasticsearch-source:
|
||||||
|
kind: elasticsearch
|
||||||
|
addresses:
|
||||||
|
- ${ELASTICSEARCH_HOST}
|
||||||
|
apikey: ${ELASTICSEARCH_APIKEY}
|
||||||
|
|
||||||
|
tools:
|
||||||
|
execute_esql_query:
|
||||||
|
kind: elasticsearch-esql
|
||||||
|
source: elasticsearch-source
|
||||||
|
description: Use this tool to execute ES|QL queries.
|
||||||
|
parameters:
|
||||||
|
- name: query
|
||||||
|
type: string
|
||||||
|
description: The ES|QL query to execute.
|
||||||
|
toolsets:
|
||||||
|
elasticsearch-tools:
|
||||||
|
- execute_esql_query
|
||||||
149
internal/sources/elasticsearch/elasticsearch.go
Normal file
149
internal/sources/elasticsearch/elasticsearch.go
Normal file
@@ -0,0 +1,149 @@
|
|||||||
|
// 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"
|
||||||
|
"net/http"
|
||||||
|
|
||||||
|
"github.com/elastic/elastic-transport-go/v8/elastictransport"
|
||||||
|
"github.com/elastic/go-elasticsearch/v8"
|
||||||
|
"github.com/elastic/go-elasticsearch/v8/esapi"
|
||||||
|
"github.com/goccy/go-yaml"
|
||||||
|
"github.com/googleapis/genai-toolbox/internal/sources"
|
||||||
|
"github.com/googleapis/genai-toolbox/internal/util"
|
||||||
|
"go.opentelemetry.io/otel/trace"
|
||||||
|
)
|
||||||
|
|
||||||
|
const SourceKind string = "elasticsearch"
|
||||||
|
|
||||||
|
// validate interface
|
||||||
|
var _ sources.SourceConfig = Config{}
|
||||||
|
|
||||||
|
func init() {
|
||||||
|
if !sources.Register(SourceKind, newConfig) {
|
||||||
|
panic(fmt.Sprintf("source kind %q already registered", SourceKind))
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func newConfig(ctx context.Context, name string, decoder *yaml.Decoder) (sources.SourceConfig, error) {
|
||||||
|
actual := Config{Name: name}
|
||||||
|
if err := decoder.DecodeContext(ctx, &actual); err != nil {
|
||||||
|
return nil, err
|
||||||
|
}
|
||||||
|
return actual, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
type Config struct {
|
||||||
|
Name string `yaml:"name" validate:"required"`
|
||||||
|
Kind string `yaml:"kind" validate:"required"`
|
||||||
|
Addresses []string `yaml:"addresses" validate:"required"`
|
||||||
|
Username string `yaml:"username"`
|
||||||
|
Password string `yaml:"password"`
|
||||||
|
APIKey string `yaml:"apikey"`
|
||||||
|
}
|
||||||
|
|
||||||
|
func (c Config) SourceConfigKind() string {
|
||||||
|
return SourceKind
|
||||||
|
}
|
||||||
|
|
||||||
|
type EsClient interface {
|
||||||
|
esapi.Transport
|
||||||
|
elastictransport.Instrumented
|
||||||
|
}
|
||||||
|
|
||||||
|
type Source struct {
|
||||||
|
Name string
|
||||||
|
Kind string
|
||||||
|
Client EsClient
|
||||||
|
}
|
||||||
|
|
||||||
|
var _ sources.Source = &Source{}
|
||||||
|
|
||||||
|
// tracerProviderAdapter adapts a Tracer to implement the TracerProvider interface
|
||||||
|
type tracerProviderAdapter struct {
|
||||||
|
trace.TracerProvider
|
||||||
|
tracer trace.Tracer
|
||||||
|
}
|
||||||
|
|
||||||
|
// Tracer implements the TracerProvider interface
|
||||||
|
func (t *tracerProviderAdapter) Tracer(name string, options ...trace.TracerOption) trace.Tracer {
|
||||||
|
return t.tracer
|
||||||
|
}
|
||||||
|
|
||||||
|
// Initialize creates a new Elasticsearch Source instance.
|
||||||
|
func (c Config) Initialize(ctx context.Context, tracer trace.Tracer) (sources.Source, error) {
|
||||||
|
tracerProvider := &tracerProviderAdapter{tracer: tracer}
|
||||||
|
|
||||||
|
ua, err := util.UserAgentFromContext(ctx)
|
||||||
|
if err != nil {
|
||||||
|
return nil, fmt.Errorf("error getting user agent from context: %w", err)
|
||||||
|
}
|
||||||
|
|
||||||
|
// Create a new Elasticsearch client with the provided configuration
|
||||||
|
cfg := elasticsearch.Config{
|
||||||
|
Addresses: c.Addresses,
|
||||||
|
Instrumentation: elasticsearch.NewOpenTelemetryInstrumentation(tracerProvider, false),
|
||||||
|
Header: http.Header{"User-Agent": []string{ua + " go-elasticsearch/" + elasticsearch.Version}},
|
||||||
|
}
|
||||||
|
|
||||||
|
// Client need either username and password or an API key
|
||||||
|
if c.Username != "" && c.Password != "" {
|
||||||
|
cfg.Username = c.Username
|
||||||
|
cfg.Password = c.Password
|
||||||
|
} else if c.APIKey != "" {
|
||||||
|
// API key will be set below
|
||||||
|
cfg.APIKey = c.APIKey
|
||||||
|
} else {
|
||||||
|
// If neither username/password nor API key is provided, we throw an error
|
||||||
|
return nil, fmt.Errorf("elasticsearch source %q requires either username/password or an API key", c.Name)
|
||||||
|
}
|
||||||
|
|
||||||
|
client, err := elasticsearch.NewBaseClient(cfg)
|
||||||
|
if err != nil {
|
||||||
|
return nil, err
|
||||||
|
}
|
||||||
|
|
||||||
|
// Test connection
|
||||||
|
res, err := esapi.InfoRequest{
|
||||||
|
Instrument: client.InstrumentationEnabled(),
|
||||||
|
}.Do(ctx, client)
|
||||||
|
|
||||||
|
if err != nil {
|
||||||
|
return nil, err
|
||||||
|
}
|
||||||
|
defer res.Body.Close()
|
||||||
|
|
||||||
|
if res.IsError() {
|
||||||
|
return nil, fmt.Errorf("elasticsearch connection failed: status %d", res.StatusCode)
|
||||||
|
}
|
||||||
|
|
||||||
|
s := &Source{
|
||||||
|
Name: c.Name,
|
||||||
|
Kind: SourceKind,
|
||||||
|
Client: client,
|
||||||
|
}
|
||||||
|
return s, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
// SourceKind returns the kind string for this source.
|
||||||
|
func (s *Source) SourceKind() string {
|
||||||
|
return SourceKind
|
||||||
|
}
|
||||||
|
|
||||||
|
func (s *Source) ElasticsearchClient() EsClient {
|
||||||
|
return s.Client
|
||||||
|
}
|
||||||
66
internal/sources/elasticsearch/elasticsearch_test.go
Normal file
66
internal/sources/elasticsearch/elasticsearch_test.go
Normal file
@@ -0,0 +1,66 @@
|
|||||||
|
// 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_test
|
||||||
|
|
||||||
|
import (
|
||||||
|
"testing"
|
||||||
|
|
||||||
|
yaml "github.com/goccy/go-yaml"
|
||||||
|
"github.com/google/go-cmp/cmp"
|
||||||
|
"github.com/googleapis/genai-toolbox/internal/server"
|
||||||
|
"github.com/googleapis/genai-toolbox/internal/sources/elasticsearch"
|
||||||
|
)
|
||||||
|
|
||||||
|
func TestParseFromYamlElasticsearch(t *testing.T) {
|
||||||
|
tcs := []struct {
|
||||||
|
desc string
|
||||||
|
in string
|
||||||
|
want server.SourceConfigs
|
||||||
|
}{
|
||||||
|
{
|
||||||
|
desc: "basic example",
|
||||||
|
in: `
|
||||||
|
sources:
|
||||||
|
my-es-instance:
|
||||||
|
kind: elasticsearch
|
||||||
|
addresses:
|
||||||
|
- http://localhost:9200
|
||||||
|
apikey: somekey
|
||||||
|
`,
|
||||||
|
want: server.SourceConfigs{
|
||||||
|
"my-es-instance": elasticsearch.Config{
|
||||||
|
Name: "my-es-instance",
|
||||||
|
Kind: elasticsearch.SourceKind,
|
||||||
|
Addresses: []string{"http://localhost:9200"},
|
||||||
|
APIKey: "somekey",
|
||||||
|
},
|
||||||
|
},
|
||||||
|
},
|
||||||
|
}
|
||||||
|
for _, tc := range tcs {
|
||||||
|
t.Run(tc.desc, func(t *testing.T) {
|
||||||
|
got := struct {
|
||||||
|
Sources server.SourceConfigs `yaml:"sources"`
|
||||||
|
}{}
|
||||||
|
err := yaml.Unmarshal([]byte(tc.in), &got)
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("failed to parse yaml: %v", err)
|
||||||
|
}
|
||||||
|
if diff := cmp.Diff(tc.want, got.Sources); diff != "" {
|
||||||
|
t.Errorf("unexpected config diff (-want +got):\n%s", diff)
|
||||||
|
}
|
||||||
|
})
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -0,0 +1,243 @@
|
|||||||
|
// 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 elasticsearchesql
|
||||||
|
|
||||||
|
import (
|
||||||
|
"bytes"
|
||||||
|
"context"
|
||||||
|
"encoding/json"
|
||||||
|
"fmt"
|
||||||
|
"time"
|
||||||
|
|
||||||
|
"github.com/elastic/go-elasticsearch/v8/esapi"
|
||||||
|
"github.com/googleapis/genai-toolbox/internal/util"
|
||||||
|
|
||||||
|
"github.com/goccy/go-yaml"
|
||||||
|
"github.com/googleapis/genai-toolbox/internal/sources"
|
||||||
|
es "github.com/googleapis/genai-toolbox/internal/sources/elasticsearch"
|
||||||
|
"github.com/googleapis/genai-toolbox/internal/tools"
|
||||||
|
)
|
||||||
|
|
||||||
|
const kind string = "elasticsearch-esql"
|
||||||
|
|
||||||
|
func init() {
|
||||||
|
if !tools.Register(kind, newConfig) {
|
||||||
|
panic(fmt.Sprintf("tool kind %q already registered", kind))
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
type compatibleSource interface {
|
||||||
|
ElasticsearchClient() es.EsClient
|
||||||
|
}
|
||||||
|
|
||||||
|
var _ compatibleSource = &es.Source{}
|
||||||
|
|
||||||
|
var compatibleSources = [...]string{es.SourceKind}
|
||||||
|
|
||||||
|
type Config struct {
|
||||||
|
Name string `yaml:"name" validate:"required"`
|
||||||
|
Kind string `yaml:"kind" validate:"required"`
|
||||||
|
Source string `yaml:"source" validate:"required"`
|
||||||
|
Description string `yaml:"description" validate:"required"`
|
||||||
|
AuthRequired []string `yaml:"authRequired" validate:"required"`
|
||||||
|
Query string `yaml:"query"`
|
||||||
|
Format string `yaml:"format"`
|
||||||
|
Timeout int `yaml:"timeout"`
|
||||||
|
Parameters tools.Parameters `yaml:"parameters"`
|
||||||
|
}
|
||||||
|
|
||||||
|
var _ tools.ToolConfig = Config{}
|
||||||
|
|
||||||
|
func (c Config) ToolConfigKind() string {
|
||||||
|
return kind
|
||||||
|
}
|
||||||
|
|
||||||
|
func newConfig(ctx context.Context, name string, decoder *yaml.Decoder) (tools.ToolConfig, error) {
|
||||||
|
actual := Config{Name: name}
|
||||||
|
if err := decoder.DecodeContext(ctx, &actual); err != nil {
|
||||||
|
return nil, err
|
||||||
|
}
|
||||||
|
return actual, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
type Tool struct {
|
||||||
|
Name string `yaml:"name"`
|
||||||
|
Kind string `yaml:"kind"`
|
||||||
|
AuthRequired []string `yaml:"authRequired"`
|
||||||
|
Parameters tools.Parameters `yaml:"parameters"`
|
||||||
|
Query string `yaml:"query"`
|
||||||
|
Format string `yaml:"format" default:"json"`
|
||||||
|
Timeout int `yaml:"timeout"`
|
||||||
|
|
||||||
|
manifest tools.Manifest
|
||||||
|
mcpManifest tools.McpManifest
|
||||||
|
EsClient es.EsClient
|
||||||
|
}
|
||||||
|
|
||||||
|
var _ tools.Tool = Tool{}
|
||||||
|
|
||||||
|
func (c Config) Initialize(srcs map[string]sources.Source) (tools.Tool, error) {
|
||||||
|
// verify source exists
|
||||||
|
src, ok := srcs[c.Source]
|
||||||
|
if !ok {
|
||||||
|
return nil, fmt.Errorf("source %q not found", c.Source)
|
||||||
|
}
|
||||||
|
|
||||||
|
// verify the source is compatible
|
||||||
|
s, ok := src.(compatibleSource)
|
||||||
|
if !ok {
|
||||||
|
return nil, fmt.Errorf("invalid source for %q tool: source kind must be one of %q", kind, compatibleSources)
|
||||||
|
}
|
||||||
|
|
||||||
|
mcpManifest := tools.GetMcpManifest(c.Name, c.Description, c.AuthRequired, c.Parameters)
|
||||||
|
|
||||||
|
return Tool{
|
||||||
|
Name: c.Name,
|
||||||
|
Kind: kind,
|
||||||
|
Parameters: c.Parameters,
|
||||||
|
Query: c.Query,
|
||||||
|
Format: c.Format,
|
||||||
|
Timeout: c.Timeout,
|
||||||
|
AuthRequired: c.AuthRequired,
|
||||||
|
EsClient: s.ElasticsearchClient(),
|
||||||
|
manifest: tools.Manifest{Description: c.Description, Parameters: c.Parameters.Manifest(), AuthRequired: c.AuthRequired},
|
||||||
|
mcpManifest: mcpManifest,
|
||||||
|
}, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
type esqlColumn struct {
|
||||||
|
Name string `json:"name"`
|
||||||
|
Type string `json:"type"`
|
||||||
|
}
|
||||||
|
|
||||||
|
type esqlResult struct {
|
||||||
|
Columns []esqlColumn `json:"columns"`
|
||||||
|
Values [][]any `json:"values"`
|
||||||
|
}
|
||||||
|
|
||||||
|
func (t Tool) Invoke(ctx context.Context, params tools.ParamValues, accessToken tools.AccessToken) (any, error) {
|
||||||
|
var cancel context.CancelFunc
|
||||||
|
if t.Timeout > 0 {
|
||||||
|
ctx, cancel = context.WithTimeout(ctx, time.Duration(t.Timeout)*time.Second)
|
||||||
|
defer cancel()
|
||||||
|
} else {
|
||||||
|
ctx, cancel = context.WithTimeout(ctx, time.Minute)
|
||||||
|
defer cancel()
|
||||||
|
}
|
||||||
|
|
||||||
|
bodyStruct := struct {
|
||||||
|
Query string `json:"query"`
|
||||||
|
Params []map[string]any `json:"params,omitempty"`
|
||||||
|
}{
|
||||||
|
Query: t.Query,
|
||||||
|
Params: 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 query, ok := paramMap["query"]; ok {
|
||||||
|
if str, ok := query.(string); ok && bodyStruct.Query == "" {
|
||||||
|
bodyStruct.Query = str
|
||||||
|
}
|
||||||
|
|
||||||
|
// Drop the query param if not a string or if the tool already has a query.
|
||||||
|
delete(paramMap, "query")
|
||||||
|
}
|
||||||
|
|
||||||
|
for _, param := range t.Parameters {
|
||||||
|
if param.GetType() == "array" {
|
||||||
|
return nil, fmt.Errorf("array parameters are not supported yet")
|
||||||
|
}
|
||||||
|
bodyStruct.Params = append(bodyStruct.Params, map[string]any{param.GetName(): paramMap[param.GetName()]})
|
||||||
|
}
|
||||||
|
|
||||||
|
body, err := json.Marshal(bodyStruct)
|
||||||
|
if err != nil {
|
||||||
|
return nil, fmt.Errorf("failed to marshal query body: %w", err)
|
||||||
|
}
|
||||||
|
res, err := esapi.EsqlQueryRequest{
|
||||||
|
Body: bytes.NewReader(body),
|
||||||
|
Format: t.Format,
|
||||||
|
FilterPath: []string{"columns", "values"},
|
||||||
|
Instrument: t.EsClient.InstrumentationEnabled(),
|
||||||
|
}.Do(ctx, t.EsClient)
|
||||||
|
|
||||||
|
if err != nil {
|
||||||
|
return nil, err
|
||||||
|
}
|
||||||
|
defer res.Body.Close()
|
||||||
|
|
||||||
|
if res.IsError() {
|
||||||
|
// Try to extract error message from response
|
||||||
|
var esErr json.RawMessage
|
||||||
|
err = util.DecodeJSON(res.Body, &esErr)
|
||||||
|
if err != nil {
|
||||||
|
return nil, fmt.Errorf("elasticsearch error: status %s", res.Status())
|
||||||
|
}
|
||||||
|
return esErr, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
var result esqlResult
|
||||||
|
err = util.DecodeJSON(res.Body, &result)
|
||||||
|
if err != nil {
|
||||||
|
return nil, fmt.Errorf("failed to decode response body: %w", err)
|
||||||
|
}
|
||||||
|
|
||||||
|
output := t.esqlToMap(result)
|
||||||
|
|
||||||
|
return output, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
// esqlToMap converts the esqlResult to a slice of maps.
|
||||||
|
func (t Tool) esqlToMap(result esqlResult) []map[string]any {
|
||||||
|
output := make([]map[string]any, 0, len(result.Values))
|
||||||
|
for _, value := range result.Values {
|
||||||
|
row := make(map[string]any)
|
||||||
|
if value == nil {
|
||||||
|
output = append(output, row)
|
||||||
|
continue
|
||||||
|
}
|
||||||
|
for i, col := range result.Columns {
|
||||||
|
if i < len(value) {
|
||||||
|
row[col.Name] = value[i]
|
||||||
|
} else {
|
||||||
|
row[col.Name] = nil
|
||||||
|
}
|
||||||
|
}
|
||||||
|
output = append(output, row)
|
||||||
|
}
|
||||||
|
return output
|
||||||
|
}
|
||||||
|
|
||||||
|
func (t Tool) ParseParams(data map[string]any, claims map[string]map[string]any) (tools.ParamValues, error) {
|
||||||
|
return tools.ParseParams(t.Parameters, data, claims)
|
||||||
|
}
|
||||||
|
|
||||||
|
func (t Tool) Manifest() tools.Manifest {
|
||||||
|
return t.manifest
|
||||||
|
}
|
||||||
|
|
||||||
|
func (t Tool) McpManifest() tools.McpManifest {
|
||||||
|
return t.mcpManifest
|
||||||
|
}
|
||||||
|
|
||||||
|
func (t Tool) Authorized(verifiedAuthServices []string) bool {
|
||||||
|
return tools.IsAuthorized(t.AuthRequired, verifiedAuthServices)
|
||||||
|
}
|
||||||
|
|
||||||
|
func (t Tool) RequiresClientAuthorization() bool {
|
||||||
|
return false
|
||||||
|
}
|
||||||
@@ -0,0 +1,261 @@
|
|||||||
|
// 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 elasticsearchesql
|
||||||
|
|
||||||
|
import (
|
||||||
|
"reflect"
|
||||||
|
"testing"
|
||||||
|
|
||||||
|
"github.com/goccy/go-yaml"
|
||||||
|
"github.com/google/go-cmp/cmp"
|
||||||
|
"github.com/googleapis/genai-toolbox/internal/server"
|
||||||
|
"github.com/googleapis/genai-toolbox/internal/testutils"
|
||||||
|
"github.com/googleapis/genai-toolbox/internal/tools"
|
||||||
|
)
|
||||||
|
|
||||||
|
func TestParseFromYamlElasticsearchEsql(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 search example",
|
||||||
|
in: `
|
||||||
|
tools:
|
||||||
|
example_tool:
|
||||||
|
kind: elasticsearch-esql
|
||||||
|
source: my-elasticsearch-instance
|
||||||
|
description: Elasticsearch ES|QL tool
|
||||||
|
query: |
|
||||||
|
FROM my-index
|
||||||
|
| KEEP first_name, last_name
|
||||||
|
`,
|
||||||
|
want: server.ToolConfigs{
|
||||||
|
"example_tool": Config{
|
||||||
|
Name: "example_tool",
|
||||||
|
Kind: "elasticsearch-esql",
|
||||||
|
Source: "my-elasticsearch-instance",
|
||||||
|
Description: "Elasticsearch ES|QL tool",
|
||||||
|
AuthRequired: []string{},
|
||||||
|
Query: "FROM my-index\n| KEEP first_name, last_name\n",
|
||||||
|
},
|
||||||
|
},
|
||||||
|
},
|
||||||
|
{
|
||||||
|
desc: "search with customizable limit parameter",
|
||||||
|
in: `
|
||||||
|
tools:
|
||||||
|
example_tool:
|
||||||
|
kind: elasticsearch-esql
|
||||||
|
source: my-elasticsearch-instance
|
||||||
|
description: Elasticsearch ES|QL tool with customizable limit
|
||||||
|
parameters:
|
||||||
|
- name: limit
|
||||||
|
type: integer
|
||||||
|
description: Limit the number of results
|
||||||
|
query: |
|
||||||
|
FROM my-index
|
||||||
|
| LIMIT ?limit
|
||||||
|
`,
|
||||||
|
want: server.ToolConfigs{
|
||||||
|
"example_tool": Config{
|
||||||
|
Name: "example_tool",
|
||||||
|
Kind: "elasticsearch-esql",
|
||||||
|
Source: "my-elasticsearch-instance",
|
||||||
|
Description: "Elasticsearch ES|QL tool with customizable limit",
|
||||||
|
AuthRequired: []string{},
|
||||||
|
Parameters: tools.Parameters{
|
||||||
|
tools.NewIntParameter("limit", "Limit the number of results"),
|
||||||
|
},
|
||||||
|
Query: "FROM my-index\n| LIMIT ?limit\n",
|
||||||
|
},
|
||||||
|
},
|
||||||
|
},
|
||||||
|
}
|
||||||
|
|
||||||
|
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)
|
||||||
|
}
|
||||||
|
})
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestTool_esqlToMap(t1 *testing.T) {
|
||||||
|
tests := []struct {
|
||||||
|
name string
|
||||||
|
result esqlResult
|
||||||
|
want []map[string]any
|
||||||
|
}{
|
||||||
|
{
|
||||||
|
name: "simple case with two rows",
|
||||||
|
result: esqlResult{
|
||||||
|
Columns: []esqlColumn{
|
||||||
|
{Name: "first_name", Type: "text"},
|
||||||
|
{Name: "last_name", Type: "text"},
|
||||||
|
},
|
||||||
|
Values: [][]any{
|
||||||
|
{"John", "Doe"},
|
||||||
|
{"Jane", "Smith"},
|
||||||
|
},
|
||||||
|
},
|
||||||
|
want: []map[string]any{
|
||||||
|
{"first_name": "John", "last_name": "Doe"},
|
||||||
|
{"first_name": "Jane", "last_name": "Smith"},
|
||||||
|
},
|
||||||
|
},
|
||||||
|
{
|
||||||
|
name: "different data types",
|
||||||
|
result: esqlResult{
|
||||||
|
Columns: []esqlColumn{
|
||||||
|
{Name: "id", Type: "integer"},
|
||||||
|
{Name: "active", Type: "boolean"},
|
||||||
|
{Name: "score", Type: "float"},
|
||||||
|
},
|
||||||
|
Values: [][]any{
|
||||||
|
{1, true, 95.5},
|
||||||
|
{2, false, 88.0},
|
||||||
|
},
|
||||||
|
},
|
||||||
|
want: []map[string]any{
|
||||||
|
{"id": 1, "active": true, "score": 95.5},
|
||||||
|
{"id": 2, "active": false, "score": 88.0},
|
||||||
|
},
|
||||||
|
},
|
||||||
|
{
|
||||||
|
name: "no rows",
|
||||||
|
result: esqlResult{
|
||||||
|
Columns: []esqlColumn{
|
||||||
|
{Name: "id", Type: "integer"},
|
||||||
|
{Name: "name", Type: "text"},
|
||||||
|
},
|
||||||
|
Values: [][]any{},
|
||||||
|
},
|
||||||
|
want: []map[string]any{},
|
||||||
|
},
|
||||||
|
{
|
||||||
|
name: "null values",
|
||||||
|
result: esqlResult{
|
||||||
|
Columns: []esqlColumn{
|
||||||
|
{Name: "id", Type: "integer"},
|
||||||
|
{Name: "name", Type: "text"},
|
||||||
|
},
|
||||||
|
Values: [][]any{
|
||||||
|
{1, nil},
|
||||||
|
{2, "Alice"},
|
||||||
|
},
|
||||||
|
},
|
||||||
|
want: []map[string]any{
|
||||||
|
{"id": 1, "name": nil},
|
||||||
|
{"id": 2, "name": "Alice"},
|
||||||
|
},
|
||||||
|
},
|
||||||
|
{
|
||||||
|
name: "missing values in a row",
|
||||||
|
result: esqlResult{
|
||||||
|
Columns: []esqlColumn{
|
||||||
|
{Name: "id", Type: "integer"},
|
||||||
|
{Name: "name", Type: "text"},
|
||||||
|
{Name: "age", Type: "integer"},
|
||||||
|
},
|
||||||
|
Values: [][]any{
|
||||||
|
{1, "Bob"},
|
||||||
|
{2, "Charlie", 30},
|
||||||
|
},
|
||||||
|
},
|
||||||
|
want: []map[string]any{
|
||||||
|
{"id": 1, "name": "Bob", "age": nil},
|
||||||
|
{"id": 2, "name": "Charlie", "age": 30},
|
||||||
|
},
|
||||||
|
},
|
||||||
|
{
|
||||||
|
name: "all null row",
|
||||||
|
result: esqlResult{
|
||||||
|
Columns: []esqlColumn{
|
||||||
|
{Name: "id", Type: "integer"},
|
||||||
|
{Name: "name", Type: "text"},
|
||||||
|
},
|
||||||
|
Values: [][]any{
|
||||||
|
nil,
|
||||||
|
},
|
||||||
|
},
|
||||||
|
want: []map[string]any{
|
||||||
|
{},
|
||||||
|
},
|
||||||
|
},
|
||||||
|
{
|
||||||
|
name: "empty columns",
|
||||||
|
result: esqlResult{
|
||||||
|
Columns: []esqlColumn{},
|
||||||
|
Values: [][]any{
|
||||||
|
{},
|
||||||
|
{},
|
||||||
|
},
|
||||||
|
},
|
||||||
|
want: []map[string]any{
|
||||||
|
{},
|
||||||
|
{},
|
||||||
|
},
|
||||||
|
},
|
||||||
|
{
|
||||||
|
name: "more values than columns",
|
||||||
|
result: esqlResult{
|
||||||
|
Columns: []esqlColumn{
|
||||||
|
{Name: "id", Type: "integer"},
|
||||||
|
},
|
||||||
|
Values: [][]any{
|
||||||
|
{1, "extra"},
|
||||||
|
},
|
||||||
|
},
|
||||||
|
want: []map[string]any{
|
||||||
|
{"id": 1},
|
||||||
|
},
|
||||||
|
},
|
||||||
|
{
|
||||||
|
name: "no columns but with values",
|
||||||
|
result: esqlResult{
|
||||||
|
Columns: []esqlColumn{},
|
||||||
|
Values: [][]any{
|
||||||
|
{1, "data"},
|
||||||
|
},
|
||||||
|
},
|
||||||
|
want: []map[string]any{
|
||||||
|
{},
|
||||||
|
},
|
||||||
|
},
|
||||||
|
}
|
||||||
|
for _, tt := range tests {
|
||||||
|
t1.Run(tt.name, func(t1 *testing.T) {
|
||||||
|
t := Tool{}
|
||||||
|
if got := t.esqlToMap(tt.result); !reflect.DeepEqual(got, tt.want) {
|
||||||
|
t1.Errorf("esqlToMap() = %v, want %v", got, tt.want)
|
||||||
|
}
|
||||||
|
})
|
||||||
|
}
|
||||||
|
}
|
||||||
309
tests/elasticsearch/elasticsearch_integration_test.go
Normal file
309
tests/elasticsearch/elasticsearch_integration_test.go
Normal file
@@ -0,0 +1,309 @@
|
|||||||
|
// 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/v8"
|
||||||
|
"github.com/elastic/go-elasticsearch/v8/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
|
||||||
|
}
|
||||||
@@ -18,6 +18,7 @@ package tests
|
|||||||
|
|
||||||
// InvokeTestConfig represents the various configuration options for RunToolInvokeTest()
|
// InvokeTestConfig represents the various configuration options for RunToolInvokeTest()
|
||||||
type InvokeTestConfig struct {
|
type InvokeTestConfig struct {
|
||||||
|
myAuthToolWant string
|
||||||
myToolId3NameAliceWant string
|
myToolId3NameAliceWant string
|
||||||
myToolById4Want string
|
myToolById4Want string
|
||||||
nullWant string
|
nullWant string
|
||||||
@@ -31,6 +32,14 @@ type InvokeTestConfig struct {
|
|||||||
|
|
||||||
type InvokeTestOption func(*InvokeTestConfig)
|
type InvokeTestOption func(*InvokeTestConfig)
|
||||||
|
|
||||||
|
// WithMyAuthToolWant represents the response value for my-auth-tool.
|
||||||
|
// e.g. tests.RunToolInvokeTest(t, select1Want, tests.WithMyAuthToolWant("custom"))
|
||||||
|
func WithMyAuthToolWant(s string) InvokeTestOption {
|
||||||
|
return func(c *InvokeTestConfig) {
|
||||||
|
c.myAuthToolWant = s
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
// WithMyToolId3NameAliceWant represents the response value for my-tool with id=3 and name=Alice.
|
// WithMyToolId3NameAliceWant represents the response value for my-tool with id=3 and name=Alice.
|
||||||
// e.g. tests.RunToolInvokeTest(t, select1Want, tests.WithMyToolId3NameAliceWant("custom"))
|
// e.g. tests.RunToolInvokeTest(t, select1Want, tests.WithMyToolId3NameAliceWant("custom"))
|
||||||
func WithMyToolId3NameAliceWant(s string) InvokeTestOption {
|
func WithMyToolId3NameAliceWant(s string) InvokeTestOption {
|
||||||
|
|||||||
@@ -364,7 +364,7 @@ func RunToolInvokeTest(t *testing.T, select1Want string, options ...InvokeTestOp
|
|||||||
enabled: configs.supportSelect1Auth,
|
enabled: configs.supportSelect1Auth,
|
||||||
requestHeader: map[string]string{"my-google-auth_token": idToken},
|
requestHeader: map[string]string{"my-google-auth_token": idToken},
|
||||||
requestBody: bytes.NewBuffer([]byte(`{}`)),
|
requestBody: bytes.NewBuffer([]byte(`{}`)),
|
||||||
wantBody: "[{\"name\":\"Alice\"}]",
|
wantBody: configs.myAuthToolWant,
|
||||||
wantStatusCode: http.StatusOK,
|
wantStatusCode: http.StatusOK,
|
||||||
},
|
},
|
||||||
{
|
{
|
||||||
|
|||||||
Reference in New Issue
Block a user