diff --git a/.ci/integration.cloudbuild.yaml b/.ci/integration.cloudbuild.yaml index 604be42499f..24690994c36 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 ce32662cea8..7309f648552 100644 --- a/cmd/root.go +++ b/cmd/root.go @@ -224,6 +224,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" @@ -270,6 +271,7 @@ 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" 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..91aabf45ddc --- /dev/null +++ b/docs/en/resources/sources/spanner-admin.md @@ -0,0 +1,43 @@ +--- +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 +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** | +| -------------- | :------: | :----------: | ---------------------------------------------------------------------------------------------------------------------------------------------- | +| 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 new file mode 100644 index 00000000000..5f98079980c --- /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 +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 + +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** | +| ----------- | :------: | :----------: | ------------------------------------------------------------ | +| 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/prebuiltconfigs/prebuiltconfigs_test.go b/internal/prebuiltconfigs/prebuiltconfigs_test.go index 8bc6ef1e41e..936cd647faf 100644 --- a/internal/prebuiltconfigs/prebuiltconfigs_test.go +++ b/internal/prebuiltconfigs/prebuiltconfigs_test.go @@ -50,6 +50,7 @@ var expectedToolSources = []string{ "serverless-spark", "singlestore", "snowflake", + "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..b58d82b5957 --- /dev/null +++ b/internal/prebuiltconfigs/tools/spanner-admin.yaml @@ -0,0 +1,27 @@ +# 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 diff --git a/internal/sources/spanneradmin/spanneradmin.go b/internal/sources/spanneradmin/spanneradmin.go new file mode 100644 index 00000000000..ac045e64553 --- /dev/null +++ b/internal/sources/spanneradmin/spanneradmin.go @@ -0,0 +1,120 @@ +// 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 SourceType string = "spanner-admin" + +// validate interface +var _ sources.SourceConfig = Config{} + +func init() { + if !sources.Register(SourceType, newConfig) { + panic(fmt.Sprintf("source type %q already registered", SourceType)) + } +} + +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"` + Type string `yaml:"type" validate:"required"` + DefaultProject string `yaml:"defaultProject"` + UseClientOAuth bool `yaml:"useClientOAuth"` +} + +func (r Config) SourceConfigType() string { + return SourceType +} + +// 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) SourceType() string { + return SourceType +} + +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 +} diff --git a/internal/sources/spanneradmin/spanneradmin_test.go b/internal/sources/spanneradmin/spanneradmin_test.go new file mode 100644 index 00000000000..444f242527c --- /dev/null +++ b/internal/sources/spanneradmin/spanneradmin_test.go @@ -0,0 +1,122 @@ +// 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 ( + "context" + "testing" + + "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) { + tcs := []struct { + desc string + in string + want server.SourceConfigs + }{ + { + desc: "basic example", + in: ` + 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", + Type: spanneradmin.SourceType, + UseClientOAuth: false, + }, + }, + }, + { + desc: "use client auth example", + in: ` + 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", + Type: spanneradmin.SourceType, + UseClientOAuth: true, + }, + }, + }, + } + for _, tc := range tcs { + t.Run(tc.desc, func(t *testing.T) { + 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) { + t.Fatalf("incorrect parse: want %v, got %v", tc.want, got) + } + }) + } +} + +func TestFailParseFromYaml(t *testing.T) { + t.Parallel() + tcs := []struct { + desc string + in string + err string + }{ + { + desc: "extra field", + in: ` + kind: sources + name: my-spanner-admin-instance + type: spanner-admin + project: test-project + `, + 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: ` + kind: sources + name: my-spanner-admin-instance + useClientOAuth: true + `, + err: "error unmarshaling sources: missing 'type' field or it is not a string", + }, + } + for _, tc := range tcs { + t.Run(tc.desc, func(t *testing.T) { + _, _, _, _, _, _, err := server.UnmarshalResourceConfig(context.Background(), testutils.FormatYaml(tc.in)) + 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..610f34d30e9 --- /dev/null +++ b/internal/tools/spanneradmin/spannercreateinstance/spannercreateinstance.go @@ -0,0 +1,237 @@ +// 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 resourceType string = "spanner-create-instance" + +func init() { + if !tools.Register(resourceType, newConfig) { + panic(fmt.Sprintf("tool type %q already registered", resourceType)) + } +} + +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"` + Type string `yaml:"type" 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) ToolConfigType() string { + return resourceType +} + +// 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", resourceType, 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.Type) + 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.Type) + if err != nil { + return false, err + } + return source.UseClientAuthorization(), nil +} + +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 new file mode 100644 index 00000000000..452759e0bb3 --- /dev/null +++ b/internal/tools/spanneradmin/spannercreateinstance/spannercreateinstance_test.go @@ -0,0 +1,106 @@ +// 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" + + "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: ` + 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", + Type: "spanner-create-instance", + Description: "a test description", + Source: "a-source", + AuthRequired: []string{}, + }, + }, + }, + } + for _, tc := range tcs { + t.Run(tc.desc, func(t *testing.T) { + // Parse contents + _, _, _, 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); 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) + } + }) + } +} diff --git a/tests/spanneradmin/spanneradmin_integration_test.go b/tests/spanneradmin/spanneradmin_integration_test.go new file mode 100644 index 00000000000..5236b29a68a --- /dev/null +++ b/tests/spanneradmin/spanneradmin_integration_test.go @@ -0,0 +1,178 @@ +// 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{ + "type": "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{ + "type": "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) + } +}