diff --git a/.ci/integration.cloudbuild.yaml b/.ci/integration.cloudbuild.yaml index 1d7bdf1f2b..3dbb2695e5 100644 --- a/.ci/integration.cloudbuild.yaml +++ b/.ci/integration.cloudbuild.yaml @@ -87,6 +87,25 @@ steps: - | go test -race -v -tags=integration,alloydb_ai_nl ./tests + - id: "bigtable" + name: golang:1 + waitFor: ["install-dependencies"] + entrypoint: /bin/bash + env: + - "GOPATH=/gopath" + - "BIGTABLE_PROJECT=$PROJECT_ID" + - "BIGTABLE_INSTANCE=$_BIGTABLE_INSTANCE" + - "SERVICE_ACCOUNT_EMAIL=$SERVICE_ACCOUNT_EMAIL" + secretEnv: + ["CLIENT_ID"] + volumes: + - name: "go" + path: "/gopath" + args: + - -c + - | + go test -race -v -tags=integration,bigtable ./tests + - id: "postgres" name: golang:1 waitFor: ["install-dependencies"] @@ -310,6 +329,7 @@ substitutions: _ALLOYDB_POSTGRES_INSTANCE: "alloydb-pg-testing-instance" _ALLOYDB_AI_NL_CLUSTER: "alloydb-ai-nl-testing" _ALLOYDB_AI_NL_INSTANCE: "alloydb-ai-nl-testing-instance" + _BIGTABLE_INSTANCE: "bigtable-testing-instance" _POSTGRES_HOST: 127.0.0.1 _POSTGRES_PORT: "5432" _SPANNER_INSTANCE: "spanner-testing" diff --git a/.golangci.yaml b/.golangci.yaml index 6dbb7a296c..d9a96734eb 100644 --- a/.golangci.yaml +++ b/.golangci.yaml @@ -45,3 +45,4 @@ run: - mysql - http - alloydb_ai_nl + - bigtable diff --git a/docs/en/resources/sources/bigtable.md b/docs/en/resources/sources/bigtable.md new file mode 100644 index 0000000000..f61b976c6f --- /dev/null +++ b/docs/en/resources/sources/bigtable.md @@ -0,0 +1,70 @@ +--- +title: "Bigtable" +type: docs +weight: 1 +description: > + Bigtable is a low-latency NoSQL database service for machine learning, operational analytics, and user-facing operations. It's a wide-column, key-value store that can scale to billions of rows and thousands of columns. With Bigtable, you can replicate your data to regions across the world for high availability and data resiliency. + +--- + +# Bigtable Source + +[Bigtable][bigtable-docs] is a low-latency NoSQL database service for machine +learning, operational analytics, and user-facing operations. It's a wide-column, +key-value store that can scale to billions of rows and thousands of columns. +With Bigtable, you can replicate your data to regions across the world for high +availability and data resiliency. + +If you are new to Bigtable, you can try to [create an instance and write data +with the cbt CLI][bigtable-quickstart-with-cli]. + +You can use [GoogleSQL statements][bigtable-googlesql] to query your Bigtable +data. GoogleSQL is an ANSI-compliant structured query language (SQL) that is +also implemented for other Google Cloud services. SQL queries are handled by +cluster nodes in the same way as NoSQL data requests. Therefore, the same best +practices apply when creating SQL queries to run against your Bigtable data, +such as avoiding full table scans or complex filters. + +[bigtable-docs]: https://cloud.google.com/bigtable/docs +[bigtable-quickstart-with-cli]: + https://cloud.google.com/bigtable/docs/create-instance-write-data-cbt-cli + +[bigtable-googlesql]: + https://cloud.google.com/bigtable/docs/googlesql-overview + +## Requirements + +### IAM Permissions + +Bigtable uses [Identity and Access Management (IAM)][iam-overview] to control +user and group access to Bigtable resources at the project, instance, table, and +backup level. Toolbox will use your [Application Default Credentials (ADC)][adc] +to authorize and authenticate when interacting with [Bigtable][bigtable-docs]. + +In addition to [setting the ADC for your server][set-adc], you need to ensure +the IAM identity has been given the correct IAM permissions for the query +provided. See [Apply IAM roles][grant-permissions] for more information on +applying IAM permissions and roles to an identity. + +[iam-overview]: https://cloud.google.com/bigtable/docs/access-control +[adc]: https://cloud.google.com/docs/authentication#adc +[set-adc]: https://cloud.google.com/docs/authentication/provide-credentials-adc +[grant-permissions]: https://cloud.google.com/bigtable/docs/access-control#iam-management-instance + +## Example + +```yaml +sources: + my-bigtable-source: + kind: "bigtable" + project: "my-project-id" + instance: "test-instance" +``` + +## Reference + +| **field** | **type** | **required** | **description** | +|-----------|:--------:|:------------:|-------------------------------------------------------------------------------| +| kind | string | true | Must be "bigtable". | +| project | string | true | Id of the GCP project that the cluster was created in (e.g. "my-project-id"). | +| instance | string | true | Name of the Bigtable instance. | diff --git a/docs/en/resources/tools/bigtable-sql.md b/docs/en/resources/tools/bigtable-sql.md new file mode 100644 index 0000000000..7eb64c8218 --- /dev/null +++ b/docs/en/resources/tools/bigtable-sql.md @@ -0,0 +1,82 @@ +--- +title: "bigtable-sql" +type: docs +weight: 1 +description: > + A "bigtable-sql" tool executes a pre-defined SQL statement against a Google + Cloud Bigtable instance. +--- + +## About + +A `bigtable-sql` tool executes a pre-defined SQL statement against a Bigtable +instance. It's compatible with any of the following sources: + +- [bigtable](../sources/bigtable.md) + +### GoogleSQL + +Bigtable supports SQL queries. The integration with Toolbox supports `googlesql` +dialect, the specified SQL statement is executed as a [data manipulation +language (DML)][bigtable-googlesql] statements, and specified parameters will +inserted according to their name: e.g. `@name`. + +[bigtable-googlesql]: https://cloud.google.com/bigtable/docs/googlesql-overview + +## Example + +```yaml +tools: + search_user_by_id_or_name: + kind: bigtable-sql + source: my-bigtable-instance + statement: | + SELECT + TO_INT64(cf[ 'id' ]) as id, + CAST(cf[ 'name' ] AS string) as name, + FROM + % s + WHERE + TO_INT64(cf[ 'id' ]) = @id + OR CAST(cf[ 'name' ] AS string) = @name; + description: | + Use this tool to get information for a specific user. + Takes an id number or a name and returns info on the user. + + Example: + {{ + "id": 123, + "name": "Alice", + }} + parameters: + - name: id + type: integer + description: User ID + - name: name + type: string + description: Name of the user +``` + +## Reference + +| **field** | **type** | **required** | **description** | +|-------------|:------------------------------------------:|:------------:|--------------------------------------------------------------------------------------------------| +| kind | string | true | Must be "bigtable-sql". | +| source | string | true | Name of the source the SQL should execute on. | +| description | string | true | Description of the tool that is passed to the LLM. | +| statement | string | true | SQL statement to execute on. | +| parameters | [parameters](_index#specifying-parameters) | false | List of [parameters](_index#specifying-parameters) that will be inserted into the SQL statement. | + +## Tips + +- [Bigtable Studio][bigtable-studio] is a useful to explore and manage your + Bigtable data. If you're unfamiliar with the query syntax, [Query + Builder][bigtable-querybuilder] lets you build a query, run it against a + table, and then view the results in the console. +- Some Python libraries limit the use of underscore columns such as `_key`. A + workaround would be to leverage Bigtable [Logical + Views][bigtable-logical-view] to rename the columns. + +[bigtable-studio]: https://cloud.google.com/bigtable/docs/manage-data-using-console +[bigtable-logical-view]: https://cloud.google.com/bigtable/docs/create-manage-logical-views +[bigtable-querybuilder]: https://cloud.google.com/bigtable/docs/query-builder diff --git a/go.mod b/go.mod index 4ef593c049..d27c12f46f 100644 --- a/go.mod +++ b/go.mod @@ -6,6 +6,7 @@ toolchain go1.23.2 require ( cloud.google.com/go/alloydbconn v1.15.1 + cloud.google.com/go/bigtable v1.36.0 cloud.google.com/go/cloudsqlconn v1.16.1 cloud.google.com/go/spanner v1.79.0 github.com/GoogleCloudPlatform/opentelemetry-operations-go/exporter/metric v0.51.0 diff --git a/go.sum b/go.sum index aa063e5479..6820999c8b 100644 --- a/go.sum +++ b/go.sum @@ -137,6 +137,8 @@ cloud.google.com/go/bigquery v1.47.0/go.mod h1:sA9XOgy0A8vQK9+MWhEQTY6Tix87M/Zur cloud.google.com/go/bigquery v1.48.0/go.mod h1:QAwSz+ipNgfL5jxiaK7weyOhzdoAy1zFm0Nf1fysJac= cloud.google.com/go/bigquery v1.49.0/go.mod h1:Sv8hMmTFFYBlt/ftw2uN6dFdQPzBlREY9yBh7Oy7/4Q= cloud.google.com/go/bigquery v1.50.0/go.mod h1:YrleYEh2pSEbgTBZYMJ5SuSr0ML3ypjRB1zgf7pvQLU= +cloud.google.com/go/bigtable v1.36.0 h1:GU4XWYb7H9XYHvksDA/jixUZUv2ZESNS48QEc/qduV4= +cloud.google.com/go/bigtable v1.36.0/go.mod h1:u98oqNAXiAufepkRGAd95lq2ap4kHGr3wLeFojvJwew= cloud.google.com/go/billing v1.4.0/go.mod h1:g9IdKBEFlItS8bTtlrZdVLWSSdSyFUZKXNS02zKMOZY= cloud.google.com/go/billing v1.5.0/go.mod h1:mztb1tBc3QekhjSgmpf/CV4LzWXLzCArwpLmP2Gm88s= cloud.google.com/go/billing v1.6.0/go.mod h1:WoXzguj+BeHXPbKfNWkqVtDdzORazmCjraY+vrxcyvI= @@ -819,6 +821,8 @@ github.com/golang/snappy v0.0.3/go.mod h1:/XxbfmMg8lxefKM7IXC3fBNl/7bRcc72aCRzEW github.com/golang/snappy v0.0.4/go.mod h1:/XxbfmMg8lxefKM7IXC3fBNl/7bRcc72aCRzEWrmP2Q= github.com/google/btree v0.0.0-20180813153112-4030bb1f1f0c/go.mod h1:lNA+9X1NB3Zf8V7Ke586lFgjr2dZNuvo3lPJSGZ5JPQ= github.com/google/btree v1.0.0/go.mod h1:lNA+9X1NB3Zf8V7Ke586lFgjr2dZNuvo3lPJSGZ5JPQ= +github.com/google/btree v1.1.3 h1:CVpQJjYgC4VbzxeGVHfvZrv1ctoYCAI8vbl07Fcxlyg= +github.com/google/btree v1.1.3/go.mod h1:qOPhT0dTNdNzV6Z/lhRX0YXUafgPLFUh+gZMl761Gm4= github.com/google/flatbuffers v2.0.8+incompatible/go.mod h1:1AeVuKshWv4vARoZatz6mlQ0JxURH0Kv5+zNeJKJCa8= github.com/google/go-cmp v0.2.0/go.mod h1:oXzfMopK8JAjlY9xF4vHSVASa0yLyX7SntLO5aqRK0M= github.com/google/go-cmp v0.3.0/go.mod h1:8QqcDgzrUqlUb/G2PQTWiueGozuR1884gddMywk6iLU= @@ -1747,6 +1751,7 @@ modernc.org/strutil v1.1.3/go.mod h1:MEHNA7PdEnEwLvspRMtWTNnp2nnyvMfkimT1NKNAGbw modernc.org/tcl v1.13.1/go.mod h1:XOLfOwzhkljL4itZkK6T72ckMgvj0BDsnKNdZVUOecw= modernc.org/token v1.0.0/go.mod h1:UGzOrNV1mAFSEB63lOFHIpNRUVMvYTc6yu1SMY/XTDM= modernc.org/z v1.5.1/go.mod h1:eWFB510QWW5Th9YGZT81s+LwvaAs3Q2yr4sP0rmLkv8= +rsc.io/binaryregexp v0.2.0 h1:HfqmD5MEmC0zvwBuF187nq9mdnXjXsSivRiXN7SmRkE= rsc.io/binaryregexp v0.2.0/go.mod h1:qTv7/COck+e2FymRvadv62gMdZztPaShugOCi3I+8D8= rsc.io/pdf v0.1.1/go.mod h1:n8OzWcQ6Sp37PL01nO98y4iUCRdTGarVfzxY20ICaU4= rsc.io/quote/v3 v3.1.0/go.mod h1:yEA65RcK8LyAZtP9Kv3t0HmxON59tX3rD+tICJqUlj0= diff --git a/internal/server/config.go b/internal/server/config.go index fc4d17ee6c..308a4fa7d0 100644 --- a/internal/server/config.go +++ b/internal/server/config.go @@ -23,6 +23,7 @@ import ( "github.com/googleapis/genai-toolbox/internal/auth/google" "github.com/googleapis/genai-toolbox/internal/sources" alloydbpgsrc "github.com/googleapis/genai-toolbox/internal/sources/alloydbpg" + bigtablesrc "github.com/googleapis/genai-toolbox/internal/sources/bigtable" cloudsqlmssqlsrc "github.com/googleapis/genai-toolbox/internal/sources/cloudsqlmssql" cloudsqlmysqlsrc "github.com/googleapis/genai-toolbox/internal/sources/cloudsqlmysql" cloudsqlpgsrc "github.com/googleapis/genai-toolbox/internal/sources/cloudsqlpg" @@ -35,6 +36,7 @@ import ( spannersrc "github.com/googleapis/genai-toolbox/internal/sources/spanner" "github.com/googleapis/genai-toolbox/internal/tools" "github.com/googleapis/genai-toolbox/internal/tools/alloydbainl" + "github.com/googleapis/genai-toolbox/internal/tools/bigtable" "github.com/googleapis/genai-toolbox/internal/tools/dgraph" httptool "github.com/googleapis/genai-toolbox/internal/tools/http" "github.com/googleapis/genai-toolbox/internal/tools/mssqlsql" @@ -161,6 +163,12 @@ func (c *SourceConfigs) UnmarshalYAML(ctx context.Context, unmarshal func(interf return fmt.Errorf("unable to parse as %q: %w", kind, err) } (*c)[name] = actual + case bigtablesrc.SourceKind: + actual := bigtablesrc.Config{Name: name} + if err := dec.DecodeContext(ctx, &actual); err != nil { + return fmt.Errorf("unable to parse as %q: %w", kind, err) + } + (*c)[name] = actual case cloudsqlpgsrc.SourceKind: actual := cloudsqlpgsrc.Config{Name: name, IPType: "public"} if err := dec.DecodeContext(ctx, &actual); err != nil { @@ -302,6 +310,12 @@ func (c *ToolConfigs) UnmarshalYAML(ctx context.Context, unmarshal func(interfac return fmt.Errorf("error creating decoder: %w", err) } switch kind { + case bigtable.ToolKind: + actual := bigtable.Config{Name: name} + if err := dec.DecodeContext(ctx, &actual); err != nil { + return fmt.Errorf("unable to parse as %q: %w", kind, err) + } + (*c)[name] = actual case postgressql.ToolKind: actual := postgressql.Config{Name: name} if err := dec.DecodeContext(ctx, &actual); err != nil { diff --git a/internal/sources/bigtable/bigtable.go b/internal/sources/bigtable/bigtable.go new file mode 100644 index 0000000000..51f4fea6a1 --- /dev/null +++ b/internal/sources/bigtable/bigtable.go @@ -0,0 +1,87 @@ +// 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 bigtable + +import ( + "context" + "fmt" + + "cloud.google.com/go/bigtable" + "github.com/googleapis/genai-toolbox/internal/sources" + "go.opentelemetry.io/otel/trace" + "google.golang.org/api/option" +) + +const SourceKind string = "bigtable" + +// validate interface +var _ sources.SourceConfig = Config{} + +type Config struct { + Name string `yaml:"name" validate:"required"` + Kind string `yaml:"kind" validate:"required"` + Project string `yaml:"project" validate:"required"` + Instance string `yaml:"instance" validate:"required"` +} + +func (r Config) SourceConfigKind() string { + return SourceKind +} + +func (r Config) Initialize(ctx context.Context, tracer trace.Tracer) (sources.Source, error) { + client, err := initBigtableClient(ctx, tracer, r.Name, r.Project, r.Instance) + if err != nil { + return nil, fmt.Errorf("unable to create client: %w", err) + } + + s := &Source{ + Name: r.Name, + Kind: SourceKind, + Client: client, + } + return s, nil +} + +var _ sources.Source = &Source{} + +type Source struct { + Name string `yaml:"name"` + Kind string `yaml:"kind"` + Client *bigtable.Client +} + +func (s *Source) SourceKind() string { + return SourceKind +} + +func (s *Source) BigtableClient() *bigtable.Client { + return s.Client +} + +func initBigtableClient(ctx context.Context, tracer trace.Tracer, name, project, instance string) (*bigtable.Client, error) { + //nolint:all // Reassigned ctx + ctx, span := sources.InitConnectionSpan(ctx, tracer, SourceKind, name) + defer span.End() + + // Set up Bigtable data operations client. + poolSize := 10 + client, err := bigtable.NewClient(ctx, project, instance, option.WithGRPCConnectionPool(poolSize)) + + if err != nil { + return nil, fmt.Errorf("unable to create bigtable.NewClient: %w", err) + } + + return client, nil +} diff --git a/internal/sources/bigtable/bigtable_test.go b/internal/sources/bigtable/bigtable_test.go new file mode 100644 index 0000000000..11a3a2981a --- /dev/null +++ b/internal/sources/bigtable/bigtable_test.go @@ -0,0 +1,116 @@ +// 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 bigtable_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" + "github.com/googleapis/genai-toolbox/internal/sources/bigtable" + "github.com/googleapis/genai-toolbox/internal/testutils" +) + +func TestParseFromYamlBigtableDb(t *testing.T) { + tcs := []struct { + desc string + in string + want server.SourceConfigs + }{ + { + desc: "can configure with a bigtable table", + in: ` + sources: + my-bigtable-instance: + kind: bigtable + project: my-project + instance: my-instance + `, + want: map[string]sources.SourceConfig{ + "my-bigtable-instance": bigtable.Config{ + Name: "my-bigtable-instance", + Kind: bigtable.SourceKind, + Project: "my-project", + Instance: "my-instance", + }, + }, + }, + } + for _, tc := range tcs { + t.Run(tc.desc, func(t *testing.T) { + got := struct { + Sources server.SourceConfigs `yaml:"sources"` + }{} + // Parse contents + err := yaml.Unmarshal(testutils.FormatYaml(tc.in), &got) + if err != nil { + t.Fatalf("unable to unmarshal: %s", err) + } + if !cmp.Equal(tc.want, got.Sources) { + t.Fatalf("incorrect parse: want %v, got %v", tc.want, got.Sources) + } + }) + } + +} + +func TestFailParseFromYaml(t *testing.T) { + tcs := []struct { + desc string + in string + err string + }{ + { + desc: "extra field", + in: ` + sources: + my-bigtable-instance: + kind: bigtable + project: my-project + instance: my-instance + foo: bar + `, + err: "unable to parse as \"bigtable\": [1:1] unknown field \"foo\"\n> 1 | foo: bar\n ^\n 2 | instance: my-instance\n 3 | kind: bigtable\n 4 | project: my-project", + }, + { + desc: "missing required field", + in: ` + sources: + my-bigtable-instance: + kind: bigtable + project: my-project + `, + err: "unable to parse as \"bigtable\": Key: 'Config.Instance' Error:Field validation for 'Instance' failed on the 'required' tag", + }, + } + for _, tc := range tcs { + t.Run(tc.desc, func(t *testing.T) { + got := struct { + Sources server.SourceConfigs `yaml:"sources"` + }{} + // Parse contents + err := yaml.Unmarshal(testutils.FormatYaml(tc.in), &got) + if err == nil { + t.Fatalf("expect parsing to fail") + } + errStr := err.Error() + if errStr != tc.err { + t.Fatalf("unexpected error: got %q, want %q", errStr, tc.err) + } + }) + } +} diff --git a/internal/tools/bigtable/bigtable.go b/internal/tools/bigtable/bigtable.go new file mode 100644 index 0000000000..565798a9b1 --- /dev/null +++ b/internal/tools/bigtable/bigtable.go @@ -0,0 +1,184 @@ +// 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 bigtable + +import ( + "context" + "fmt" + + "cloud.google.com/go/bigtable" + "github.com/googleapis/genai-toolbox/internal/sources" + bigtabledb "github.com/googleapis/genai-toolbox/internal/sources/bigtable" + "github.com/googleapis/genai-toolbox/internal/tools" +) + +const ToolKind string = "bigtable-sql" + +type compatibleSource interface { + BigtableClient() *bigtable.Client +} + +// validate compatible sources are still compatible +var _ compatibleSource = &bigtabledb.Source{} + +var compatibleSources = [...]string{bigtabledb.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"` + Statement string `yaml:"statement" validate:"required"` + AuthRequired []string `yaml:"authRequired"` + Parameters tools.Parameters `yaml:"parameters"` +} + +// validate interface +var _ tools.ToolConfig = Config{} + +func (cfg Config) ToolConfigKind() string { + return ToolKind +} + +func (cfg Config) Initialize(srcs map[string]sources.Source) (tools.Tool, error) { + // verify source exists + rawS, ok := srcs[cfg.Source] + if !ok { + return nil, fmt.Errorf("no source named %q configured", cfg.Source) + } + + // verify the source is compatible + s, ok := rawS.(compatibleSource) + if !ok { + return nil, fmt.Errorf("invalid source for %q tool: source kind must be one of %q", ToolKind, compatibleSources) + } + + mcpManifest := tools.McpManifest{ + Name: cfg.Name, + Description: cfg.Description, + InputSchema: cfg.Parameters.McpManifest(), + } + + // finish tool setup + t := Tool{ + Name: cfg.Name, + Kind: ToolKind, + Parameters: cfg.Parameters, + Statement: cfg.Statement, + AuthRequired: cfg.AuthRequired, + Client: s.BigtableClient(), + manifest: tools.Manifest{Description: cfg.Description, Parameters: cfg.Parameters.Manifest()}, + mcpManifest: mcpManifest, + } + return t, nil +} + +// validate interface +var _ tools.Tool = Tool{} + +type Tool struct { + Name string `yaml:"name"` + Kind string `yaml:"kind"` + AuthRequired []string `yaml:"authRequired"` + Parameters tools.Parameters `yaml:"parameters"` + + Client *bigtable.Client + Statement string + manifest tools.Manifest + mcpManifest tools.McpManifest +} + +func getMapParamsType(tparams tools.Parameters, params tools.ParamValues) (map[string]bigtable.SQLType, error) { + paramTypeMap := make(map[string]string) + for _, p := range tparams { + paramTypeMap[p.GetName()] = p.GetType() + } + + btParams := make(map[string]bigtable.SQLType) + for _, p := range params { + switch paramTypeMap[p.Name] { + case "boolean": + btParams[p.Name] = bigtable.BoolSQLType{} + case "string": + btParams[p.Name] = bigtable.StringSQLType{} + case "integer": + btParams[p.Name] = bigtable.Int64SQLType{} + case "float": + btParams[p.Name] = bigtable.Float64SQLType{} + case "array": + btParams[p.Name] = bigtable.ArraySQLType{} + } + } + + return btParams, nil +} + +func (t Tool) Invoke(params tools.ParamValues) ([]any, error) { + mapParamsType, err := getMapParamsType(t.Parameters, params) + if err != nil { + return nil, fmt.Errorf("fail to get map params: %w", err) + } + + ps, err := t.Client.PrepareStatement( + context.Background(), + t.Statement, + mapParamsType, + ) + if err != nil { + return nil, fmt.Errorf("unable to prepare statement: %w", err) + } + + bs, err := ps.Bind(params.AsMap()) + if err != nil { + return nil, fmt.Errorf("unable to bind: %w", err) + } + + var out []any + err = bs.Execute(context.Background(), func(resultRow bigtable.ResultRow) bool { + vMap := make(map[string]any) + cols := resultRow.Metadata.Columns + + for _, c := range cols { + var columValue any + err = resultRow.GetByName(c.Name, &columValue) + vMap[c.Name] = columValue + } + + out = append(out, vMap) + + return true + }) + if err != nil { + return nil, fmt.Errorf("unable to execute client: %w", err) + } + + return out, nil +} + +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) +} diff --git a/internal/tools/bigtable/bigtable_test.go b/internal/tools/bigtable/bigtable_test.go new file mode 100644 index 0000000000..3205845c93 --- /dev/null +++ b/internal/tools/bigtable/bigtable_test.go @@ -0,0 +1,83 @@ +// 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 bigtable_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/testutils" + "github.com/googleapis/genai-toolbox/internal/tools" + "github.com/googleapis/genai-toolbox/internal/tools/bigtable" +) + +func TestParseFromYamlBigtable(t *testing.T) { + ctx, err := testutils.ContextWithNewLogger() + if err != nil { + t.Fatalf("unexpected error: %s", err) + } + tcs := []struct { + desc string + in string + want server.ToolConfigs + }{ + { + desc: "basic example", + in: ` + tools: + example_tool: + kind: bigtable-sql + source: my-pg-instance + description: some description + statement: | + SELECT * FROM SQL_STATEMENT; + parameters: + - name: country + type: string + description: some description + `, + want: server.ToolConfigs{ + "example_tool": bigtable.Config{ + Name: "example_tool", + Kind: bigtable.ToolKind, + Source: "my-pg-instance", + Description: "some description", + Statement: "SELECT * FROM SQL_STATEMENT;\n", + Parameters: []tools.Parameter{ + tools.NewStringParameter("country", "some description"), + }, + }, + }, + }, + } + 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) + } + }) + } + +} diff --git a/tests/bigtable_integration_test.go b/tests/bigtable_integration_test.go new file mode 100644 index 0000000000..21391af608 --- /dev/null +++ b/tests/bigtable_integration_test.go @@ -0,0 +1,213 @@ +//go:build integration && bigtable + +// 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 tests + +import ( + "bytes" + "context" + "encoding/binary" + "fmt" + "log" + "os" + "regexp" + "slices" + "strings" + "testing" + "time" + + "cloud.google.com/go/bigtable" + "github.com/google/uuid" +) + +var ( + BIGTABLE_SOURCE_KIND = "bigtable" + BIGTABLE_TOOL_KIND = "bigtable-sql" + BIGTABLE_PROJECT = os.Getenv("BIGTABLE_PROJECT") + BIGTABLE_INSTANCE = os.Getenv("BIGTABLE_INSTANCE") +) + +func getBigtableVars(t *testing.T) map[string]any { + switch "" { + case BIGTABLE_PROJECT: + t.Fatal("'BIGTABLE_PROJECT' not set") + case BIGTABLE_INSTANCE: + t.Fatal("'BIGTABLE_INSTANCE' not set") + } + + return map[string]any{ + "kind": BIGTABLE_SOURCE_KIND, + "project": BIGTABLE_PROJECT, + "instance": BIGTABLE_INSTANCE, + } +} + +type TestRow struct { + RowKey string + ColumnName string + Data []byte +} + +func TestBigtableToolEndpoints(t *testing.T) { + sourceConfig := getBigtableVars(t) + ctx, cancel := context.WithTimeout(context.Background(), time.Minute) + defer cancel() + + var args []string + + tableName := "param_table" + strings.Replace(uuid.New().String(), "-", "", -1) + tableNameAuth := "auth_table_" + strings.Replace(uuid.New().String(), "-", "", -1) + + columnFamilyName := "cf" + muts, rowKeys := getTestData(columnFamilyName) + + // Do not change the shape of statement without checking tests/common_test.go. + // The structure and value of seed data has to match https://github.com/googleapis/genai-toolbox/blob/4dba0df12dc438eca3cb476ef52aa17cdf232c12/tests/common_test.go#L200-L251 + param_test_statement := fmt.Sprintf("SELECT TO_INT64(cf['id']) as id, CAST(cf['name'] AS string) as name, FROM %s WHERE TO_INT64(cf['id']) = @id OR CAST(cf['name'] AS string) = @name;", tableName) + teardownTable1 := SetupBtTable(t, ctx, sourceConfig["project"].(string), sourceConfig["instance"].(string), tableName, columnFamilyName, muts, rowKeys) + defer teardownTable1(t) + + // Do not change the shape of statement without checking tests/common_test.go. + // The structure and value of seed data has to match https://github.com/googleapis/genai-toolbox/blob/4dba0df12dc438eca3cb476ef52aa17cdf232c12/tests/common_test.go#L200-L251 + auth_tool_statement := fmt.Sprintf("SELECT CAST(cf['name'] AS string) as name FROM %s WHERE CAST(cf['email'] AS string) = @email;", tableNameAuth) + teardownTable2 := SetupBtTable(t, ctx, sourceConfig["project"].(string), sourceConfig["instance"].(string), tableNameAuth, columnFamilyName, muts, rowKeys) + defer teardownTable2(t) + + // Write config into a file and pass it to command + toolsFile := GetToolsConfig(sourceConfig, BIGTABLE_TOOL_KIND, param_test_statement, auth_tool_statement) + cmd, cleanup, err := StartCmd(ctx, toolsFile, args...) + if err != nil { + t.Fatalf("command initialization returned an error: %s", err) + } + defer cleanup() + + waitCtx, cancel := context.WithTimeout(ctx, 10*time.Second) + defer cancel() + out, err := cmd.WaitForString(waitCtx, regexp.MustCompile(`Server ready to serve`)) + if err != nil { + t.Logf("toolbox command logs: \n%s", out) + t.Fatalf("toolbox didn't start successfully: %s", err) + } + + RunToolGetTest(t) + + // Actual test parameters are set in https://github.com/googleapis/genai-toolbox/blob/52b09a67cb40ac0c5f461598b4673136699a3089/tests/tool_test.go#L250 + select_1_want := "[{$col1:1}]" + RunToolInvokeTest(t, select_1_want) +} + +func getTestData(columnFamilyName string) ([]*bigtable.Mutation, []string) { + muts := []*bigtable.Mutation{} + rowKeys := []string{} + + var ids [3][]byte + for i := range ids { + binary1 := new(bytes.Buffer) + if err := binary.Write(binary1, binary.BigEndian, int64(i+1)); err != nil { + log.Fatalf("Unable to encode id: %v", err) + } + ids[i] = binary1.Bytes() + } + + now := bigtable.Time(time.Now()) + for rowKey, mutData := range map[string]map[string][]byte{ + // Do not change the test data without checking tests/common_test.go. + // The structure and value of seed data has to match https://github.com/googleapis/genai-toolbox/blob/4dba0df12dc438eca3cb476ef52aa17cdf232c12/tests/common_test.go#L200-L251 + // Expected values are defined in https://github.com/googleapis/genai-toolbox/blob/52b09a67cb40ac0c5f461598b4673136699a3089/tests/tool_test.go#L229-L310 + "row-01": { + "name": []byte("Alice"), + "email": []byte(SERVICE_ACCOUNT_EMAIL), + "id": ids[0], + }, + "row-02": { + "name": []byte("Jane"), + "email": []byte("janedoe@gmail.com"), + "id": ids[1], + }, + "row-03": { + "name": []byte("Sid"), + "id": ids[2], + }, + } { + mut := bigtable.NewMutation() + for col, v := range mutData { + mut.Set(columnFamilyName, col, now, v) + } + muts = append(muts, mut) + rowKeys = append(rowKeys, rowKey) + } + return muts, rowKeys +} + +func SetupBtTable(t *testing.T, ctx context.Context, projectId string, instance string, tableName string, columnFamilyName string, muts []*bigtable.Mutation, rowKeys []string) func(*testing.T) { + // Creating clients + adminClient, err := bigtable.NewAdminClient(ctx, projectId, instance) + if err != nil { + t.Fatalf("NewAdminClient: %v", err) + } + + client, err := bigtable.NewClient(ctx, projectId, instance) + if err != nil { + log.Fatalf("Could not create data operations client: %v", err) + } + defer client.Close() + + // Creating tables + tables, err := adminClient.Tables(ctx) + if err != nil { + log.Fatalf("Could not fetch table list: %v", err) + } + + if !slices.Contains(tables, tableName) { + log.Printf("Creating table %s", tableName) + if err := adminClient.CreateTable(ctx, tableName); err != nil { + log.Fatalf("Could not create table %s: %v", tableName, err) + } + } + + tblInfo, err := adminClient.TableInfo(ctx, tableName) + if err != nil { + log.Fatalf("Could not read info for table %s: %v", tableName, err) + } + + // Creating column family + if !slices.Contains(tblInfo.Families, columnFamilyName) { + if err := adminClient.CreateColumnFamily(ctx, tableName, columnFamilyName); err != nil { + log.Fatalf("Could not create column family %s: %v", columnFamilyName, err) + } + } + + tbl := client.Open(tableName) + rowErrs, err := tbl.ApplyBulk(ctx, rowKeys, muts) + if err != nil { + log.Fatalf("Could not apply bulk row mutation: %v", err) + } + if rowErrs != nil { + for _, rowErr := range rowErrs { + log.Printf("Error writing row: %v", rowErr) + } + log.Fatalf("Could not write some rows") + } + + // Writing data + return func(t *testing.T) { + // tear down test + if err = adminClient.DeleteTable(ctx, tableName); err != nil { + log.Fatalf("Teardown failed. Could not delete table %s: %v", tableName, err) + } + defer adminClient.Close() + } +}