From 302b564ca1f9208b0f7c7212e58255b2d3bc1c66 Mon Sep 17 00:00:00 2001 From: gRedHeadphone Date: Fri, 19 Dec 2025 09:59:05 +0000 Subject: [PATCH 01/17] feat(spanner-admin): spanner admin source + create instance tool --- .ci/integration.cloudbuild.yaml | 20 ++ cmd/root.go | 2 + docs/SPANNERADMIN_README.md | 59 +++++ docs/en/resources/sources/spanner-admin.md | 42 ++++ .../spanneradmin/spanner-create-instance.md | 52 +++++ .../prebuiltconfigs/prebuiltconfigs_test.go | 1 + .../prebuiltconfigs/tools/spanner-admin.yaml | 27 +++ internal/sources/spanneradmin/spanneradmin.go | 116 ++++++++++ .../sources/spanneradmin/spanneradmin_test.go | 135 +++++++++++ .../spannercreateinstance.go | 215 ++++++++++++++++++ .../spannercreateinstance_test.go | 110 +++++++++ .../spanneradmin_integration_test.go | 178 +++++++++++++++ 12 files changed, 957 insertions(+) create mode 100644 docs/SPANNERADMIN_README.md create mode 100644 docs/en/resources/sources/spanner-admin.md create mode 100644 docs/en/resources/tools/spanneradmin/spanner-create-instance.md create mode 100644 internal/prebuiltconfigs/tools/spanner-admin.yaml create mode 100644 internal/sources/spanneradmin/spanneradmin.go create mode 100644 internal/sources/spanneradmin/spanneradmin_test.go create mode 100644 internal/tools/spanneradmin/spannercreateinstance/spannercreateinstance.go create mode 100644 internal/tools/spanneradmin/spannercreateinstance/spannercreateinstance_test.go create mode 100644 tests/spanneradmin/spanneradmin_integration_test.go diff --git a/.ci/integration.cloudbuild.yaml b/.ci/integration.cloudbuild.yaml index 14742514bcc..082cef97753 100644 --- a/.ci/integration.cloudbuild.yaml +++ b/.ci/integration.cloudbuild.yaml @@ -340,6 +340,26 @@ 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"] diff --git a/cmd/root.go b/cmd/root.go index 4a34cf457d9..b50ccde03a3 100644 --- a/cmd/root.go +++ b/cmd/root.go @@ -217,6 +217,7 @@ 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" @@ -261,6 +262,7 @@ import ( _ "github.com/googleapis/genai-toolbox/internal/sources/serverlessspark" _ "github.com/googleapis/genai-toolbox/internal/sources/singlestore" _ "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" diff --git a/docs/SPANNERADMIN_README.md b/docs/SPANNERADMIN_README.md new file mode 100644 index 00000000000..7995880700b --- /dev/null +++ b/docs/SPANNERADMIN_README.md @@ -0,0 +1,59 @@ +# 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). diff --git a/docs/en/resources/sources/spanner-admin.md b/docs/en/resources/sources/spanner-admin.md new file mode 100644 index 00000000000..8926315d73d --- /dev/null +++ b/docs/en/resources/sources/spanner-admin.md @@ -0,0 +1,42 @@ +--- +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`. | diff --git a/docs/en/resources/tools/spanneradmin/spanner-create-instance.md b/docs/en/resources/tools/spanneradmin/spanner-create-instance.md new file mode 100644 index 00000000000..3e683d7356b --- /dev/null +++ b/docs/en/resources/tools/spanneradmin/spanner-create-instance.md @@ -0,0 +1,52 @@ +--- +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. | diff --git a/internal/prebuiltconfigs/prebuiltconfigs_test.go b/internal/prebuiltconfigs/prebuiltconfigs_test.go index d2247030a46..79215cd84b7 100644 --- a/internal/prebuiltconfigs/prebuiltconfigs_test.go +++ b/internal/prebuiltconfigs/prebuiltconfigs_test.go @@ -49,6 +49,7 @@ var expectedToolSources = []string{ "postgres", "serverless-spark", "singlestore", + "spanner-admin", "spanner-postgres", "spanner", "sqlite", diff --git a/internal/prebuiltconfigs/tools/spanner-admin.yaml b/internal/prebuiltconfigs/tools/spanner-admin.yaml new file mode 100644 index 00000000000..c68336124a8 --- /dev/null +++ b/internal/prebuiltconfigs/tools/spanner-admin.yaml @@ -0,0 +1,27 @@ +# Copyright 2025 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +sources: + 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 diff --git a/internal/sources/spanneradmin/spanneradmin.go b/internal/sources/spanneradmin/spanneradmin.go new file mode 100644 index 00000000000..815edd63e12 --- /dev/null +++ b/internal/sources/spanneradmin/spanneradmin.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 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) 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 +} diff --git a/internal/sources/spanneradmin/spanneradmin_test.go b/internal/sources/spanneradmin/spanneradmin_test.go new file mode 100644 index 00000000000..d42c43ed999 --- /dev/null +++ b/internal/sources/spanneradmin/spanneradmin_test.go @@ -0,0 +1,135 @@ +// 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 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) + } + }) + } +} diff --git a/internal/tools/spanneradmin/spannercreateinstance/spannercreateinstance.go b/internal/tools/spanneradmin/spannercreateinstance/spannercreateinstance.go new file mode 100644 index 00000000000..d48a9148bca --- /dev/null +++ b/internal/tools/spanneradmin/spannercreateinstance/spannercreateinstance.go @@ -0,0 +1,215 @@ +// 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 spannercreateinstance + +import ( + "context" + "fmt" + + "cloud.google.com/go/spanner/admin/instance/apiv1/instancepb" + "github.com/goccy/go-yaml" + "github.com/googleapis/genai-toolbox/internal/sources" + "github.com/googleapis/genai-toolbox/internal/sources/spanneradmin" + "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 +} + +// 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.(*spanneradmin.Source) + if !ok { + return nil, fmt.Errorf("invalid source for %q tool: source kind must be `spanner-admin`", kind) + } + + project := s.DefaultProject + 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"), + parameters.NewIntParameter("processingUnits", "The number of processing units"), + 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, + Source: s, + 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 + Source *spanneradmin.Source + 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") + } + + client, err := t.Source.GetClient(ctx, string(accessToken)) + if err != nil { + return nil, err + } + if t.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) +} + +// 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 { + return t.Source.UseClientAuthorization() +} + +func (t Tool) GetAuthTokenHeaderName() string { + return "Authorization" +} diff --git a/internal/tools/spanneradmin/spannercreateinstance/spannercreateinstance_test.go b/internal/tools/spanneradmin/spannercreateinstance/spannercreateinstance_test.go new file mode 100644 index 00000000000..9465dc0ba30 --- /dev/null +++ b/internal/tools/spanneradmin/spannercreateinstance/spannercreateinstance_test.go @@ -0,0 +1,110 @@ +// 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 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 TestInvoke_NodeCountAndProcessingUnitsValidation(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) + } + }) + } +} diff --git a/tests/spanneradmin/spanneradmin_integration_test.go b/tests/spanneradmin/spanneradmin_integration_test.go new file mode 100644 index 00000000000..f886e485768 --- /dev/null +++ b/tests/spanneradmin/spanneradmin_integration_test.go @@ -0,0 +1,178 @@ +// 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 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) + } +} From 5b7f5a30392b4cf33d626b4e25e02b457696bb27 Mon Sep 17 00:00:00 2001 From: gRedHeadphone Date: Fri, 19 Dec 2025 10:10:59 +0000 Subject: [PATCH 02/17] chore: update parameter description & format docs table spanner create instance --- .../spanneradmin/spanner-create-instance.md | 18 +++++++++--------- .../spannercreateinstance.go | 4 ++-- 2 files changed, 11 insertions(+), 11 deletions(-) diff --git a/docs/en/resources/tools/spanneradmin/spanner-create-instance.md b/docs/en/resources/tools/spanneradmin/spanner-create-instance.md index 3e683d7356b..80ed99f785f 100644 --- a/docs/en/resources/tools/spanneradmin/spanner-create-instance.md +++ b/docs/en/resources/tools/spanneradmin/spanner-create-instance.md @@ -33,15 +33,15 @@ tools: 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`). | +| **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 diff --git a/internal/tools/spanneradmin/spannercreateinstance/spannercreateinstance.go b/internal/tools/spanneradmin/spannercreateinstance/spannercreateinstance.go index d48a9148bca..5ea31fcd7a8 100644 --- a/internal/tools/spanneradmin/spannercreateinstance/spannercreateinstance.go +++ b/internal/tools/spanneradmin/spannercreateinstance/spannercreateinstance.go @@ -83,8 +83,8 @@ func (cfg Config) Initialize(srcs map[string]sources.Source) (tools.Tool, error) 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"), - parameters.NewIntParameter("processingUnits", "The number of processing units"), + 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)"), } From e6de2170cf9e74da377b1a442b4a1fb2fa277612 Mon Sep 17 00:00:00 2001 From: gRedHeadphone Date: Mon, 22 Dec 2025 10:52:11 +0530 Subject: [PATCH 03/17] Update internal/sources/spanneradmin/spanneradmin.go Co-authored-by: Yuan Teoh <45984206+Yuan325@users.noreply.github.com> --- internal/sources/spanneradmin/spanneradmin.go | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/internal/sources/spanneradmin/spanneradmin.go b/internal/sources/spanneradmin/spanneradmin.go index 815edd63e12..9ffede63a47 100644 --- a/internal/sources/spanneradmin/spanneradmin.go +++ b/internal/sources/spanneradmin/spanneradmin.go @@ -95,6 +95,10 @@ 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} From 4f62db499d99800cceab5d361723f469c6bbcae5 Mon Sep 17 00:00:00 2001 From: gRedHeadphone Date: Mon, 22 Dec 2025 10:52:38 +0530 Subject: [PATCH 04/17] Update internal/tools/spanneradmin/spannercreateinstance/spannercreateinstance.go Co-authored-by: Yuan Teoh <45984206+Yuan325@users.noreply.github.com> --- .../spannercreateinstance/spannercreateinstance.go | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/internal/tools/spanneradmin/spannercreateinstance/spannercreateinstance.go b/internal/tools/spanneradmin/spannercreateinstance/spannercreateinstance.go index 5ea31fcd7a8..a16061ff44c 100644 --- a/internal/tools/spanneradmin/spannercreateinstance/spannercreateinstance.go +++ b/internal/tools/spanneradmin/spannercreateinstance/spannercreateinstance.go @@ -42,6 +42,12 @@ func newConfig(ctx context.Context, name string, decoder *yaml.Decoder) (tools.T 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"` From e831f71421286da07b217d02a4e94a0611ae7902 Mon Sep 17 00:00:00 2001 From: gRedHeadphone Date: Mon, 22 Dec 2025 10:53:12 +0530 Subject: [PATCH 05/17] Update internal/tools/spanneradmin/spannercreateinstance/spannercreateinstance.go Co-authored-by: Yuan Teoh <45984206+Yuan325@users.noreply.github.com> --- .../spanneradmin/spannercreateinstance/spannercreateinstance.go | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/internal/tools/spanneradmin/spannercreateinstance/spannercreateinstance.go b/internal/tools/spanneradmin/spannercreateinstance/spannercreateinstance.go index a16061ff44c..03f741f31a8 100644 --- a/internal/tools/spanneradmin/spannercreateinstance/spannercreateinstance.go +++ b/internal/tools/spanneradmin/spannercreateinstance/spannercreateinstance.go @@ -73,7 +73,7 @@ func (cfg Config) Initialize(srcs map[string]sources.Source) (tools.Tool, error) } s, ok := rawS.(*spanneradmin.Source) if !ok { - return nil, fmt.Errorf("invalid source for %q tool: source kind must be `spanner-admin`", kind) + return nil, fmt.Errorf("invalid source for %q tool: source %q not compatible", kind, cfg.Source) } project := s.DefaultProject From 330dd843bf49fdf74c26853ad77d2d69fa085892 Mon Sep 17 00:00:00 2001 From: gRedHeadphone Date: Mon, 22 Dec 2025 10:53:30 +0530 Subject: [PATCH 06/17] Update internal/tools/spanneradmin/spannercreateinstance/spannercreateinstance.go Co-authored-by: Yuan Teoh <45984206+Yuan325@users.noreply.github.com> --- .../spanneradmin/spannercreateinstance/spannercreateinstance.go | 1 - 1 file changed, 1 deletion(-) diff --git a/internal/tools/spanneradmin/spannercreateinstance/spannercreateinstance.go b/internal/tools/spanneradmin/spannercreateinstance/spannercreateinstance.go index 03f741f31a8..f75d19ff72b 100644 --- a/internal/tools/spanneradmin/spannercreateinstance/spannercreateinstance.go +++ b/internal/tools/spanneradmin/spannercreateinstance/spannercreateinstance.go @@ -104,7 +104,6 @@ func (cfg Config) Initialize(srcs map[string]sources.Source) (tools.Tool, error) return Tool{ Config: cfg, - Source: s, AllParams: allParameters, manifest: tools.Manifest{Description: description, Parameters: paramManifest, AuthRequired: cfg.AuthRequired}, mcpManifest: mcpManifest, From 7d79f4909a360827a2eaffffe9a6c1a435e3f21b Mon Sep 17 00:00:00 2001 From: gRedHeadphone Date: Mon, 22 Dec 2025 10:53:43 +0530 Subject: [PATCH 07/17] Update internal/tools/spanneradmin/spannercreateinstance/spannercreateinstance.go Co-authored-by: Yuan Teoh <45984206+Yuan325@users.noreply.github.com> --- .../spanneradmin/spannercreateinstance/spannercreateinstance.go | 1 - 1 file changed, 1 deletion(-) diff --git a/internal/tools/spanneradmin/spannercreateinstance/spannercreateinstance.go b/internal/tools/spanneradmin/spannercreateinstance/spannercreateinstance.go index f75d19ff72b..50d6a8d28ff 100644 --- a/internal/tools/spanneradmin/spannercreateinstance/spannercreateinstance.go +++ b/internal/tools/spanneradmin/spannercreateinstance/spannercreateinstance.go @@ -113,7 +113,6 @@ func (cfg Config) Initialize(srcs map[string]sources.Source) (tools.Tool, error) // Tool represents the create-instance tool. type Tool struct { Config - Source *spanneradmin.Source AllParams parameters.Parameters manifest tools.Manifest mcpManifest tools.McpManifest From 09c979d8db6dda88a04120598a262ab8e1b62205 Mon Sep 17 00:00:00 2001 From: gRedHeadphone Date: Mon, 22 Dec 2025 10:54:26 +0530 Subject: [PATCH 08/17] Update internal/tools/spanneradmin/spannercreateinstance/spannercreateinstance.go Co-authored-by: Yuan Teoh <45984206+Yuan325@users.noreply.github.com> --- .../spannercreateinstance/spannercreateinstance.go | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/internal/tools/spanneradmin/spannercreateinstance/spannercreateinstance.go b/internal/tools/spanneradmin/spannercreateinstance/spannercreateinstance.go index 50d6a8d28ff..4ba799041b1 100644 --- a/internal/tools/spanneradmin/spannercreateinstance/spannercreateinstance.go +++ b/internal/tools/spanneradmin/spannercreateinstance/spannercreateinstance.go @@ -124,6 +124,11 @@ func (t Tool) ToConfig() tools.ToolConfig { // Invoke executes the tool's logic. func (t Tool) Invoke(ctx context.Context, resourceMgr tools.SourceProvider, params parameters.ParamValues, accessToken tools.AccessToken) (any, error) { + source, err := tools.GetCompatibleSource[compatibleSource](resourceMgr, t.Source, t.Name, t.Kind) + if err != nil { + return nil, err + } + paramsMap := params.AsMap() project, _ := paramsMap["project"].(string) From a9edd5e32da3a59ede23ffdf71688c7d3833f5ea Mon Sep 17 00:00:00 2001 From: gRedHeadphone Date: Mon, 22 Dec 2025 10:54:40 +0530 Subject: [PATCH 09/17] Update internal/tools/spanneradmin/spannercreateinstance/spannercreateinstance.go Co-authored-by: Yuan Teoh <45984206+Yuan325@users.noreply.github.com> --- .../spanneradmin/spannercreateinstance/spannercreateinstance.go | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/internal/tools/spanneradmin/spannercreateinstance/spannercreateinstance.go b/internal/tools/spanneradmin/spannercreateinstance/spannercreateinstance.go index 4ba799041b1..2200aa852b7 100644 --- a/internal/tools/spanneradmin/spannercreateinstance/spannercreateinstance.go +++ b/internal/tools/spanneradmin/spannercreateinstance/spannercreateinstance.go @@ -143,7 +143,7 @@ func (t Tool) Invoke(ctx context.Context, resourceMgr tools.SourceProvider, para return nil, fmt.Errorf("one of nodeCount or processingUnits must be positive, and the other must be 0") } - client, err := t.Source.GetClient(ctx, string(accessToken)) + client, err := source.GetClient(ctx, string(accessToken)) if err != nil { return nil, err } From 4170fe309ff4b55f746ed7a7317f6b5b421cd19c Mon Sep 17 00:00:00 2001 From: gRedHeadphone Date: Mon, 22 Dec 2025 10:54:58 +0530 Subject: [PATCH 10/17] Update internal/tools/spanneradmin/spannercreateinstance/spannercreateinstance.go Co-authored-by: Yuan Teoh <45984206+Yuan325@users.noreply.github.com> --- .../spanneradmin/spannercreateinstance/spannercreateinstance.go | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/internal/tools/spanneradmin/spannercreateinstance/spannercreateinstance.go b/internal/tools/spanneradmin/spannercreateinstance/spannercreateinstance.go index 2200aa852b7..c2a2b8cba2c 100644 --- a/internal/tools/spanneradmin/spannercreateinstance/spannercreateinstance.go +++ b/internal/tools/spanneradmin/spannercreateinstance/spannercreateinstance.go @@ -147,7 +147,7 @@ func (t Tool) Invoke(ctx context.Context, resourceMgr tools.SourceProvider, para if err != nil { return nil, err } - if t.Source.UseClientAuthorization() { + if source.UseClientAuthorization() { defer client.Close() } From f9f34b1005ad1a1d3bdf4216f7d73bf623545ed6 Mon Sep 17 00:00:00 2001 From: gRedHeadphone Date: Mon, 22 Dec 2025 10:55:26 +0530 Subject: [PATCH 11/17] Update internal/tools/spanneradmin/spannercreateinstance/spannercreateinstance.go Co-authored-by: Yuan Teoh <45984206+Yuan325@users.noreply.github.com> --- .../spannercreateinstance/spannercreateinstance.go | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/internal/tools/spanneradmin/spannercreateinstance/spannercreateinstance.go b/internal/tools/spanneradmin/spannercreateinstance/spannercreateinstance.go index c2a2b8cba2c..5fccb7dce4d 100644 --- a/internal/tools/spanneradmin/spannercreateinstance/spannercreateinstance.go +++ b/internal/tools/spanneradmin/spannercreateinstance/spannercreateinstance.go @@ -215,8 +215,12 @@ func (t Tool) Authorized(verifiedAuthServices []string) bool { return true } -func (t Tool) RequiresClientAuthorization(resourceMgr tools.SourceProvider) bool { - return t.Source.UseClientAuthorization() +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() string { From 62095feadbe5ec2ba6b44044f3cec3893206ad45 Mon Sep 17 00:00:00 2001 From: gRedHeadphone Date: Mon, 22 Dec 2025 10:55:45 +0530 Subject: [PATCH 12/17] Update internal/tools/spanneradmin/spannercreateinstance/spannercreateinstance.go Co-authored-by: Yuan Teoh <45984206+Yuan325@users.noreply.github.com> --- .../spannercreateinstance/spannercreateinstance.go | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/internal/tools/spanneradmin/spannercreateinstance/spannercreateinstance.go b/internal/tools/spanneradmin/spannercreateinstance/spannercreateinstance.go index 5fccb7dce4d..9f5ab51e349 100644 --- a/internal/tools/spanneradmin/spannercreateinstance/spannercreateinstance.go +++ b/internal/tools/spanneradmin/spannercreateinstance/spannercreateinstance.go @@ -223,6 +223,6 @@ func (t Tool) RequiresClientAuthorization(resourceMgr tools.SourceProvider) (boo return source.UseClientAuthorization(), nil } -func (t Tool) GetAuthTokenHeaderName() string { - return "Authorization" +func (t Tool) GetAuthTokenHeaderName(resourceMgr tools.SourceProvider) (string, error) { + return "Authorization", nil } From 511c4651f3ca4cc9a177c170fa9449b50248ec4f Mon Sep 17 00:00:00 2001 From: gRedHeadphone Date: Mon, 22 Dec 2025 10:56:08 +0530 Subject: [PATCH 13/17] Update internal/tools/spanneradmin/spannercreateinstance/spannercreateinstance.go Co-authored-by: Yuan Teoh <45984206+Yuan325@users.noreply.github.com> --- .../spanneradmin/spannercreateinstance/spannercreateinstance.go | 1 - 1 file changed, 1 deletion(-) diff --git a/internal/tools/spanneradmin/spannercreateinstance/spannercreateinstance.go b/internal/tools/spanneradmin/spannercreateinstance/spannercreateinstance.go index 9f5ab51e349..21be8647826 100644 --- a/internal/tools/spanneradmin/spannercreateinstance/spannercreateinstance.go +++ b/internal/tools/spanneradmin/spannercreateinstance/spannercreateinstance.go @@ -21,7 +21,6 @@ import ( "cloud.google.com/go/spanner/admin/instance/apiv1/instancepb" "github.com/goccy/go-yaml" "github.com/googleapis/genai-toolbox/internal/sources" - "github.com/googleapis/genai-toolbox/internal/sources/spanneradmin" "github.com/googleapis/genai-toolbox/internal/tools" "github.com/googleapis/genai-toolbox/internal/util/parameters" ) From c67b0cf0fca1b5292eff85f359313ff4a7f7b4dd Mon Sep 17 00:00:00 2001 From: gRedHeadphone Date: Mon, 22 Dec 2025 10:56:20 +0530 Subject: [PATCH 14/17] Update internal/tools/spanneradmin/spannercreateinstance/spannercreateinstance.go Co-authored-by: Yuan Teoh <45984206+Yuan325@users.noreply.github.com> --- .../spanneradmin/spannercreateinstance/spannercreateinstance.go | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/internal/tools/spanneradmin/spannercreateinstance/spannercreateinstance.go b/internal/tools/spanneradmin/spannercreateinstance/spannercreateinstance.go index 21be8647826..a765729ace7 100644 --- a/internal/tools/spanneradmin/spannercreateinstance/spannercreateinstance.go +++ b/internal/tools/spanneradmin/spannercreateinstance/spannercreateinstance.go @@ -70,7 +70,7 @@ func (cfg Config) Initialize(srcs map[string]sources.Source) (tools.Tool, error) if !ok { return nil, fmt.Errorf("no source named %q configured", cfg.Source) } - s, ok := rawS.(*spanneradmin.Source) + s, ok := rawS.(compatibleSource) if !ok { return nil, fmt.Errorf("invalid source for %q tool: source %q not compatible", kind, cfg.Source) } From af62ddf9c183139e04ac446bfe0852509adf5a79 Mon Sep 17 00:00:00 2001 From: gRedHeadphone Date: Mon, 22 Dec 2025 05:49:49 +0000 Subject: [PATCH 15/17] chore: minor fixes --- .../spannercreateinstance/spannercreateinstance.go | 13 +++++++------ .../spannercreateinstance_test.go | 2 +- 2 files changed, 8 insertions(+), 7 deletions(-) diff --git a/internal/tools/spanneradmin/spannercreateinstance/spannercreateinstance.go b/internal/tools/spanneradmin/spannercreateinstance/spannercreateinstance.go index a765729ace7..fcc5ae9bc09 100644 --- a/internal/tools/spanneradmin/spannercreateinstance/spannercreateinstance.go +++ b/internal/tools/spanneradmin/spannercreateinstance/spannercreateinstance.go @@ -18,6 +18,7 @@ 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/sources" @@ -75,7 +76,7 @@ func (cfg Config) Initialize(srcs map[string]sources.Source) (tools.Tool, error) return nil, fmt.Errorf("invalid source for %q tool: source %q not compatible", kind, cfg.Source) } - project := s.DefaultProject + project := s.GetDefaultProject() var projectParam parameters.Parameter if project != "" { projectParam = parameters.NewStringParameterWithDefault("project", project, "The GCP project ID.") @@ -123,11 +124,6 @@ func (t Tool) ToConfig() tools.ToolConfig { // Invoke executes the tool's logic. func (t Tool) Invoke(ctx context.Context, resourceMgr tools.SourceProvider, params parameters.ParamValues, accessToken tools.AccessToken) (any, error) { - source, err := tools.GetCompatibleSource[compatibleSource](resourceMgr, t.Source, t.Name, t.Kind) - if err != nil { - return nil, err - } - paramsMap := params.AsMap() project, _ := paramsMap["project"].(string) @@ -142,6 +138,11 @@ func (t Tool) Invoke(ctx context.Context, resourceMgr tools.SourceProvider, para 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 diff --git a/internal/tools/spanneradmin/spannercreateinstance/spannercreateinstance_test.go b/internal/tools/spanneradmin/spannercreateinstance/spannercreateinstance_test.go index 9465dc0ba30..0dceb932829 100644 --- a/internal/tools/spanneradmin/spannercreateinstance/spannercreateinstance_test.go +++ b/internal/tools/spanneradmin/spannercreateinstance/spannercreateinstance_test.go @@ -73,7 +73,7 @@ func TestParseFromYaml(t *testing.T) { } } -func TestInvoke_NodeCountAndProcessingUnitsValidation(t *testing.T) { +func TestInvokeNodeCountAndProcessingUnitsValidation(t *testing.T) { t.Parallel() testCases := []struct { name string From de7c65eb1c78932b8a2525c4b76bebde50f15615 Mon Sep 17 00:00:00 2001 From: gRedHeadphone Date: Fri, 9 Jan 2026 05:21:19 +0000 Subject: [PATCH 16/17] chore: header update + main merge fixes --- internal/prebuiltconfigs/prebuiltconfigs_test.go | 2 +- internal/prebuiltconfigs/tools/spanner-admin.yaml | 2 +- internal/sources/spanneradmin/spanneradmin.go | 2 +- internal/sources/spanneradmin/spanneradmin_test.go | 2 +- .../spannercreateinstance/spannercreateinstance.go | 7 ++++++- .../spannercreateinstance/spannercreateinstance_test.go | 2 +- tests/spanneradmin/spanneradmin_integration_test.go | 2 +- 7 files changed, 12 insertions(+), 7 deletions(-) diff --git a/internal/prebuiltconfigs/prebuiltconfigs_test.go b/internal/prebuiltconfigs/prebuiltconfigs_test.go index 8e83a512e3c..936cd647faf 100644 --- a/internal/prebuiltconfigs/prebuiltconfigs_test.go +++ b/internal/prebuiltconfigs/prebuiltconfigs_test.go @@ -49,8 +49,8 @@ var expectedToolSources = []string{ "postgres", "serverless-spark", "singlestore", - "spanner-admin", "snowflake", + "spanner-admin", "spanner-postgres", "spanner", "sqlite", diff --git a/internal/prebuiltconfigs/tools/spanner-admin.yaml b/internal/prebuiltconfigs/tools/spanner-admin.yaml index c68336124a8..b58d82b5957 100644 --- a/internal/prebuiltconfigs/tools/spanner-admin.yaml +++ b/internal/prebuiltconfigs/tools/spanner-admin.yaml @@ -1,4 +1,4 @@ -# Copyright 2025 Google LLC +# 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. diff --git a/internal/sources/spanneradmin/spanneradmin.go b/internal/sources/spanneradmin/spanneradmin.go index 9ffede63a47..3c89e8e62ec 100644 --- a/internal/sources/spanneradmin/spanneradmin.go +++ b/internal/sources/spanneradmin/spanneradmin.go @@ -1,4 +1,4 @@ -// Copyright 2025 Google LLC +// 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. diff --git a/internal/sources/spanneradmin/spanneradmin_test.go b/internal/sources/spanneradmin/spanneradmin_test.go index d42c43ed999..bdbf843c143 100644 --- a/internal/sources/spanneradmin/spanneradmin_test.go +++ b/internal/sources/spanneradmin/spanneradmin_test.go @@ -1,4 +1,4 @@ -// Copyright 2025 Google LLC +// 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. diff --git a/internal/tools/spanneradmin/spannercreateinstance/spannercreateinstance.go b/internal/tools/spanneradmin/spannercreateinstance/spannercreateinstance.go index fcc5ae9bc09..68fea9b7ea3 100644 --- a/internal/tools/spanneradmin/spannercreateinstance/spannercreateinstance.go +++ b/internal/tools/spanneradmin/spannercreateinstance/spannercreateinstance.go @@ -1,4 +1,4 @@ -// Copyright 2025 Google LLC +// 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. @@ -21,6 +21,7 @@ import ( 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" @@ -200,6 +201,10 @@ func (t Tool) ParseParams(data map[string]any, claims map[string]map[string]any) 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 diff --git a/internal/tools/spanneradmin/spannercreateinstance/spannercreateinstance_test.go b/internal/tools/spanneradmin/spannercreateinstance/spannercreateinstance_test.go index 0dceb932829..a302f271731 100644 --- a/internal/tools/spanneradmin/spannercreateinstance/spannercreateinstance_test.go +++ b/internal/tools/spanneradmin/spannercreateinstance/spannercreateinstance_test.go @@ -1,4 +1,4 @@ -// Copyright 2025 Google LLC +// 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. diff --git a/tests/spanneradmin/spanneradmin_integration_test.go b/tests/spanneradmin/spanneradmin_integration_test.go index f886e485768..e76f945cb0c 100644 --- a/tests/spanneradmin/spanneradmin_integration_test.go +++ b/tests/spanneradmin/spanneradmin_integration_test.go @@ -1,4 +1,4 @@ -// Copyright 2025 Google LLC +// 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. From caa33d41b23e2a01c883ffd37e95f3ddcbff6bc0 Mon Sep 17 00:00:00 2001 From: gRedHeadphone Date: Wed, 28 Jan 2026 09:35:41 +0000 Subject: [PATCH 17/17] chore: fix for configuration file changes --- docs/en/resources/sources/spanner-admin.md | 17 ++--- .../spanneradmin/spanner-create-instance.md | 20 +++--- internal/sources/spanneradmin/spanneradmin.go | 16 ++--- .../sources/spanneradmin/spanneradmin_test.go | 63 ++++++++----------- .../spannercreateinstance.go | 22 ++++--- .../spannercreateinstance_test.go | 20 +++--- .../spanneradmin_integration_test.go | 4 +- 7 files changed, 75 insertions(+), 87 deletions(-) diff --git a/docs/en/resources/sources/spanner-admin.md b/docs/en/resources/sources/spanner-admin.md index 8926315d73d..91aabf45ddc 100644 --- a/docs/en/resources/sources/spanner-admin.md +++ b/docs/en/resources/sources/spanner-admin.md @@ -24,19 +24,20 @@ Authentication can be handled in two ways: ## Example ```yaml -sources: - my-spanner-admin: - kind: spanner-admin - - my-oauth-spanner-admin: - kind: spanner-admin - useClientOAuth: true +kind: sources +name: my-spanner-admin +type: spanner-admin +--- +kind: sources +name: my-oauth-spanner-admin +type: spanner-admin +useClientOAuth: true ``` ## Reference | **field** | **type** | **required** | **description** | | -------------- | :------: | :----------: | ---------------------------------------------------------------------------------------------------------------------------------------------- | -| kind | string | true | Must be "spanner-admin". | +| type | 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`. | diff --git a/docs/en/resources/tools/spanneradmin/spanner-create-instance.md b/docs/en/resources/tools/spanneradmin/spanner-create-instance.md index 80ed99f785f..5f98079980c 100644 --- a/docs/en/resources/tools/spanneradmin/spanner-create-instance.md +++ b/docs/en/resources/tools/spanneradmin/spanner-create-instance.md @@ -18,15 +18,15 @@ 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." +kind: sources +name: my-spanner-admin-source +type: spanner-admin +--- +kind: tools +name: create_my_spanner_instance +type: spanner-create-instance +source: my-spanner-admin-source +description: "Creates a Spanner instance." ``` ## Parameters @@ -47,6 +47,6 @@ The `spanner-create-instance` tool has the following parameters: | **field** | **type** | **required** | **description** | | ----------- | :------: | :----------: | ------------------------------------------------------------ | -| kind | string | true | Must be `spanner-create-instance`. | +| type | 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. | diff --git a/internal/sources/spanneradmin/spanneradmin.go b/internal/sources/spanneradmin/spanneradmin.go index 3c89e8e62ec..ac045e64553 100644 --- a/internal/sources/spanneradmin/spanneradmin.go +++ b/internal/sources/spanneradmin/spanneradmin.go @@ -27,14 +27,14 @@ import ( "google.golang.org/api/option" ) -const SourceKind string = "spanner-admin" +const SourceType 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)) + if !sources.Register(SourceType, newConfig) { + panic(fmt.Sprintf("source type %q already registered", SourceType)) } } @@ -48,13 +48,13 @@ func newConfig(ctx context.Context, name string, decoder *yaml.Decoder) (sources type Config struct { Name string `yaml:"name" validate:"required"` - Kind string `yaml:"kind" validate:"required"` + Type string `yaml:"type" validate:"required"` DefaultProject string `yaml:"defaultProject"` UseClientOAuth bool `yaml:"useClientOAuth"` } -func (r Config) SourceConfigKind() string { - return SourceKind +func (r Config) SourceConfigType() string { + return SourceType } // Initialize initializes a Spanner Admin Source instance. @@ -87,8 +87,8 @@ type Source struct { Client *instance.InstanceAdminClient } -func (s *Source) SourceKind() string { - return SourceKind +func (s *Source) SourceType() string { + return SourceType } func (s *Source) ToConfig() sources.SourceConfig { diff --git a/internal/sources/spanneradmin/spanneradmin_test.go b/internal/sources/spanneradmin/spanneradmin_test.go index bdbf843c143..444f242527c 100644 --- a/internal/sources/spanneradmin/spanneradmin_test.go +++ b/internal/sources/spanneradmin/spanneradmin_test.go @@ -15,9 +15,9 @@ package spanneradmin_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/sources" @@ -26,7 +26,6 @@ import ( ) func TestParseFromYamlSpannerAdmin(t *testing.T) { - t.Parallel() tcs := []struct { desc string in string @@ -35,14 +34,14 @@ func TestParseFromYamlSpannerAdmin(t *testing.T) { { desc: "basic example", in: ` - sources: - my-spanner-admin-instance: - kind: spanner-admin + kind: sources + name: my-spanner-admin-instance + type: spanner-admin `, want: map[string]sources.SourceConfig{ "my-spanner-admin-instance": spanneradmin.Config{ Name: "my-spanner-admin-instance", - Kind: spanneradmin.SourceKind, + Type: spanneradmin.SourceType, UseClientOAuth: false, }, }, @@ -50,34 +49,28 @@ func TestParseFromYamlSpannerAdmin(t *testing.T) { { desc: "use client auth example", in: ` - sources: - my-spanner-admin-instance: - kind: spanner-admin - useClientOAuth: true + kind: sources + name: my-spanner-admin-instance + type: spanner-admin + useClientOAuth: true `, want: map[string]sources.SourceConfig{ "my-spanner-admin-instance": spanneradmin.Config{ Name: "my-spanner-admin-instance", - Kind: spanneradmin.SourceKind, + Type: spanneradmin.SourceType, 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) + got, _, _, _, _, _, err := server.UnmarshalResourceConfig(context.Background(), testutils.FormatYaml(tc.in)) 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) + if !cmp.Equal(tc.want, got) { + t.Fatalf("incorrect parse: want %v, got %v", tc.want, got) } }) } @@ -93,36 +86,30 @@ func TestFailParseFromYaml(t *testing.T) { { desc: "extra field", in: ` - sources: - my-spanner-admin-instance: - kind: spanner-admin - project: test-project + kind: sources + name: my-spanner-admin-instance + type: 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 + err: `error unmarshaling sources: unable to parse source "my-spanner-admin-instance" as "spanner-admin": [2:1] unknown field "project" + 1 | name: my-spanner-admin-instance > 2 | project: test-project ^ -`, + 3 | type: spanner-admin`, }, { desc: "missing required field", in: ` - sources: - my-spanner-admin-instance: - useClientOAuth: true + kind: sources + name: my-spanner-admin-instance + useClientOAuth: true `, - err: "missing 'kind' field for source \"my-spanner-admin-instance\"", + err: "error unmarshaling sources: missing 'type' field or it is not a string", }, } 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) + _, _, _, _, _, _, err := server.UnmarshalResourceConfig(context.Background(), testutils.FormatYaml(tc.in)) if err == nil { t.Fatalf("expect parsing to fail") } diff --git a/internal/tools/spanneradmin/spannercreateinstance/spannercreateinstance.go b/internal/tools/spanneradmin/spannercreateinstance/spannercreateinstance.go index 68fea9b7ea3..610f34d30e9 100644 --- a/internal/tools/spanneradmin/spannercreateinstance/spannercreateinstance.go +++ b/internal/tools/spanneradmin/spannercreateinstance/spannercreateinstance.go @@ -27,11 +27,11 @@ import ( "github.com/googleapis/genai-toolbox/internal/util/parameters" ) -const kind string = "spanner-create-instance" +const resourceType string = "spanner-create-instance" func init() { - if !tools.Register(kind, newConfig) { - panic(fmt.Sprintf("tool kind %q already registered", kind)) + if !tools.Register(resourceType, newConfig) { + panic(fmt.Sprintf("tool type %q already registered", resourceType)) } } @@ -52,7 +52,7 @@ type compatibleSource interface { // Config defines the configuration for the create-instance tool. type Config struct { Name string `yaml:"name" validate:"required"` - Kind string `yaml:"kind" validate:"required"` + Type string `yaml:"type" validate:"required"` Description string `yaml:"description"` Source string `yaml:"source" validate:"required"` AuthRequired []string `yaml:"authRequired"` @@ -62,8 +62,8 @@ type Config struct { var _ tools.ToolConfig = Config{} // ToolConfigKind returns the kind of the tool. -func (cfg Config) ToolConfigKind() string { - return kind +func (cfg Config) ToolConfigType() string { + return resourceType } // Initialize initializes the tool from the configuration. @@ -74,7 +74,7 @@ func (cfg Config) Initialize(srcs map[string]sources.Source) (tools.Tool, error) } s, ok := rawS.(compatibleSource) if !ok { - return nil, fmt.Errorf("invalid source for %q tool: source %q not compatible", kind, cfg.Source) + return nil, fmt.Errorf("invalid source for %q tool: source %q not compatible", resourceType, cfg.Source) } project := s.GetDefaultProject() @@ -139,7 +139,7 @@ func (t Tool) Invoke(ctx context.Context, resourceMgr tools.SourceProvider, para 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) + source, err := tools.GetCompatibleSource[compatibleSource](resourceMgr, t.Source, t.Name, t.Type) if err != nil { return nil, err } @@ -221,7 +221,7 @@ func (t Tool) Authorized(verifiedAuthServices []string) bool { } func (t Tool) RequiresClientAuthorization(resourceMgr tools.SourceProvider) (bool, error) { - source, err := tools.GetCompatibleSource[compatibleSource](resourceMgr, t.Source, t.Name, t.Kind) + source, err := tools.GetCompatibleSource[compatibleSource](resourceMgr, t.Source, t.Name, t.Type) if err != nil { return false, err } @@ -231,3 +231,7 @@ func (t Tool) RequiresClientAuthorization(resourceMgr tools.SourceProvider) (boo func (t Tool) GetAuthTokenHeaderName(resourceMgr tools.SourceProvider) (string, error) { return "Authorization", nil } + +func (t Tool) GetParameters() parameters.Parameters { + return t.AllParams +} diff --git a/internal/tools/spanneradmin/spannercreateinstance/spannercreateinstance_test.go b/internal/tools/spanneradmin/spannercreateinstance/spannercreateinstance_test.go index a302f271731..452759e0bb3 100644 --- a/internal/tools/spanneradmin/spannercreateinstance/spannercreateinstance_test.go +++ b/internal/tools/spanneradmin/spannercreateinstance/spannercreateinstance_test.go @@ -18,7 +18,6 @@ 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" @@ -39,16 +38,16 @@ func TestParseFromYaml(t *testing.T) { { desc: "basic example", in: ` - tools: - create-instance-tool: - kind: spanner-create-instance - description: a test description - source: a-source + kind: tools + name: create-instance-tool + type: 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", + Type: "spanner-create-instance", Description: "a test description", Source: "a-source", AuthRequired: []string{}, @@ -58,15 +57,12 @@ func TestParseFromYaml(t *testing.T) { } 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) + _, _, _, got, _, _, err := server.UnmarshalResourceConfig(ctx, testutils.FormatYaml(tc.in)) if err != nil { t.Fatalf("unable to unmarshal: %s", err) } - if diff := cmp.Diff(tc.want, got.Tools); diff != "" { + if diff := cmp.Diff(tc.want, got); diff != "" { t.Fatalf("incorrect parse: diff %v", diff) } }) diff --git a/tests/spanneradmin/spanneradmin_integration_test.go b/tests/spanneradmin/spanneradmin_integration_test.go index e76f945cb0c..5236b29a68a 100644 --- a/tests/spanneradmin/spanneradmin_integration_test.go +++ b/tests/spanneradmin/spanneradmin_integration_test.go @@ -44,7 +44,7 @@ func getSpannerAdminVars(t *testing.T) map[string]any { } return map[string]any{ - "kind": "spanner-admin", + "type": "spanner-admin", "defaultProject": SpannerProject, } } @@ -90,7 +90,7 @@ func TestSpannerAdminCreateInstance(t *testing.T) { }, "tools": map[string]any{ "create-instance-tool": map[string]any{ - "kind": "spanner-create-instance", + "type": "spanner-create-instance", "source": "my-spanner-admin", "description": "Creates a Spanner instance.", },