mirror of
https://github.com/googleapis/genai-toolbox.git
synced 2026-01-11 16:38:15 -05:00
Compare commits
2 Commits
spanner-cr
...
main
| Author | SHA1 | Date | |
|---|---|---|---|
|
4abf0c39e7 | ||
|
|
dd7b9de623 |
@@ -340,26 +340,6 @@ steps:
|
||||
spanner \
|
||||
spanner || echo "Integration tests failed." # ignore test failures
|
||||
|
||||
- id: "spanner-admin"
|
||||
name: golang:1
|
||||
waitFor: ["compile-test-binary"]
|
||||
entrypoint: /bin/bash
|
||||
env:
|
||||
- "GOPATH=/gopath"
|
||||
- "SPANNER_PROJECT=$PROJECT_ID"
|
||||
- "SERVICE_ACCOUNT_EMAIL=$SERVICE_ACCOUNT_EMAIL"
|
||||
secretEnv: ["CLIENT_ID"]
|
||||
volumes:
|
||||
- name: "go"
|
||||
path: "/gopath"
|
||||
args:
|
||||
- -c
|
||||
- |
|
||||
.ci/test_with_coverage.sh \
|
||||
"Spanner Admin" \
|
||||
spanneradmin \
|
||||
spanneradmin || echo "Integration tests failed."
|
||||
|
||||
- id: "neo4j"
|
||||
name: golang:1
|
||||
waitFor: ["compile-test-binary"]
|
||||
|
||||
17
DEVELOPER.md
17
DEVELOPER.md
@@ -379,6 +379,23 @@ to approve PRs for main. TeamSync is used to create this team from the MDB
|
||||
Group `toolbox-contributors`. Googlers who are developing for MCP-Toolbox
|
||||
but aren't part of the core team should join this group.
|
||||
|
||||
### Issue/PR Triage and SLO
|
||||
After an issue is created, maintainers will assign the following labels:
|
||||
* `Priority` (defaulted to P0)
|
||||
* `Type` (if applicable)
|
||||
* `Product` (if applicable)
|
||||
|
||||
All incoming issues and PRs will follow the following SLO:
|
||||
| Type | Priority | Objective |
|
||||
|-----------------|----------|------------------------------------------------------------------------|
|
||||
| Feature Request | P0 | Must respond within **5 days** |
|
||||
| Process | P0 | Must respond within **5 days** |
|
||||
| Bugs | P0 | Must respond within **5 days**, and resolve/closure within **14 days** |
|
||||
| Bugs | P1 | Must respond within **7 days**, and resolve/closure within **90 days** |
|
||||
| Bugs | P2 | Must respond within **30 days**
|
||||
|
||||
_Types that are not listed in the table do not adhere to any SLO._
|
||||
|
||||
### Releasing
|
||||
|
||||
Toolbox has two types of releases: versioned and continuous. It uses Google
|
||||
|
||||
@@ -221,7 +221,6 @@ import (
|
||||
_ "github.com/googleapis/genai-toolbox/internal/tools/spanner/spannerlistgraphs"
|
||||
_ "github.com/googleapis/genai-toolbox/internal/tools/spanner/spannerlisttables"
|
||||
_ "github.com/googleapis/genai-toolbox/internal/tools/spanner/spannersql"
|
||||
_ "github.com/googleapis/genai-toolbox/internal/tools/spanneradmin/spannercreateinstance"
|
||||
_ "github.com/googleapis/genai-toolbox/internal/tools/sqlite/sqliteexecutesql"
|
||||
_ "github.com/googleapis/genai-toolbox/internal/tools/sqlite/sqlitesql"
|
||||
_ "github.com/googleapis/genai-toolbox/internal/tools/tidb/tidbexecutesql"
|
||||
@@ -268,7 +267,6 @@ import (
|
||||
_ "github.com/googleapis/genai-toolbox/internal/sources/singlestore"
|
||||
_ "github.com/googleapis/genai-toolbox/internal/sources/snowflake"
|
||||
_ "github.com/googleapis/genai-toolbox/internal/sources/spanner"
|
||||
_ "github.com/googleapis/genai-toolbox/internal/sources/spanneradmin"
|
||||
_ "github.com/googleapis/genai-toolbox/internal/sources/sqlite"
|
||||
_ "github.com/googleapis/genai-toolbox/internal/sources/tidb"
|
||||
_ "github.com/googleapis/genai-toolbox/internal/sources/trino"
|
||||
|
||||
@@ -1,59 +0,0 @@
|
||||
# Cloud Spanner Admin MCP Server
|
||||
|
||||
The Cloud Spanner Admin Model Context Protocol (MCP) Server gives AI-powered development tools the ability to manage your Google Cloud Spanner infrastructure. It supports creating instances.
|
||||
|
||||
## Features
|
||||
|
||||
An editor configured to use the Cloud Spanner Admin MCP server can use its AI capabilities to help you:
|
||||
|
||||
- **Provision & Manage Infrastructure** - Create Cloud Spanner instances
|
||||
|
||||
## Prerequisites
|
||||
|
||||
* [Node.js](https://nodejs.org/) installed.
|
||||
* A Google Cloud project with the **Cloud Spanner Admin API** enabled.
|
||||
* Ensure [Application Default Credentials](https://cloud.google.com/docs/authentication/gcloud) are available in your environment.
|
||||
* IAM Permissions:
|
||||
* Cloud Spanner Admin (`roles/spanner.admin`)
|
||||
|
||||
## Install & Configuration
|
||||
|
||||
In the Antigravity MCP Store, click the "Install" button.
|
||||
|
||||
You'll now be able to see all enabled tools in the "Tools" tab.
|
||||
|
||||
> [!NOTE]
|
||||
> If you encounter issues with Windows Defender blocking the execution, you may need to configure an allowlist. See [Configure exclusions for Microsoft Defender Antivirus](https://learn.microsoft.com/en-us/microsoft-365/security/defender-endpoint/configure-exclusions-microsoft-defender-antivirus?view=o365-worldwide) for more details.
|
||||
|
||||
## Usage
|
||||
|
||||
Once configured, the MCP server will automatically provide Cloud Spanner Admin capabilities to your AI assistant. You can:
|
||||
|
||||
* "Create a new Spanner instance named 'my-spanner-instance' in the 'my-gcp-project' project with config 'regional-us-central1', edition 'ENTERPRISE', and 1 node."
|
||||
|
||||
## Server Capabilities
|
||||
|
||||
The Cloud Spanner Admin MCP server provides the following tools:
|
||||
|
||||
| Tool Name | Description |
|
||||
|:------------------|:---------------------------------|
|
||||
| `create_instance` | Create a Cloud Spanner instance. |
|
||||
|
||||
## Custom MCP Server Configuration
|
||||
|
||||
Add the following configuration to your MCP client (e.g., `settings.json` for Gemini CLI, `mcp_config.json` for Antigravity):
|
||||
|
||||
```json
|
||||
{
|
||||
"mcpServers": {
|
||||
"spanner-admin": {
|
||||
"command": "npx",
|
||||
"args": ["-y", "@toolbox-sdk/server", "--prebuilt", "spanner-admin", "--stdio"]
|
||||
}
|
||||
}
|
||||
}
|
||||
```
|
||||
|
||||
## Documentation
|
||||
|
||||
For more information, visit the [Cloud Spanner Admin API documentation](https://cloud.google.com/spanner/docs/reference/rpc/google.spanner.admin.instance.v1).
|
||||
@@ -134,6 +134,7 @@ sources:
|
||||
# scopes: # Optional: List of OAuth scopes to request.
|
||||
# - "https://www.googleapis.com/auth/bigquery"
|
||||
# - "https://www.googleapis.com/auth/drive.readonly"
|
||||
# maxQueryResultRows: 50 # Optional: Limits the number of rows returned by queries. Defaults to 50.
|
||||
```
|
||||
|
||||
Initialize a BigQuery source that uses the client's access token:
|
||||
@@ -153,6 +154,7 @@ sources:
|
||||
# scopes: # Optional: List of OAuth scopes to request.
|
||||
# - "https://www.googleapis.com/auth/bigquery"
|
||||
# - "https://www.googleapis.com/auth/drive.readonly"
|
||||
# maxQueryResultRows: 50 # Optional: Limits the number of rows returned by queries. Defaults to 50.
|
||||
```
|
||||
|
||||
## Reference
|
||||
@@ -167,3 +169,4 @@ sources:
|
||||
| useClientOAuth | bool | false | If true, forwards the client's OAuth access token from the "Authorization" header to downstream queries. **Note:** This cannot be used with `writeMode: protected`. |
|
||||
| scopes | []string | false | A list of OAuth 2.0 scopes to use for the credentials. If not provided, default scopes are used. |
|
||||
| impersonateServiceAccount | string | false | Service account email to impersonate when making BigQuery and Dataplex API calls. The authenticated principal must have the `roles/iam.serviceAccountTokenCreator` role on the target service account. [Learn More](https://cloud.google.com/iam/docs/service-account-impersonation) |
|
||||
| maxQueryResultRows | int | false | The maximum number of rows to return from a query. Defaults to 50. |
|
||||
|
||||
@@ -1,42 +0,0 @@
|
||||
---
|
||||
title: Spanner Admin
|
||||
type: docs
|
||||
weight: 1
|
||||
description: "A \"spanner-admin\" source provides a client for the Cloud Spanner Admin API.\n"
|
||||
alias: [/resources/sources/spanner-admin]
|
||||
---
|
||||
|
||||
## About
|
||||
|
||||
The `spanner-admin` source provides a client to interact with the [Google
|
||||
Cloud Spanner Admin API](https://cloud.google.com/spanner/docs/reference/rpc/google.spanner.admin.instance.v1). This
|
||||
allows tools to perform administrative tasks on Spanner instances, such as
|
||||
creating instances.
|
||||
|
||||
Authentication can be handled in two ways:
|
||||
|
||||
1. **Application Default Credentials (ADC):** By default, the source uses ADC
|
||||
to authenticate with the API.
|
||||
2. **Client-side OAuth:** If `useClientOAuth` is set to `true`, the source will
|
||||
expect an OAuth 2.0 access token to be provided by the client (e.g., a web
|
||||
browser) for each request.
|
||||
|
||||
## Example
|
||||
|
||||
```yaml
|
||||
sources:
|
||||
my-spanner-admin:
|
||||
kind: spanner-admin
|
||||
|
||||
my-oauth-spanner-admin:
|
||||
kind: spanner-admin
|
||||
useClientOAuth: true
|
||||
```
|
||||
|
||||
## Reference
|
||||
|
||||
| **field** | **type** | **required** | **description** |
|
||||
| -------------- | :------: | :----------: | ---------------------------------------------------------------------------------------------------------------------------------------------- |
|
||||
| kind | string | true | Must be "spanner-admin". |
|
||||
| defaultProject | string | false | The Google Cloud project ID to use for Spanner infrastructure tools. |
|
||||
| useClientOAuth | boolean | false | If true, the source will use client-side OAuth for authorization. Otherwise, it will use Application Default Credentials. Defaults to `false`. |
|
||||
@@ -1,52 +0,0 @@
|
||||
---
|
||||
title: spanner-create-instance
|
||||
type: docs
|
||||
weight: 2
|
||||
description: "Create a Cloud Spanner instance."
|
||||
---
|
||||
|
||||
The `spanner-create-instance` tool creates a new Cloud Spanner instance in a
|
||||
specified Google Cloud project.
|
||||
|
||||
{{< notice info >}}
|
||||
This tool uses the `spanner-admin` source.
|
||||
{{< /notice >}}
|
||||
|
||||
## Configuration
|
||||
|
||||
Here is an example of how to configure the `spanner-create-instance` tool in
|
||||
your `tools.yaml` file:
|
||||
|
||||
```yaml
|
||||
sources:
|
||||
my-spanner-admin-source:
|
||||
kind: spanner-admin
|
||||
|
||||
tools:
|
||||
create_my_spanner_instance:
|
||||
kind: spanner-create-instance
|
||||
source: my-spanner-admin-source
|
||||
description: "Creates a Spanner instance."
|
||||
```
|
||||
|
||||
## Parameters
|
||||
|
||||
The `spanner-create-instance` tool has the following parameters:
|
||||
|
||||
| **field** | **type** | **required** | **description** |
|
||||
| --------------- | :------: | :----------: | ------------------------------------------------------------------------------------ |
|
||||
| project | string | true | The Google Cloud project ID. |
|
||||
| instanceId | string | true | The ID of the instance to create. |
|
||||
| displayName | string | true | The display name of the instance. |
|
||||
| config | string | true | The instance configuration (e.g., `regional-us-central1`). |
|
||||
| nodeCount | integer | true | The number of nodes. Mutually exclusive with `processingUnits` (one must be 0). |
|
||||
| processingUnits | integer | true | The number of processing units. Mutually exclusive with `nodeCount` (one must be 0). |
|
||||
| edition | string | false | The edition of the instance (`STANDARD`, `ENTERPRISE`, `ENTERPRISE_PLUS`). |
|
||||
|
||||
## Reference
|
||||
|
||||
| **field** | **type** | **required** | **description** |
|
||||
| ----------- | :------: | :----------: | ------------------------------------------------------------ |
|
||||
| kind | string | true | Must be `spanner-create-instance`. |
|
||||
| source | string | true | The name of the `spanner-admin` source to use for this tool. |
|
||||
| description | string | false | A description of the tool that is passed to the agent. |
|
||||
@@ -50,7 +50,6 @@ var expectedToolSources = []string{
|
||||
"serverless-spark",
|
||||
"singlestore",
|
||||
"snowflake",
|
||||
"spanner-admin",
|
||||
"spanner-postgres",
|
||||
"spanner",
|
||||
"sqlite",
|
||||
|
||||
@@ -19,6 +19,7 @@ sources:
|
||||
location: ${BIGQUERY_LOCATION:}
|
||||
useClientOAuth: ${BIGQUERY_USE_CLIENT_OAUTH:false}
|
||||
scopes: ${BIGQUERY_SCOPES:}
|
||||
maxQueryResultRows: ${BIGQUERY_MAX_QUERY_RESULT_ROWS:50}
|
||||
|
||||
tools:
|
||||
analyze_contribution:
|
||||
|
||||
@@ -1,27 +0,0 @@
|
||||
# Copyright 2026 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:
|
||||
spanner-admin-source:
|
||||
kind: spanner-admin
|
||||
defaultProject: ${SPANNER_PROJECT:}
|
||||
|
||||
tools:
|
||||
create_instance:
|
||||
kind: spanner-create-instance
|
||||
source: spanner-admin-source
|
||||
|
||||
toolsets:
|
||||
spanner_admin_tools:
|
||||
- create_instance
|
||||
@@ -89,6 +89,7 @@ type Config struct {
|
||||
UseClientOAuth bool `yaml:"useClientOAuth"`
|
||||
ImpersonateServiceAccount string `yaml:"impersonateServiceAccount"`
|
||||
Scopes StringOrStringSlice `yaml:"scopes"`
|
||||
MaxQueryResultRows int `yaml:"maxQueryResultRows"`
|
||||
}
|
||||
|
||||
// StringOrStringSlice is a custom type that can unmarshal both a single string
|
||||
@@ -127,6 +128,10 @@ func (r Config) Initialize(ctx context.Context, tracer trace.Tracer) (sources.So
|
||||
r.WriteMode = WriteModeAllowed
|
||||
}
|
||||
|
||||
if r.MaxQueryResultRows == 0 {
|
||||
r.MaxQueryResultRows = 50
|
||||
}
|
||||
|
||||
if r.WriteMode == WriteModeProtected && r.UseClientOAuth {
|
||||
// The protected mode only allows write operations to the session's temporary datasets.
|
||||
// when using client OAuth, a new session is created every
|
||||
@@ -150,7 +155,7 @@ func (r Config) Initialize(ctx context.Context, tracer trace.Tracer) (sources.So
|
||||
Client: client,
|
||||
RestService: restService,
|
||||
TokenSource: tokenSource,
|
||||
MaxQueryResultRows: 50,
|
||||
MaxQueryResultRows: r.MaxQueryResultRows,
|
||||
ClientCreator: clientCreator,
|
||||
}
|
||||
|
||||
@@ -567,7 +572,7 @@ func (s *Source) RunSQL(ctx context.Context, bqClient *bigqueryapi.Client, state
|
||||
}
|
||||
|
||||
var out []any
|
||||
for {
|
||||
for s.MaxQueryResultRows <= 0 || len(out) < s.MaxQueryResultRows {
|
||||
var val []bigqueryapi.Value
|
||||
err = it.Next(&val)
|
||||
if err == iterator.Done {
|
||||
|
||||
@@ -21,9 +21,12 @@ import (
|
||||
|
||||
yaml "github.com/goccy/go-yaml"
|
||||
"github.com/google/go-cmp/cmp"
|
||||
"go.opentelemetry.io/otel/trace/noop"
|
||||
|
||||
"github.com/googleapis/genai-toolbox/internal/server"
|
||||
"github.com/googleapis/genai-toolbox/internal/sources/bigquery"
|
||||
"github.com/googleapis/genai-toolbox/internal/testutils"
|
||||
"github.com/googleapis/genai-toolbox/internal/util"
|
||||
)
|
||||
|
||||
func TestParseFromYamlBigQuery(t *testing.T) {
|
||||
@@ -154,6 +157,26 @@ func TestParseFromYamlBigQuery(t *testing.T) {
|
||||
},
|
||||
},
|
||||
},
|
||||
{
|
||||
desc: "with max query result rows example",
|
||||
in: `
|
||||
sources:
|
||||
my-instance:
|
||||
kind: bigquery
|
||||
project: my-project
|
||||
location: us
|
||||
maxQueryResultRows: 10
|
||||
`,
|
||||
want: server.SourceConfigs{
|
||||
"my-instance": bigquery.Config{
|
||||
Name: "my-instance",
|
||||
Kind: bigquery.SourceKind,
|
||||
Project: "my-project",
|
||||
Location: "us",
|
||||
MaxQueryResultRows: 10,
|
||||
},
|
||||
},
|
||||
},
|
||||
}
|
||||
for _, tc := range tcs {
|
||||
t.Run(tc.desc, func(t *testing.T) {
|
||||
@@ -220,6 +243,59 @@ func TestFailParseFromYaml(t *testing.T) {
|
||||
}
|
||||
}
|
||||
|
||||
func TestInitialize_MaxQueryResultRows(t *testing.T) {
|
||||
ctx, err := testutils.ContextWithNewLogger()
|
||||
if err != nil {
|
||||
t.Fatalf("unexpected error: %s", err)
|
||||
}
|
||||
ctx = util.WithUserAgent(ctx, "test-agent")
|
||||
tracer := noop.NewTracerProvider().Tracer("")
|
||||
|
||||
tcs := []struct {
|
||||
desc string
|
||||
cfg bigquery.Config
|
||||
want int
|
||||
}{
|
||||
{
|
||||
desc: "default value",
|
||||
cfg: bigquery.Config{
|
||||
Name: "test-default",
|
||||
Kind: bigquery.SourceKind,
|
||||
Project: "test-project",
|
||||
UseClientOAuth: true,
|
||||
},
|
||||
want: 50,
|
||||
},
|
||||
{
|
||||
desc: "configured value",
|
||||
cfg: bigquery.Config{
|
||||
Name: "test-configured",
|
||||
Kind: bigquery.SourceKind,
|
||||
Project: "test-project",
|
||||
UseClientOAuth: true,
|
||||
MaxQueryResultRows: 100,
|
||||
},
|
||||
want: 100,
|
||||
},
|
||||
}
|
||||
|
||||
for _, tc := range tcs {
|
||||
t.Run(tc.desc, func(t *testing.T) {
|
||||
src, err := tc.cfg.Initialize(ctx, tracer)
|
||||
if err != nil {
|
||||
t.Fatalf("Initialize failed: %v", err)
|
||||
}
|
||||
bqSrc, ok := src.(*bigquery.Source)
|
||||
if !ok {
|
||||
t.Fatalf("Expected *bigquery.Source, got %T", src)
|
||||
}
|
||||
if bqSrc.MaxQueryResultRows != tc.want {
|
||||
t.Errorf("MaxQueryResultRows = %d, want %d", bqSrc.MaxQueryResultRows, tc.want)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestNormalizeValue(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
|
||||
@@ -1,120 +0,0 @@
|
||||
// Copyright 2026 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 spanneradmin
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
|
||||
instance "cloud.google.com/go/spanner/admin/instance/apiv1"
|
||||
"github.com/goccy/go-yaml"
|
||||
"github.com/googleapis/genai-toolbox/internal/sources"
|
||||
"github.com/googleapis/genai-toolbox/internal/util"
|
||||
"go.opentelemetry.io/otel/trace"
|
||||
"golang.org/x/oauth2"
|
||||
"google.golang.org/api/option"
|
||||
)
|
||||
|
||||
const SourceKind string = "spanner-admin"
|
||||
|
||||
// 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"`
|
||||
DefaultProject string `yaml:"defaultProject"`
|
||||
UseClientOAuth bool `yaml:"useClientOAuth"`
|
||||
}
|
||||
|
||||
func (r Config) SourceConfigKind() string {
|
||||
return SourceKind
|
||||
}
|
||||
|
||||
// Initialize initializes a Spanner Admin Source instance.
|
||||
func (r Config) Initialize(ctx context.Context, tracer trace.Tracer) (sources.Source, error) {
|
||||
var client *instance.InstanceAdminClient
|
||||
|
||||
if !r.UseClientOAuth {
|
||||
ua, err := util.UserAgentFromContext(ctx)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("error in User Agent retrieval: %s", err)
|
||||
}
|
||||
// Use Application Default Credentials
|
||||
client, err = instance.NewInstanceAdminClient(ctx, option.WithUserAgent(ua))
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("error creating new spanner instance admin client: %w", err)
|
||||
}
|
||||
}
|
||||
|
||||
s := &Source{
|
||||
Config: r,
|
||||
Client: client,
|
||||
}
|
||||
return s, nil
|
||||
}
|
||||
|
||||
var _ sources.Source = &Source{}
|
||||
|
||||
type Source struct {
|
||||
Config
|
||||
Client *instance.InstanceAdminClient
|
||||
}
|
||||
|
||||
func (s *Source) SourceKind() string {
|
||||
return SourceKind
|
||||
}
|
||||
|
||||
func (s *Source) ToConfig() sources.SourceConfig {
|
||||
return s.Config
|
||||
}
|
||||
|
||||
func (s *Source) GetDefaultProject() string {
|
||||
return s.DefaultProject
|
||||
}
|
||||
|
||||
func (s *Source) GetClient(ctx context.Context, accessToken string) (*instance.InstanceAdminClient, error) {
|
||||
if s.UseClientOAuth {
|
||||
token := &oauth2.Token{AccessToken: accessToken}
|
||||
ua, err := util.UserAgentFromContext(ctx)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
client, err := instance.NewInstanceAdminClient(ctx, option.WithTokenSource(oauth2.StaticTokenSource(token)), option.WithUserAgent(ua))
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("error creating new spanner instance admin client: %w", err)
|
||||
}
|
||||
return client, nil
|
||||
}
|
||||
return s.Client, nil
|
||||
}
|
||||
|
||||
func (s *Source) UseClientAuthorization() bool {
|
||||
return s.UseClientOAuth
|
||||
}
|
||||
@@ -1,135 +0,0 @@
|
||||
// Copyright 2026 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 spanneradmin_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/spanneradmin"
|
||||
"github.com/googleapis/genai-toolbox/internal/testutils"
|
||||
)
|
||||
|
||||
func TestParseFromYamlSpannerAdmin(t *testing.T) {
|
||||
t.Parallel()
|
||||
tcs := []struct {
|
||||
desc string
|
||||
in string
|
||||
want server.SourceConfigs
|
||||
}{
|
||||
{
|
||||
desc: "basic example",
|
||||
in: `
|
||||
sources:
|
||||
my-spanner-admin-instance:
|
||||
kind: spanner-admin
|
||||
`,
|
||||
want: map[string]sources.SourceConfig{
|
||||
"my-spanner-admin-instance": spanneradmin.Config{
|
||||
Name: "my-spanner-admin-instance",
|
||||
Kind: spanneradmin.SourceKind,
|
||||
UseClientOAuth: false,
|
||||
},
|
||||
},
|
||||
},
|
||||
{
|
||||
desc: "use client auth example",
|
||||
in: `
|
||||
sources:
|
||||
my-spanner-admin-instance:
|
||||
kind: spanner-admin
|
||||
useClientOAuth: true
|
||||
`,
|
||||
want: map[string]sources.SourceConfig{
|
||||
"my-spanner-admin-instance": spanneradmin.Config{
|
||||
Name: "my-spanner-admin-instance",
|
||||
Kind: spanneradmin.SourceKind,
|
||||
UseClientOAuth: true,
|
||||
},
|
||||
},
|
||||
},
|
||||
}
|
||||
for _, tc := range tcs {
|
||||
tc := tc
|
||||
t.Run(tc.desc, func(t *testing.T) {
|
||||
t.Parallel()
|
||||
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) {
|
||||
t.Parallel()
|
||||
tcs := []struct {
|
||||
desc string
|
||||
in string
|
||||
err string
|
||||
}{
|
||||
{
|
||||
desc: "extra field",
|
||||
in: `
|
||||
sources:
|
||||
my-spanner-admin-instance:
|
||||
kind: spanner-admin
|
||||
project: test-project
|
||||
`,
|
||||
err: `unable to parse source "my-spanner-admin-instance" as "spanner-admin": [2:1] unknown field "project"
|
||||
1 | kind: spanner-admin
|
||||
> 2 | project: test-project
|
||||
^
|
||||
`,
|
||||
},
|
||||
{
|
||||
desc: "missing required field",
|
||||
in: `
|
||||
sources:
|
||||
my-spanner-admin-instance:
|
||||
useClientOAuth: true
|
||||
`,
|
||||
err: "missing 'kind' field for source \"my-spanner-admin-instance\"",
|
||||
},
|
||||
}
|
||||
for _, tc := range tcs {
|
||||
tc := tc
|
||||
t.Run(tc.desc, func(t *testing.T) {
|
||||
t.Parallel()
|
||||
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)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
@@ -1,233 +0,0 @@
|
||||
// Copyright 2026 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 spannercreateinstance
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
|
||||
instance "cloud.google.com/go/spanner/admin/instance/apiv1"
|
||||
"cloud.google.com/go/spanner/admin/instance/apiv1/instancepb"
|
||||
"github.com/goccy/go-yaml"
|
||||
"github.com/googleapis/genai-toolbox/internal/embeddingmodels"
|
||||
"github.com/googleapis/genai-toolbox/internal/sources"
|
||||
"github.com/googleapis/genai-toolbox/internal/tools"
|
||||
"github.com/googleapis/genai-toolbox/internal/util/parameters"
|
||||
)
|
||||
|
||||
const kind string = "spanner-create-instance"
|
||||
|
||||
func init() {
|
||||
if !tools.Register(kind, newConfig) {
|
||||
panic(fmt.Sprintf("tool kind %q already registered", 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 compatibleSource interface {
|
||||
GetDefaultProject() string
|
||||
GetClient(context.Context, string) (*instance.InstanceAdminClient, error)
|
||||
UseClientAuthorization() bool
|
||||
}
|
||||
|
||||
// Config defines the configuration for the create-instance tool.
|
||||
type Config struct {
|
||||
Name string `yaml:"name" validate:"required"`
|
||||
Kind string `yaml:"kind" validate:"required"`
|
||||
Description string `yaml:"description"`
|
||||
Source string `yaml:"source" validate:"required"`
|
||||
AuthRequired []string `yaml:"authRequired"`
|
||||
}
|
||||
|
||||
// validate interface
|
||||
var _ tools.ToolConfig = Config{}
|
||||
|
||||
// ToolConfigKind returns the kind of the tool.
|
||||
func (cfg Config) ToolConfigKind() string {
|
||||
return kind
|
||||
}
|
||||
|
||||
// Initialize initializes the tool from the configuration.
|
||||
func (cfg Config) Initialize(srcs map[string]sources.Source) (tools.Tool, error) {
|
||||
rawS, ok := srcs[cfg.Source]
|
||||
if !ok {
|
||||
return nil, fmt.Errorf("no source named %q configured", cfg.Source)
|
||||
}
|
||||
s, ok := rawS.(compatibleSource)
|
||||
if !ok {
|
||||
return nil, fmt.Errorf("invalid source for %q tool: source %q not compatible", kind, cfg.Source)
|
||||
}
|
||||
|
||||
project := s.GetDefaultProject()
|
||||
var projectParam parameters.Parameter
|
||||
if project != "" {
|
||||
projectParam = parameters.NewStringParameterWithDefault("project", project, "The GCP project ID.")
|
||||
} else {
|
||||
projectParam = parameters.NewStringParameter("project", "The project ID")
|
||||
}
|
||||
|
||||
allParameters := parameters.Parameters{
|
||||
projectParam,
|
||||
parameters.NewStringParameter("instanceId", "The ID of the instance"),
|
||||
parameters.NewStringParameter("displayName", "The display name of the instance"),
|
||||
parameters.NewStringParameter("config", "The instance configuration (e.g., regional-us-central1)"),
|
||||
parameters.NewIntParameter("nodeCount", "The number of nodes, mutually exclusive with processingUnits (one must be 0)"),
|
||||
parameters.NewIntParameter("processingUnits", "The number of processing units, mutually exclusive with nodeCount (one must be 0)"),
|
||||
parameters.NewStringParameter("edition", "The edition of the instance (STANDARD, ENTERPRISE, ENTERPRISE_PLUS)"),
|
||||
}
|
||||
|
||||
paramManifest := allParameters.Manifest()
|
||||
|
||||
description := cfg.Description
|
||||
if description == "" {
|
||||
description = "Creates a Spanner instance."
|
||||
}
|
||||
mcpManifest := tools.GetMcpManifest(cfg.Name, description, cfg.AuthRequired, allParameters, nil)
|
||||
|
||||
return Tool{
|
||||
Config: cfg,
|
||||
AllParams: allParameters,
|
||||
manifest: tools.Manifest{Description: description, Parameters: paramManifest, AuthRequired: cfg.AuthRequired},
|
||||
mcpManifest: mcpManifest,
|
||||
}, nil
|
||||
}
|
||||
|
||||
// Tool represents the create-instance tool.
|
||||
type Tool struct {
|
||||
Config
|
||||
AllParams parameters.Parameters
|
||||
manifest tools.Manifest
|
||||
mcpManifest tools.McpManifest
|
||||
}
|
||||
|
||||
func (t Tool) ToConfig() tools.ToolConfig {
|
||||
return t.Config
|
||||
}
|
||||
|
||||
// Invoke executes the tool's logic.
|
||||
func (t Tool) Invoke(ctx context.Context, resourceMgr tools.SourceProvider, params parameters.ParamValues, accessToken tools.AccessToken) (any, error) {
|
||||
paramsMap := params.AsMap()
|
||||
|
||||
project, _ := paramsMap["project"].(string)
|
||||
instanceId, _ := paramsMap["instanceId"].(string)
|
||||
displayName, _ := paramsMap["displayName"].(string)
|
||||
config, _ := paramsMap["config"].(string)
|
||||
nodeCount, _ := paramsMap["nodeCount"].(int)
|
||||
processingUnits, _ := paramsMap["processingUnits"].(int)
|
||||
editionStr, _ := paramsMap["edition"].(string)
|
||||
|
||||
if (nodeCount > 0 && processingUnits > 0) || (nodeCount == 0 && processingUnits == 0) {
|
||||
return nil, fmt.Errorf("one of nodeCount or processingUnits must be positive, and the other must be 0")
|
||||
}
|
||||
|
||||
source, err := tools.GetCompatibleSource[compatibleSource](resourceMgr, t.Source, t.Name, t.Kind)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
client, err := source.GetClient(ctx, string(accessToken))
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if source.UseClientAuthorization() {
|
||||
defer client.Close()
|
||||
}
|
||||
|
||||
parent := fmt.Sprintf("projects/%s", project)
|
||||
instanceConfig := fmt.Sprintf("projects/%s/instanceConfigs/%s", project, config)
|
||||
|
||||
var edition instancepb.Instance_Edition
|
||||
switch editionStr {
|
||||
case "STANDARD":
|
||||
edition = instancepb.Instance_STANDARD
|
||||
case "ENTERPRISE":
|
||||
edition = instancepb.Instance_ENTERPRISE
|
||||
case "ENTERPRISE_PLUS":
|
||||
edition = instancepb.Instance_ENTERPRISE_PLUS
|
||||
default:
|
||||
edition = instancepb.Instance_EDITION_UNSPECIFIED
|
||||
}
|
||||
|
||||
// Construct the instance object
|
||||
instance := &instancepb.Instance{
|
||||
Config: instanceConfig,
|
||||
DisplayName: displayName,
|
||||
Edition: edition,
|
||||
NodeCount: int32(nodeCount),
|
||||
ProcessingUnits: int32(processingUnits),
|
||||
}
|
||||
|
||||
req := &instancepb.CreateInstanceRequest{
|
||||
Parent: parent,
|
||||
InstanceId: instanceId,
|
||||
Instance: instance,
|
||||
}
|
||||
|
||||
op, err := client.CreateInstance(ctx, req)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("failed to create instance: %w", err)
|
||||
}
|
||||
|
||||
// Wait for the operation to complete
|
||||
resp, err := op.Wait(ctx)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("failed to wait for create instance operation: %w", err)
|
||||
}
|
||||
|
||||
return resp, nil
|
||||
}
|
||||
|
||||
// ParseParams parses the parameters for the tool.
|
||||
func (t Tool) ParseParams(data map[string]any, claims map[string]map[string]any) (parameters.ParamValues, error) {
|
||||
return parameters.ParseParams(t.AllParams, data, claims)
|
||||
}
|
||||
|
||||
func (t Tool) EmbedParams(ctx context.Context, paramValues parameters.ParamValues, embeddingModelsMap map[string]embeddingmodels.EmbeddingModel) (parameters.ParamValues, error) {
|
||||
return parameters.EmbedParams(ctx, t.AllParams, paramValues, embeddingModelsMap, nil)
|
||||
}
|
||||
|
||||
// Manifest returns the tool's manifest.
|
||||
func (t Tool) Manifest() tools.Manifest {
|
||||
return t.manifest
|
||||
}
|
||||
|
||||
// McpManifest returns the tool's MCP manifest.
|
||||
func (t Tool) McpManifest() tools.McpManifest {
|
||||
return t.mcpManifest
|
||||
}
|
||||
|
||||
// Authorized checks if the tool is authorized.
|
||||
func (t Tool) Authorized(verifiedAuthServices []string) bool {
|
||||
return true
|
||||
}
|
||||
|
||||
func (t Tool) RequiresClientAuthorization(resourceMgr tools.SourceProvider) (bool, error) {
|
||||
source, err := tools.GetCompatibleSource[compatibleSource](resourceMgr, t.Source, t.Name, t.Kind)
|
||||
if err != nil {
|
||||
return false, err
|
||||
}
|
||||
return source.UseClientAuthorization(), nil
|
||||
}
|
||||
|
||||
func (t Tool) GetAuthTokenHeaderName(resourceMgr tools.SourceProvider) (string, error) {
|
||||
return "Authorization", nil
|
||||
}
|
||||
@@ -1,110 +0,0 @@
|
||||
// Copyright 2026 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 spannercreateinstance_test
|
||||
|
||||
import (
|
||||
"context"
|
||||
"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/spanneradmin/spannercreateinstance"
|
||||
"github.com/googleapis/genai-toolbox/internal/util/parameters"
|
||||
)
|
||||
|
||||
func TestParseFromYaml(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:
|
||||
create-instance-tool:
|
||||
kind: spanner-create-instance
|
||||
description: a test description
|
||||
source: a-source
|
||||
`,
|
||||
want: server.ToolConfigs{
|
||||
"create-instance-tool": spannercreateinstance.Config{
|
||||
Name: "create-instance-tool",
|
||||
Kind: "spanner-create-instance",
|
||||
Description: "a test description",
|
||||
Source: "a-source",
|
||||
AuthRequired: []string{},
|
||||
},
|
||||
},
|
||||
},
|
||||
}
|
||||
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 TestInvokeNodeCountAndProcessingUnitsValidation(t *testing.T) {
|
||||
t.Parallel()
|
||||
testCases := []struct {
|
||||
name string
|
||||
params parameters.ParamValues
|
||||
wantErr string
|
||||
}{
|
||||
{
|
||||
name: "Both positive",
|
||||
params: parameters.ParamValues{
|
||||
{Name: "nodeCount", Value: 1},
|
||||
{Name: "processingUnits", Value: 1000},
|
||||
},
|
||||
wantErr: "one of nodeCount or processingUnits must be positive, and the other must be 0",
|
||||
},
|
||||
{
|
||||
name: "Both zero",
|
||||
params: parameters.ParamValues{
|
||||
{Name: "nodeCount", Value: 0},
|
||||
{Name: "processingUnits", Value: 0},
|
||||
},
|
||||
wantErr: "one of nodeCount or processingUnits must be positive, and the other must be 0",
|
||||
},
|
||||
}
|
||||
|
||||
for _, tc := range testCases {
|
||||
t.Run(tc.name, func(t *testing.T) {
|
||||
tool := spannercreateinstance.Tool{}
|
||||
_, err := tool.Invoke(context.Background(), nil, tc.params, "")
|
||||
if err == nil || err.Error() != tc.wantErr {
|
||||
t.Errorf("Invoke() error = %v, wantErr %v", err, tc.wantErr)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
@@ -1,178 +0,0 @@
|
||||
// Copyright 2026 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 spanneradmin
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"io"
|
||||
"net/http"
|
||||
"os"
|
||||
"regexp"
|
||||
"strings"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
instance "cloud.google.com/go/spanner/admin/instance/apiv1"
|
||||
"cloud.google.com/go/spanner/admin/instance/apiv1/instancepb"
|
||||
"github.com/google/uuid"
|
||||
"github.com/googleapis/genai-toolbox/internal/testutils"
|
||||
"github.com/googleapis/genai-toolbox/tests"
|
||||
)
|
||||
|
||||
var (
|
||||
SpannerProject = os.Getenv("SPANNER_PROJECT")
|
||||
)
|
||||
|
||||
func getSpannerAdminVars(t *testing.T) map[string]any {
|
||||
if SpannerProject == "" {
|
||||
t.Fatal("'SPANNER_PROJECT' not set")
|
||||
}
|
||||
|
||||
return map[string]any{
|
||||
"kind": "spanner-admin",
|
||||
"defaultProject": SpannerProject,
|
||||
}
|
||||
}
|
||||
|
||||
func TestSpannerAdminCreateInstance(t *testing.T) {
|
||||
sourceConfig := getSpannerAdminVars(t)
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 20*time.Minute)
|
||||
defer cancel()
|
||||
|
||||
shortUuid := strings.ReplaceAll(uuid.New().String(), "-", "")[:10]
|
||||
instanceId := "test-inst-" + shortUuid
|
||||
|
||||
displayName := "Test Instance " + shortUuid
|
||||
instanceConfig := "regional-us-central1"
|
||||
nodeCount := 1
|
||||
edition := "ENTERPRISE"
|
||||
|
||||
// Setup Admin Client for verification and cleanup
|
||||
adminClient, err := instance.NewInstanceAdminClient(ctx)
|
||||
if err != nil {
|
||||
t.Fatalf("unable to create Spanner instance admin client: %s", err)
|
||||
}
|
||||
defer adminClient.Close()
|
||||
|
||||
// Teardown function
|
||||
defer func() {
|
||||
err := adminClient.DeleteInstance(ctx, &instancepb.DeleteInstanceRequest{
|
||||
Name: fmt.Sprintf("projects/%s/instances/%s", SpannerProject, instanceId),
|
||||
})
|
||||
if err != nil {
|
||||
// If it fails, it might not have been created, log it but don't fail if it's "not found"
|
||||
t.Logf("cleanup: failed to delete instance %s: %s", instanceId, err)
|
||||
} else {
|
||||
t.Logf("cleanup: deleted instance %s", instanceId)
|
||||
}
|
||||
}()
|
||||
|
||||
// Construct Tools Config
|
||||
|
||||
toolsConfig := map[string]any{
|
||||
"sources": map[string]any{
|
||||
"my-spanner-admin": sourceConfig,
|
||||
},
|
||||
"tools": map[string]any{
|
||||
"create-instance-tool": map[string]any{
|
||||
"kind": "spanner-create-instance",
|
||||
"source": "my-spanner-admin",
|
||||
"description": "Creates a Spanner instance.",
|
||||
},
|
||||
},
|
||||
}
|
||||
|
||||
// Start Toolbox Server
|
||||
cmd, cleanup, err := tests.StartCmd(ctx, toolsConfig)
|
||||
if err != nil {
|
||||
t.Fatalf("command initialization returned an error: %s", err)
|
||||
}
|
||||
defer cleanup()
|
||||
|
||||
waitCtx, cancelWait := context.WithTimeout(ctx, 10*time.Second)
|
||||
defer cancelWait()
|
||||
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)
|
||||
}
|
||||
|
||||
// Prepare Invocation Payload
|
||||
|
||||
payload := map[string]any{
|
||||
"project": SpannerProject,
|
||||
"instanceId": instanceId,
|
||||
"displayName": displayName,
|
||||
"config": instanceConfig,
|
||||
"nodeCount": nodeCount,
|
||||
"edition": edition,
|
||||
"processingUnits": 0,
|
||||
}
|
||||
|
||||
payloadBytes, err := json.Marshal(payload)
|
||||
if err != nil {
|
||||
t.Fatalf("failed to marshal payload: %s", err)
|
||||
}
|
||||
|
||||
// Invoke Tool
|
||||
invokeUrl := "http://127.0.0.1:5000/api/tool/create-instance-tool/invoke"
|
||||
req, err := http.NewRequest(http.MethodPost, invokeUrl, bytes.NewBuffer(payloadBytes))
|
||||
if err != nil {
|
||||
t.Fatalf("unable to create request: %s", err)
|
||||
}
|
||||
req.Header.Add("Content-type", "application/json")
|
||||
|
||||
t.Logf("Invoking create-instance-tool for instance: %s", instanceId)
|
||||
resp, err := http.DefaultClient.Do(req)
|
||||
if err != nil {
|
||||
t.Fatalf("unable to send request: %s", err)
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
|
||||
if resp.StatusCode != http.StatusOK {
|
||||
bodyBytes, _ := io.ReadAll(resp.Body)
|
||||
t.Fatalf("response status code is not 200, got %d: %s", resp.StatusCode, string(bodyBytes))
|
||||
}
|
||||
|
||||
// Check Response
|
||||
var body map[string]interface{}
|
||||
err = json.NewDecoder(resp.Body).Decode(&body)
|
||||
if err != nil {
|
||||
t.Fatalf("error parsing response body")
|
||||
}
|
||||
|
||||
// Verify Instance Exists via Admin Client
|
||||
t.Logf("Verifying instance %s exists...", instanceId)
|
||||
instanceName := fmt.Sprintf("projects/%s/instances/%s", SpannerProject, instanceId)
|
||||
gotInstance, err := adminClient.GetInstance(ctx, &instancepb.GetInstanceRequest{
|
||||
Name: instanceName,
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("failed to get instance from admin client: %s", err)
|
||||
}
|
||||
|
||||
if gotInstance.Name != instanceName {
|
||||
t.Errorf("expected instance name %s, got %s", instanceName, gotInstance.Name)
|
||||
}
|
||||
if gotInstance.DisplayName != displayName {
|
||||
t.Errorf("expected display name %s, got %s", displayName, gotInstance.DisplayName)
|
||||
}
|
||||
if gotInstance.NodeCount != int32(nodeCount) {
|
||||
t.Errorf("expected node count %d, got %d", nodeCount, gotInstance.NodeCount)
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user