mirror of
https://github.com/googleapis/genai-toolbox.git
synced 2026-01-09 15:38:08 -05:00
feat: Add Bigtable source and tool (#418)
# Add Bigtable support
A `bigtable` source can be added as the following example
```
sources:
test-bigtable-source:
kind: "bigtable"
project: "sample-project"
instance: "sample-instance"
```
A `bigtable` tool can be added as below
```
tools:
get-test-tool-data:
kind: bigtable-sql
source: test-bigtable-source
description: Some description
statement: SELECT * FROM `test-table` WHERE address['state'] = @state;
parameters:
- name: state
type: string
description: Filter by state
```
---------
Co-authored-by: Yuan <45984206+Yuan325@users.noreply.github.com>
Co-authored-by: Kurtis Van Gent <31518063+kurtisvg@users.noreply.github.com>
Co-authored-by: Wenxin Du <117315983+duwenxin99@users.noreply.github.com>
This commit is contained in:
An Nguyen
@@ -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"
|
||||
|
||||
@@ -45,3 +45,4 @@ run:
|
||||
- mysql
|
||||
- http
|
||||
- alloydb_ai_nl
|
||||
- bigtable
|
||||
|
||||
70
docs/en/resources/sources/bigtable.md
Normal file
70
docs/en/resources/sources/bigtable.md
Normal file
@@ -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. |
|
||||
82
docs/en/resources/tools/bigtable-sql.md
Normal file
82
docs/en/resources/tools/bigtable-sql.md
Normal file
@@ -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
|
||||
1
go.mod
1
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
|
||||
|
||||
5
go.sum
5
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=
|
||||
|
||||
@@ -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 {
|
||||
|
||||
87
internal/sources/bigtable/bigtable.go
Normal file
87
internal/sources/bigtable/bigtable.go
Normal file
@@ -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
|
||||
}
|
||||
116
internal/sources/bigtable/bigtable_test.go
Normal file
116
internal/sources/bigtable/bigtable_test.go
Normal file
@@ -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)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
184
internal/tools/bigtable/bigtable.go
Normal file
184
internal/tools/bigtable/bigtable.go
Normal file
@@ -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)
|
||||
}
|
||||
83
internal/tools/bigtable/bigtable_test.go
Normal file
83
internal/tools/bigtable/bigtable_test.go
Normal file
@@ -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)
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
}
|
||||
213
tests/bigtable_integration_test.go
Normal file
213
tests/bigtable_integration_test.go
Normal file
@@ -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()
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user