mirror of
https://github.com/googleapis/genai-toolbox.git
synced 2026-01-27 16:28:15 -05:00
Compare commits
4 Commits
spanner-cr
...
averikitsc
| Author | SHA1 | Date | |
|---|---|---|---|
|
d4cb4e81b4 | ||
|
|
f6474739e3 | ||
|
|
1d7c498116 | ||
|
|
9294ce39c8 |
@@ -340,26 +340,6 @@ steps:
|
||||
spanner \
|
||||
spanner || echo "Integration tests failed." # ignore test failures
|
||||
|
||||
- id: "spanner-admin"
|
||||
name: golang:1
|
||||
waitFor: ["compile-test-binary"]
|
||||
entrypoint: /bin/bash
|
||||
env:
|
||||
- "GOPATH=/gopath"
|
||||
- "SPANNER_PROJECT=$PROJECT_ID"
|
||||
- "SERVICE_ACCOUNT_EMAIL=$SERVICE_ACCOUNT_EMAIL"
|
||||
secretEnv: ["CLIENT_ID"]
|
||||
volumes:
|
||||
- name: "go"
|
||||
path: "/gopath"
|
||||
args:
|
||||
- -c
|
||||
- |
|
||||
.ci/test_with_coverage.sh \
|
||||
"Spanner Admin" \
|
||||
spanneradmin \
|
||||
spanneradmin || echo "Integration tests failed."
|
||||
|
||||
- id: "neo4j"
|
||||
name: golang:1
|
||||
waitFor: ["compile-test-binary"]
|
||||
|
||||
4
.github/workflows/deploy_dev_docs.yaml
vendored
4
.github/workflows/deploy_dev_docs.yaml
vendored
@@ -51,12 +51,12 @@ jobs:
|
||||
extended: true
|
||||
|
||||
- name: Setup Node
|
||||
uses: actions/setup-node@395ad3262231945c25e8478fd5baf05154b1d79f # v6
|
||||
uses: actions/setup-node@6044e13b5dc448c55e2357c09f80417699197238 # v6
|
||||
with:
|
||||
node-version: "22"
|
||||
|
||||
- name: Cache dependencies
|
||||
uses: actions/cache@9255dc7a253b0ccc959486e2bca901246202afeb # v5
|
||||
uses: actions/cache@8b402f58fbc84540c8b491a91e594a4576fec3d7 # v5
|
||||
with:
|
||||
path: ~/.npm
|
||||
key: ${{ runner.os }}-node-${{ hashFiles('**/package-lock.json') }}
|
||||
|
||||
@@ -57,7 +57,7 @@ jobs:
|
||||
with:
|
||||
hugo-version: "0.145.0"
|
||||
extended: true
|
||||
- uses: actions/setup-node@395ad3262231945c25e8478fd5baf05154b1d79f # v6
|
||||
- uses: actions/setup-node@6044e13b5dc448c55e2357c09f80417699197238 # v6
|
||||
with:
|
||||
node-version: "22"
|
||||
|
||||
|
||||
2
.github/workflows/deploy_versioned_docs.yaml
vendored
2
.github/workflows/deploy_versioned_docs.yaml
vendored
@@ -44,7 +44,7 @@ jobs:
|
||||
extended: true
|
||||
|
||||
- name: Setup Node
|
||||
uses: actions/setup-node@395ad3262231945c25e8478fd5baf05154b1d79f # v6
|
||||
uses: actions/setup-node@6044e13b5dc448c55e2357c09f80417699197238 # v6
|
||||
with:
|
||||
node-version: "22"
|
||||
|
||||
|
||||
4
.github/workflows/docs_preview_deploy.yaml
vendored
4
.github/workflows/docs_preview_deploy.yaml
vendored
@@ -62,12 +62,12 @@ jobs:
|
||||
extended: true
|
||||
|
||||
- name: Setup Node
|
||||
uses: actions/setup-node@395ad3262231945c25e8478fd5baf05154b1d79f # v6
|
||||
uses: actions/setup-node@6044e13b5dc448c55e2357c09f80417699197238 # v6
|
||||
with:
|
||||
node-version: "22"
|
||||
|
||||
- name: Cache dependencies
|
||||
uses: actions/cache@9255dc7a253b0ccc959486e2bca901246202afeb # v5
|
||||
uses: actions/cache@8b402f58fbc84540c8b491a91e594a4576fec3d7 # v5
|
||||
with:
|
||||
path: ~/.npm
|
||||
key: ${{ runner.os }}-node-${{ hashFiles('**/package-lock.json') }}
|
||||
|
||||
3
.github/workflows/link_checker_workflow.yaml
vendored
3
.github/workflows/link_checker_workflow.yaml
vendored
@@ -25,7 +25,7 @@ jobs:
|
||||
uses: actions/checkout@8e8c483db84b4bee98b60c0593521ed34d9990e8 # v6
|
||||
|
||||
- name: Restore lychee cache
|
||||
uses: actions/cache@9255dc7a253b0ccc959486e2bca901246202afeb # v5
|
||||
uses: actions/cache@8b402f58fbc84540c8b491a91e594a4576fec3d7 # v5
|
||||
with:
|
||||
path: .lycheecache
|
||||
key: cache-lychee-${{ github.sha }}
|
||||
@@ -39,6 +39,7 @@ jobs:
|
||||
--no-progress
|
||||
--cache
|
||||
--max-cache-age 1d
|
||||
--exclude '^neo4j\+.*' --exclude '^bolt://.*'
|
||||
README.md
|
||||
docs/
|
||||
output: /tmp/foo.txt
|
||||
|
||||
4
.github/workflows/lint.yaml
vendored
4
.github/workflows/lint.yaml
vendored
@@ -51,11 +51,11 @@ jobs:
|
||||
console.log('Failed to remove label. Another job may have already removed it!');
|
||||
}
|
||||
- name: Setup Go
|
||||
uses: actions/setup-go@4dc6199c7b1a012772edbd06daecab0f50c9053c # v6.1.0
|
||||
uses: actions/setup-go@7a3fe6cf4cb3a834922a1244abfce67bcef6a0c5 # v6.2.0
|
||||
with:
|
||||
go-version: "1.25"
|
||||
- name: Checkout code
|
||||
uses: actions/checkout@8e8c483db84b4bee98b60c0593521ed34d9990e8 # v6.0.1
|
||||
uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2
|
||||
with:
|
||||
ref: ${{ github.event.pull_request.head.sha }}
|
||||
repository: ${{ github.event.pull_request.head.repo.full_name }}
|
||||
|
||||
2
.github/workflows/sync-labels.yaml
vendored
2
.github/workflows/sync-labels.yaml
vendored
@@ -29,7 +29,7 @@ jobs:
|
||||
issues: 'write'
|
||||
pull-requests: 'write'
|
||||
steps:
|
||||
- uses: actions/checkout@8e8c483db84b4bee98b60c0593521ed34d9990e8 # v6.0.1
|
||||
- uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2
|
||||
- uses: micnncim/action-label-syncer@3abd5ab72fda571e69fffd97bd4e0033dd5f495c # v1.3.0
|
||||
env:
|
||||
GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
|
||||
|
||||
4
.github/workflows/tests.yaml
vendored
4
.github/workflows/tests.yaml
vendored
@@ -57,12 +57,12 @@ jobs:
|
||||
}
|
||||
|
||||
- name: Setup Go
|
||||
uses: actions/setup-go@4dc6199c7b1a012772edbd06daecab0f50c9053c # v6.1.0
|
||||
uses: actions/setup-go@7a3fe6cf4cb3a834922a1244abfce67bcef6a0c5 # v6.2.0
|
||||
with:
|
||||
go-version: "1.24"
|
||||
|
||||
- name: Checkout code
|
||||
uses: actions/checkout@8e8c483db84b4bee98b60c0593521ed34d9990e8 # v6.0.1
|
||||
uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2
|
||||
with:
|
||||
ref: ${{ github.event.pull_request.head.sha }}
|
||||
repository: ${{ github.event.pull_request.head.repo.full_name }}
|
||||
|
||||
@@ -39,7 +39,7 @@ https://dev.mysql.com/doc/refman/8.4/en/user-names.html
|
||||
# npmjs links can occasionally trigger rate limiting during high-frequency CI builds
|
||||
https://www.npmjs.com/package/@toolbox-sdk/core
|
||||
https://www.npmjs.com/package/@toolbox-sdk/adk
|
||||
|
||||
https://www.oceanbase.com/
|
||||
|
||||
# Ignore social media and blog profiles to reduce external request overhead
|
||||
https://medium.com/@mcp_toolbox
|
||||
https://medium.com/@mcp_toolbox
|
||||
|
||||
@@ -223,7 +223,6 @@ import (
|
||||
_ "github.com/googleapis/genai-toolbox/internal/tools/spanner/spannerlistgraphs"
|
||||
_ "github.com/googleapis/genai-toolbox/internal/tools/spanner/spannerlisttables"
|
||||
_ "github.com/googleapis/genai-toolbox/internal/tools/spanner/spannersql"
|
||||
_ "github.com/googleapis/genai-toolbox/internal/tools/spanneradmin/spannercreateinstance"
|
||||
_ "github.com/googleapis/genai-toolbox/internal/tools/sqlite/sqliteexecutesql"
|
||||
_ "github.com/googleapis/genai-toolbox/internal/tools/sqlite/sqlitesql"
|
||||
_ "github.com/googleapis/genai-toolbox/internal/tools/tidb/tidbexecutesql"
|
||||
@@ -270,7 +269,6 @@ import (
|
||||
_ "github.com/googleapis/genai-toolbox/internal/sources/singlestore"
|
||||
_ "github.com/googleapis/genai-toolbox/internal/sources/snowflake"
|
||||
_ "github.com/googleapis/genai-toolbox/internal/sources/spanner"
|
||||
_ "github.com/googleapis/genai-toolbox/internal/sources/spanneradmin"
|
||||
_ "github.com/googleapis/genai-toolbox/internal/sources/sqlite"
|
||||
_ "github.com/googleapis/genai-toolbox/internal/sources/tidb"
|
||||
_ "github.com/googleapis/genai-toolbox/internal/sources/trino"
|
||||
|
||||
@@ -1,59 +0,0 @@
|
||||
# Cloud Spanner Admin MCP Server
|
||||
|
||||
The Cloud Spanner Admin Model Context Protocol (MCP) Server gives AI-powered development tools the ability to manage your Google Cloud Spanner infrastructure. It supports creating instances.
|
||||
|
||||
## Features
|
||||
|
||||
An editor configured to use the Cloud Spanner Admin MCP server can use its AI capabilities to help you:
|
||||
|
||||
- **Provision & Manage Infrastructure** - Create Cloud Spanner instances
|
||||
|
||||
## Prerequisites
|
||||
|
||||
* [Node.js](https://nodejs.org/) installed.
|
||||
* A Google Cloud project with the **Cloud Spanner Admin API** enabled.
|
||||
* Ensure [Application Default Credentials](https://cloud.google.com/docs/authentication/gcloud) are available in your environment.
|
||||
* IAM Permissions:
|
||||
* Cloud Spanner Admin (`roles/spanner.admin`)
|
||||
|
||||
## Install & Configuration
|
||||
|
||||
In the Antigravity MCP Store, click the "Install" button.
|
||||
|
||||
You'll now be able to see all enabled tools in the "Tools" tab.
|
||||
|
||||
> [!NOTE]
|
||||
> If you encounter issues with Windows Defender blocking the execution, you may need to configure an allowlist. See [Configure exclusions for Microsoft Defender Antivirus](https://learn.microsoft.com/en-us/microsoft-365/security/defender-endpoint/configure-exclusions-microsoft-defender-antivirus?view=o365-worldwide) for more details.
|
||||
|
||||
## Usage
|
||||
|
||||
Once configured, the MCP server will automatically provide Cloud Spanner Admin capabilities to your AI assistant. You can:
|
||||
|
||||
* "Create a new Spanner instance named 'my-spanner-instance' in the 'my-gcp-project' project with config 'regional-us-central1', edition 'ENTERPRISE', and 1 node."
|
||||
|
||||
## Server Capabilities
|
||||
|
||||
The Cloud Spanner Admin MCP server provides the following tools:
|
||||
|
||||
| Tool Name | Description |
|
||||
|:------------------|:---------------------------------|
|
||||
| `create_instance` | Create a Cloud Spanner instance. |
|
||||
|
||||
## Custom MCP Server Configuration
|
||||
|
||||
Add the following configuration to your MCP client (e.g., `settings.json` for Gemini CLI, `mcp_config.json` for Antigravity):
|
||||
|
||||
```json
|
||||
{
|
||||
"mcpServers": {
|
||||
"spanner-admin": {
|
||||
"command": "npx",
|
||||
"args": ["-y", "@toolbox-sdk/server", "--prebuilt", "spanner-admin", "--stdio"]
|
||||
}
|
||||
}
|
||||
}
|
||||
```
|
||||
|
||||
## Documentation
|
||||
|
||||
For more information, visit the [Cloud Spanner Admin API documentation](https://cloud.google.com/spanner/docs/reference/rpc/google.spanner.admin.instance.v1).
|
||||
@@ -377,10 +377,10 @@ See [Usage Examples](../reference/cli.md#examples).
|
||||
entries.
|
||||
* **Dataplex Editor** (`roles/dataplex.editor`) to modify entries.
|
||||
* **Tools:**
|
||||
* `dataplex_search_entries`: Searches for entries in Dataplex Catalog.
|
||||
* `dataplex_lookup_entry`: Retrieves a specific entry from Dataplex
|
||||
* `search_entries`: Searches for entries in Dataplex Catalog.
|
||||
* `lookup_entry`: Retrieves a specific entry from Dataplex
|
||||
Catalog.
|
||||
* `dataplex_search_aspect_types`: Finds aspect types relevant to the
|
||||
* `search_aspect_types`: Finds aspect types relevant to the
|
||||
query.
|
||||
|
||||
## Firestore
|
||||
|
||||
@@ -1,42 +0,0 @@
|
||||
---
|
||||
title: Spanner Admin
|
||||
type: docs
|
||||
weight: 1
|
||||
description: "A \"spanner-admin\" source provides a client for the Cloud Spanner Admin API.\n"
|
||||
alias: [/resources/sources/spanner-admin]
|
||||
---
|
||||
|
||||
## About
|
||||
|
||||
The `spanner-admin` source provides a client to interact with the [Google
|
||||
Cloud Spanner Admin API](https://cloud.google.com/spanner/docs/reference/rpc/google.spanner.admin.instance.v1). This
|
||||
allows tools to perform administrative tasks on Spanner instances, such as
|
||||
creating instances.
|
||||
|
||||
Authentication can be handled in two ways:
|
||||
|
||||
1. **Application Default Credentials (ADC):** By default, the source uses ADC
|
||||
to authenticate with the API.
|
||||
2. **Client-side OAuth:** If `useClientOAuth` is set to `true`, the source will
|
||||
expect an OAuth 2.0 access token to be provided by the client (e.g., a web
|
||||
browser) for each request.
|
||||
|
||||
## Example
|
||||
|
||||
```yaml
|
||||
sources:
|
||||
my-spanner-admin:
|
||||
kind: spanner-admin
|
||||
|
||||
my-oauth-spanner-admin:
|
||||
kind: spanner-admin
|
||||
useClientOAuth: true
|
||||
```
|
||||
|
||||
## Reference
|
||||
|
||||
| **field** | **type** | **required** | **description** |
|
||||
| -------------- | :------: | :----------: | ---------------------------------------------------------------------------------------------------------------------------------------------- |
|
||||
| kind | string | true | Must be "spanner-admin". |
|
||||
| defaultProject | string | false | The Google Cloud project ID to use for Spanner infrastructure tools. |
|
||||
| useClientOAuth | boolean | false | If true, the source will use client-side OAuth for authorization. Otherwise, it will use Application Default Credentials. Defaults to `false`. |
|
||||
@@ -1,52 +0,0 @@
|
||||
---
|
||||
title: spanner-create-instance
|
||||
type: docs
|
||||
weight: 2
|
||||
description: "Create a Cloud Spanner instance."
|
||||
---
|
||||
|
||||
The `spanner-create-instance` tool creates a new Cloud Spanner instance in a
|
||||
specified Google Cloud project.
|
||||
|
||||
{{< notice info >}}
|
||||
This tool uses the `spanner-admin` source.
|
||||
{{< /notice >}}
|
||||
|
||||
## Configuration
|
||||
|
||||
Here is an example of how to configure the `spanner-create-instance` tool in
|
||||
your `tools.yaml` file:
|
||||
|
||||
```yaml
|
||||
sources:
|
||||
my-spanner-admin-source:
|
||||
kind: spanner-admin
|
||||
|
||||
tools:
|
||||
create_my_spanner_instance:
|
||||
kind: spanner-create-instance
|
||||
source: my-spanner-admin-source
|
||||
description: "Creates a Spanner instance."
|
||||
```
|
||||
|
||||
## Parameters
|
||||
|
||||
The `spanner-create-instance` tool has the following parameters:
|
||||
|
||||
| **field** | **type** | **required** | **description** |
|
||||
| --------------- | :------: | :----------: | ------------------------------------------------------------------------------------ |
|
||||
| project | string | true | The Google Cloud project ID. |
|
||||
| instanceId | string | true | The ID of the instance to create. |
|
||||
| displayName | string | true | The display name of the instance. |
|
||||
| config | string | true | The instance configuration (e.g., `regional-us-central1`). |
|
||||
| nodeCount | integer | true | The number of nodes. Mutually exclusive with `processingUnits` (one must be 0). |
|
||||
| processingUnits | integer | true | The number of processing units. Mutually exclusive with `nodeCount` (one must be 0). |
|
||||
| edition | string | false | The edition of the instance (`STANDARD`, `ENTERPRISE`, `ENTERPRISE_PLUS`). |
|
||||
|
||||
## Reference
|
||||
|
||||
| **field** | **type** | **required** | **description** |
|
||||
| ----------- | :------: | :----------: | ------------------------------------------------------------ |
|
||||
| kind | string | true | Must be `spanner-create-instance`. |
|
||||
| source | string | true | The name of the `spanner-admin` source to use for this tool. |
|
||||
| description | string | false | A description of the tool that is passed to the agent. |
|
||||
2
go.mod
2
go.mod
@@ -63,6 +63,7 @@ require (
|
||||
google.golang.org/api v0.256.0
|
||||
google.golang.org/genai v1.37.0
|
||||
google.golang.org/genproto v0.0.0-20251022142026-3a174f9686a8
|
||||
google.golang.org/grpc v1.76.0
|
||||
google.golang.org/protobuf v1.36.10
|
||||
modernc.org/sqlite v1.40.0
|
||||
)
|
||||
@@ -229,7 +230,6 @@ require (
|
||||
golang.org/x/xerrors v0.0.0-20240903120638-7835f813f4da // indirect
|
||||
google.golang.org/genproto/googleapis/api v0.0.0-20251111163417-95abcf5c77ba // indirect
|
||||
google.golang.org/genproto/googleapis/rpc v0.0.0-20251103181224-f26f9409b101 // indirect
|
||||
google.golang.org/grpc v1.76.0 // indirect
|
||||
gopkg.in/inf.v0 v0.9.1 // indirect
|
||||
gopkg.in/ini.v1 v1.67.0 // indirect
|
||||
modernc.org/libc v1.66.10 // indirect
|
||||
|
||||
@@ -50,7 +50,6 @@ var expectedToolSources = []string{
|
||||
"serverless-spark",
|
||||
"singlestore",
|
||||
"snowflake",
|
||||
"spanner-admin",
|
||||
"spanner-postgres",
|
||||
"spanner",
|
||||
"sqlite",
|
||||
|
||||
@@ -1,27 +0,0 @@
|
||||
# Copyright 2026 Google LLC
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
|
||||
sources:
|
||||
spanner-admin-source:
|
||||
kind: spanner-admin
|
||||
defaultProject: ${SPANNER_PROJECT:}
|
||||
|
||||
tools:
|
||||
create_instance:
|
||||
kind: spanner-create-instance
|
||||
source: spanner-admin-source
|
||||
|
||||
toolsets:
|
||||
spanner_admin_tools:
|
||||
- create_instance
|
||||
@@ -26,7 +26,9 @@ import (
|
||||
"github.com/googleapis/genai-toolbox/internal/util"
|
||||
"go.opentelemetry.io/otel/trace"
|
||||
"golang.org/x/oauth2/google"
|
||||
"google.golang.org/api/iterator"
|
||||
"google.golang.org/api/option"
|
||||
grpcstatus "google.golang.org/grpc/status"
|
||||
)
|
||||
|
||||
const SourceKind string = "dataplex"
|
||||
@@ -173,9 +175,18 @@ func (s *Source) SearchAspectTypes(ctx context.Context, query string, pageSize i
|
||||
var results []*dataplexpb.AspectType
|
||||
for {
|
||||
entry, err := it.Next()
|
||||
if err != nil {
|
||||
|
||||
if err == iterator.Done {
|
||||
break
|
||||
}
|
||||
if err != nil {
|
||||
if st, ok := grpcstatus.FromError(err); ok {
|
||||
errorCode := st.Code()
|
||||
errorMessage := st.Message()
|
||||
return nil, fmt.Errorf("failed to search aspect types with error code: %q message: %s", errorCode.String(), errorMessage)
|
||||
}
|
||||
return nil, fmt.Errorf("failed to search aspect types: %w", err)
|
||||
}
|
||||
|
||||
// Create an instance of exponential backoff with default values for retrying GetAspectType calls
|
||||
// InitialInterval, RandomizationFactor, Multiplier, MaxInterval = 500 ms, 0.5, 1.5, 60 s
|
||||
@@ -214,9 +225,17 @@ func (s *Source) SearchEntries(ctx context.Context, query string, pageSize int,
|
||||
var results []*dataplexpb.SearchEntriesResult
|
||||
for {
|
||||
entry, err := it.Next()
|
||||
if err != nil {
|
||||
if err == iterator.Done {
|
||||
break
|
||||
}
|
||||
if err != nil {
|
||||
if st, ok := grpcstatus.FromError(err); ok {
|
||||
errorCode := st.Code()
|
||||
errorMessage := st.Message()
|
||||
return nil, fmt.Errorf("failed to search entries with error code: %q message: %s", errorCode.String(), errorMessage)
|
||||
}
|
||||
return nil, fmt.Errorf("failed to search entries: %w", err)
|
||||
}
|
||||
results = append(results, entry)
|
||||
}
|
||||
return results, nil
|
||||
|
||||
@@ -1,120 +0,0 @@
|
||||
// Copyright 2026 Google LLC
|
||||
//
|
||||
// Licensed under the Apache License, Version 2.0 (the "License");
|
||||
// you may not use this file except in compliance with the License.
|
||||
// You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software
|
||||
// distributed under the License is distributed on an "AS IS" BASIS,
|
||||
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
// See the License for the specific language governing permissions and
|
||||
// limitations under the License.
|
||||
|
||||
package spanneradmin
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
|
||||
instance "cloud.google.com/go/spanner/admin/instance/apiv1"
|
||||
"github.com/goccy/go-yaml"
|
||||
"github.com/googleapis/genai-toolbox/internal/sources"
|
||||
"github.com/googleapis/genai-toolbox/internal/util"
|
||||
"go.opentelemetry.io/otel/trace"
|
||||
"golang.org/x/oauth2"
|
||||
"google.golang.org/api/option"
|
||||
)
|
||||
|
||||
const SourceKind string = "spanner-admin"
|
||||
|
||||
// validate interface
|
||||
var _ sources.SourceConfig = Config{}
|
||||
|
||||
func init() {
|
||||
if !sources.Register(SourceKind, newConfig) {
|
||||
panic(fmt.Sprintf("source kind %q already registered", SourceKind))
|
||||
}
|
||||
}
|
||||
|
||||
func newConfig(ctx context.Context, name string, decoder *yaml.Decoder) (sources.SourceConfig, error) {
|
||||
actual := Config{Name: name}
|
||||
if err := decoder.DecodeContext(ctx, &actual); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return actual, nil
|
||||
}
|
||||
|
||||
type Config struct {
|
||||
Name string `yaml:"name" validate:"required"`
|
||||
Kind string `yaml:"kind" validate:"required"`
|
||||
DefaultProject string `yaml:"defaultProject"`
|
||||
UseClientOAuth bool `yaml:"useClientOAuth"`
|
||||
}
|
||||
|
||||
func (r Config) SourceConfigKind() string {
|
||||
return SourceKind
|
||||
}
|
||||
|
||||
// Initialize initializes a Spanner Admin Source instance.
|
||||
func (r Config) Initialize(ctx context.Context, tracer trace.Tracer) (sources.Source, error) {
|
||||
var client *instance.InstanceAdminClient
|
||||
|
||||
if !r.UseClientOAuth {
|
||||
ua, err := util.UserAgentFromContext(ctx)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("error in User Agent retrieval: %s", err)
|
||||
}
|
||||
// Use Application Default Credentials
|
||||
client, err = instance.NewInstanceAdminClient(ctx, option.WithUserAgent(ua))
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("error creating new spanner instance admin client: %w", err)
|
||||
}
|
||||
}
|
||||
|
||||
s := &Source{
|
||||
Config: r,
|
||||
Client: client,
|
||||
}
|
||||
return s, nil
|
||||
}
|
||||
|
||||
var _ sources.Source = &Source{}
|
||||
|
||||
type Source struct {
|
||||
Config
|
||||
Client *instance.InstanceAdminClient
|
||||
}
|
||||
|
||||
func (s *Source) SourceKind() string {
|
||||
return SourceKind
|
||||
}
|
||||
|
||||
func (s *Source) ToConfig() sources.SourceConfig {
|
||||
return s.Config
|
||||
}
|
||||
|
||||
func (s *Source) GetDefaultProject() string {
|
||||
return s.DefaultProject
|
||||
}
|
||||
|
||||
func (s *Source) GetClient(ctx context.Context, accessToken string) (*instance.InstanceAdminClient, error) {
|
||||
if s.UseClientOAuth {
|
||||
token := &oauth2.Token{AccessToken: accessToken}
|
||||
ua, err := util.UserAgentFromContext(ctx)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
client, err := instance.NewInstanceAdminClient(ctx, option.WithTokenSource(oauth2.StaticTokenSource(token)), option.WithUserAgent(ua))
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("error creating new spanner instance admin client: %w", err)
|
||||
}
|
||||
return client, nil
|
||||
}
|
||||
return s.Client, nil
|
||||
}
|
||||
|
||||
func (s *Source) UseClientAuthorization() bool {
|
||||
return s.UseClientOAuth
|
||||
}
|
||||
@@ -1,135 +0,0 @@
|
||||
// Copyright 2026 Google LLC
|
||||
//
|
||||
// Licensed under the Apache License, Version 2.0 (the "License");
|
||||
// you may not use this file except in compliance with the License.
|
||||
// You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software
|
||||
// distributed under the License is distributed on an "AS IS" BASIS,
|
||||
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
// See the License for the specific language governing permissions and
|
||||
// limitations under the License.
|
||||
|
||||
package spanneradmin_test
|
||||
|
||||
import (
|
||||
"testing"
|
||||
|
||||
yaml "github.com/goccy/go-yaml"
|
||||
"github.com/google/go-cmp/cmp"
|
||||
"github.com/googleapis/genai-toolbox/internal/server"
|
||||
"github.com/googleapis/genai-toolbox/internal/sources"
|
||||
"github.com/googleapis/genai-toolbox/internal/sources/spanneradmin"
|
||||
"github.com/googleapis/genai-toolbox/internal/testutils"
|
||||
)
|
||||
|
||||
func TestParseFromYamlSpannerAdmin(t *testing.T) {
|
||||
t.Parallel()
|
||||
tcs := []struct {
|
||||
desc string
|
||||
in string
|
||||
want server.SourceConfigs
|
||||
}{
|
||||
{
|
||||
desc: "basic example",
|
||||
in: `
|
||||
sources:
|
||||
my-spanner-admin-instance:
|
||||
kind: spanner-admin
|
||||
`,
|
||||
want: map[string]sources.SourceConfig{
|
||||
"my-spanner-admin-instance": spanneradmin.Config{
|
||||
Name: "my-spanner-admin-instance",
|
||||
Kind: spanneradmin.SourceKind,
|
||||
UseClientOAuth: false,
|
||||
},
|
||||
},
|
||||
},
|
||||
{
|
||||
desc: "use client auth example",
|
||||
in: `
|
||||
sources:
|
||||
my-spanner-admin-instance:
|
||||
kind: spanner-admin
|
||||
useClientOAuth: true
|
||||
`,
|
||||
want: map[string]sources.SourceConfig{
|
||||
"my-spanner-admin-instance": spanneradmin.Config{
|
||||
Name: "my-spanner-admin-instance",
|
||||
Kind: spanneradmin.SourceKind,
|
||||
UseClientOAuth: true,
|
||||
},
|
||||
},
|
||||
},
|
||||
}
|
||||
for _, tc := range tcs {
|
||||
tc := tc
|
||||
t.Run(tc.desc, func(t *testing.T) {
|
||||
t.Parallel()
|
||||
got := struct {
|
||||
Sources server.SourceConfigs `yaml:"sources"`
|
||||
}{}
|
||||
// Parse contents
|
||||
err := yaml.Unmarshal(testutils.FormatYaml(tc.in), &got)
|
||||
if err != nil {
|
||||
t.Fatalf("unable to unmarshal: %s", err)
|
||||
}
|
||||
if !cmp.Equal(tc.want, got.Sources) {
|
||||
t.Fatalf("incorrect parse: want %v, got %v", tc.want, got.Sources)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestFailParseFromYaml(t *testing.T) {
|
||||
t.Parallel()
|
||||
tcs := []struct {
|
||||
desc string
|
||||
in string
|
||||
err string
|
||||
}{
|
||||
{
|
||||
desc: "extra field",
|
||||
in: `
|
||||
sources:
|
||||
my-spanner-admin-instance:
|
||||
kind: spanner-admin
|
||||
project: test-project
|
||||
`,
|
||||
err: `unable to parse source "my-spanner-admin-instance" as "spanner-admin": [2:1] unknown field "project"
|
||||
1 | kind: spanner-admin
|
||||
> 2 | project: test-project
|
||||
^
|
||||
`,
|
||||
},
|
||||
{
|
||||
desc: "missing required field",
|
||||
in: `
|
||||
sources:
|
||||
my-spanner-admin-instance:
|
||||
useClientOAuth: true
|
||||
`,
|
||||
err: "missing 'kind' field for source \"my-spanner-admin-instance\"",
|
||||
},
|
||||
}
|
||||
for _, tc := range tcs {
|
||||
tc := tc
|
||||
t.Run(tc.desc, func(t *testing.T) {
|
||||
t.Parallel()
|
||||
got := struct {
|
||||
Sources server.SourceConfigs `yaml:"sources"`
|
||||
}{}
|
||||
// Parse contents
|
||||
err := yaml.Unmarshal(testutils.FormatYaml(tc.in), &got)
|
||||
if err == nil {
|
||||
t.Fatalf("expect parsing to fail")
|
||||
}
|
||||
errStr := err.Error()
|
||||
if errStr != tc.err {
|
||||
t.Fatalf("unexpected error: got %q, want %q", errStr, tc.err)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
@@ -1,233 +0,0 @@
|
||||
// Copyright 2026 Google LLC
|
||||
//
|
||||
// Licensed under the Apache License, Version 2.0 (the "License");
|
||||
// you may not use this file except in compliance with the License.
|
||||
// You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software
|
||||
// distributed under the License is distributed on an "AS IS" BASIS,
|
||||
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
// See the License for the specific language governing permissions and
|
||||
// limitations under the License.
|
||||
|
||||
package spannercreateinstance
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
|
||||
instance "cloud.google.com/go/spanner/admin/instance/apiv1"
|
||||
"cloud.google.com/go/spanner/admin/instance/apiv1/instancepb"
|
||||
"github.com/goccy/go-yaml"
|
||||
"github.com/googleapis/genai-toolbox/internal/embeddingmodels"
|
||||
"github.com/googleapis/genai-toolbox/internal/sources"
|
||||
"github.com/googleapis/genai-toolbox/internal/tools"
|
||||
"github.com/googleapis/genai-toolbox/internal/util/parameters"
|
||||
)
|
||||
|
||||
const kind string = "spanner-create-instance"
|
||||
|
||||
func init() {
|
||||
if !tools.Register(kind, newConfig) {
|
||||
panic(fmt.Sprintf("tool kind %q already registered", kind))
|
||||
}
|
||||
}
|
||||
|
||||
func newConfig(ctx context.Context, name string, decoder *yaml.Decoder) (tools.ToolConfig, error) {
|
||||
actual := Config{Name: name}
|
||||
if err := decoder.DecodeContext(ctx, &actual); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return actual, nil
|
||||
}
|
||||
|
||||
type compatibleSource interface {
|
||||
GetDefaultProject() string
|
||||
GetClient(context.Context, string) (*instance.InstanceAdminClient, error)
|
||||
UseClientAuthorization() bool
|
||||
}
|
||||
|
||||
// Config defines the configuration for the create-instance tool.
|
||||
type Config struct {
|
||||
Name string `yaml:"name" validate:"required"`
|
||||
Kind string `yaml:"kind" validate:"required"`
|
||||
Description string `yaml:"description"`
|
||||
Source string `yaml:"source" validate:"required"`
|
||||
AuthRequired []string `yaml:"authRequired"`
|
||||
}
|
||||
|
||||
// validate interface
|
||||
var _ tools.ToolConfig = Config{}
|
||||
|
||||
// ToolConfigKind returns the kind of the tool.
|
||||
func (cfg Config) ToolConfigKind() string {
|
||||
return kind
|
||||
}
|
||||
|
||||
// Initialize initializes the tool from the configuration.
|
||||
func (cfg Config) Initialize(srcs map[string]sources.Source) (tools.Tool, error) {
|
||||
rawS, ok := srcs[cfg.Source]
|
||||
if !ok {
|
||||
return nil, fmt.Errorf("no source named %q configured", cfg.Source)
|
||||
}
|
||||
s, ok := rawS.(compatibleSource)
|
||||
if !ok {
|
||||
return nil, fmt.Errorf("invalid source for %q tool: source %q not compatible", kind, cfg.Source)
|
||||
}
|
||||
|
||||
project := s.GetDefaultProject()
|
||||
var projectParam parameters.Parameter
|
||||
if project != "" {
|
||||
projectParam = parameters.NewStringParameterWithDefault("project", project, "The GCP project ID.")
|
||||
} else {
|
||||
projectParam = parameters.NewStringParameter("project", "The project ID")
|
||||
}
|
||||
|
||||
allParameters := parameters.Parameters{
|
||||
projectParam,
|
||||
parameters.NewStringParameter("instanceId", "The ID of the instance"),
|
||||
parameters.NewStringParameter("displayName", "The display name of the instance"),
|
||||
parameters.NewStringParameter("config", "The instance configuration (e.g., regional-us-central1)"),
|
||||
parameters.NewIntParameter("nodeCount", "The number of nodes, mutually exclusive with processingUnits (one must be 0)"),
|
||||
parameters.NewIntParameter("processingUnits", "The number of processing units, mutually exclusive with nodeCount (one must be 0)"),
|
||||
parameters.NewStringParameter("edition", "The edition of the instance (STANDARD, ENTERPRISE, ENTERPRISE_PLUS)"),
|
||||
}
|
||||
|
||||
paramManifest := allParameters.Manifest()
|
||||
|
||||
description := cfg.Description
|
||||
if description == "" {
|
||||
description = "Creates a Spanner instance."
|
||||
}
|
||||
mcpManifest := tools.GetMcpManifest(cfg.Name, description, cfg.AuthRequired, allParameters, nil)
|
||||
|
||||
return Tool{
|
||||
Config: cfg,
|
||||
AllParams: allParameters,
|
||||
manifest: tools.Manifest{Description: description, Parameters: paramManifest, AuthRequired: cfg.AuthRequired},
|
||||
mcpManifest: mcpManifest,
|
||||
}, nil
|
||||
}
|
||||
|
||||
// Tool represents the create-instance tool.
|
||||
type Tool struct {
|
||||
Config
|
||||
AllParams parameters.Parameters
|
||||
manifest tools.Manifest
|
||||
mcpManifest tools.McpManifest
|
||||
}
|
||||
|
||||
func (t Tool) ToConfig() tools.ToolConfig {
|
||||
return t.Config
|
||||
}
|
||||
|
||||
// Invoke executes the tool's logic.
|
||||
func (t Tool) Invoke(ctx context.Context, resourceMgr tools.SourceProvider, params parameters.ParamValues, accessToken tools.AccessToken) (any, error) {
|
||||
paramsMap := params.AsMap()
|
||||
|
||||
project, _ := paramsMap["project"].(string)
|
||||
instanceId, _ := paramsMap["instanceId"].(string)
|
||||
displayName, _ := paramsMap["displayName"].(string)
|
||||
config, _ := paramsMap["config"].(string)
|
||||
nodeCount, _ := paramsMap["nodeCount"].(int)
|
||||
processingUnits, _ := paramsMap["processingUnits"].(int)
|
||||
editionStr, _ := paramsMap["edition"].(string)
|
||||
|
||||
if (nodeCount > 0 && processingUnits > 0) || (nodeCount == 0 && processingUnits == 0) {
|
||||
return nil, fmt.Errorf("one of nodeCount or processingUnits must be positive, and the other must be 0")
|
||||
}
|
||||
|
||||
source, err := tools.GetCompatibleSource[compatibleSource](resourceMgr, t.Source, t.Name, t.Kind)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
client, err := source.GetClient(ctx, string(accessToken))
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if source.UseClientAuthorization() {
|
||||
defer client.Close()
|
||||
}
|
||||
|
||||
parent := fmt.Sprintf("projects/%s", project)
|
||||
instanceConfig := fmt.Sprintf("projects/%s/instanceConfigs/%s", project, config)
|
||||
|
||||
var edition instancepb.Instance_Edition
|
||||
switch editionStr {
|
||||
case "STANDARD":
|
||||
edition = instancepb.Instance_STANDARD
|
||||
case "ENTERPRISE":
|
||||
edition = instancepb.Instance_ENTERPRISE
|
||||
case "ENTERPRISE_PLUS":
|
||||
edition = instancepb.Instance_ENTERPRISE_PLUS
|
||||
default:
|
||||
edition = instancepb.Instance_EDITION_UNSPECIFIED
|
||||
}
|
||||
|
||||
// Construct the instance object
|
||||
instance := &instancepb.Instance{
|
||||
Config: instanceConfig,
|
||||
DisplayName: displayName,
|
||||
Edition: edition,
|
||||
NodeCount: int32(nodeCount),
|
||||
ProcessingUnits: int32(processingUnits),
|
||||
}
|
||||
|
||||
req := &instancepb.CreateInstanceRequest{
|
||||
Parent: parent,
|
||||
InstanceId: instanceId,
|
||||
Instance: instance,
|
||||
}
|
||||
|
||||
op, err := client.CreateInstance(ctx, req)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("failed to create instance: %w", err)
|
||||
}
|
||||
|
||||
// Wait for the operation to complete
|
||||
resp, err := op.Wait(ctx)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("failed to wait for create instance operation: %w", err)
|
||||
}
|
||||
|
||||
return resp, nil
|
||||
}
|
||||
|
||||
// ParseParams parses the parameters for the tool.
|
||||
func (t Tool) ParseParams(data map[string]any, claims map[string]map[string]any) (parameters.ParamValues, error) {
|
||||
return parameters.ParseParams(t.AllParams, data, claims)
|
||||
}
|
||||
|
||||
func (t Tool) EmbedParams(ctx context.Context, paramValues parameters.ParamValues, embeddingModelsMap map[string]embeddingmodels.EmbeddingModel) (parameters.ParamValues, error) {
|
||||
return parameters.EmbedParams(ctx, t.AllParams, paramValues, embeddingModelsMap, nil)
|
||||
}
|
||||
|
||||
// Manifest returns the tool's manifest.
|
||||
func (t Tool) Manifest() tools.Manifest {
|
||||
return t.manifest
|
||||
}
|
||||
|
||||
// McpManifest returns the tool's MCP manifest.
|
||||
func (t Tool) McpManifest() tools.McpManifest {
|
||||
return t.mcpManifest
|
||||
}
|
||||
|
||||
// Authorized checks if the tool is authorized.
|
||||
func (t Tool) Authorized(verifiedAuthServices []string) bool {
|
||||
return true
|
||||
}
|
||||
|
||||
func (t Tool) RequiresClientAuthorization(resourceMgr tools.SourceProvider) (bool, error) {
|
||||
source, err := tools.GetCompatibleSource[compatibleSource](resourceMgr, t.Source, t.Name, t.Kind)
|
||||
if err != nil {
|
||||
return false, err
|
||||
}
|
||||
return source.UseClientAuthorization(), nil
|
||||
}
|
||||
|
||||
func (t Tool) GetAuthTokenHeaderName(resourceMgr tools.SourceProvider) (string, error) {
|
||||
return "Authorization", nil
|
||||
}
|
||||
@@ -1,110 +0,0 @@
|
||||
// Copyright 2026 Google LLC
|
||||
//
|
||||
// Licensed under the Apache License, Version 2.0 (the "License");
|
||||
// you may not use this file except in compliance with the License.
|
||||
// You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software
|
||||
// distributed under the License is distributed on an "AS IS" BASIS,
|
||||
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
// See the License for the specific language governing permissions and
|
||||
// limitations under the License.
|
||||
|
||||
package spannercreateinstance_test
|
||||
|
||||
import (
|
||||
"context"
|
||||
"testing"
|
||||
|
||||
yaml "github.com/goccy/go-yaml"
|
||||
"github.com/google/go-cmp/cmp"
|
||||
"github.com/googleapis/genai-toolbox/internal/server"
|
||||
"github.com/googleapis/genai-toolbox/internal/testutils"
|
||||
"github.com/googleapis/genai-toolbox/internal/tools/spanneradmin/spannercreateinstance"
|
||||
"github.com/googleapis/genai-toolbox/internal/util/parameters"
|
||||
)
|
||||
|
||||
func TestParseFromYaml(t *testing.T) {
|
||||
ctx, err := testutils.ContextWithNewLogger()
|
||||
if err != nil {
|
||||
t.Fatalf("unexpected error: %s", err)
|
||||
}
|
||||
tcs := []struct {
|
||||
desc string
|
||||
in string
|
||||
want server.ToolConfigs
|
||||
}{
|
||||
{
|
||||
desc: "basic example",
|
||||
in: `
|
||||
tools:
|
||||
create-instance-tool:
|
||||
kind: spanner-create-instance
|
||||
description: a test description
|
||||
source: a-source
|
||||
`,
|
||||
want: server.ToolConfigs{
|
||||
"create-instance-tool": spannercreateinstance.Config{
|
||||
Name: "create-instance-tool",
|
||||
Kind: "spanner-create-instance",
|
||||
Description: "a test description",
|
||||
Source: "a-source",
|
||||
AuthRequired: []string{},
|
||||
},
|
||||
},
|
||||
},
|
||||
}
|
||||
for _, tc := range tcs {
|
||||
t.Run(tc.desc, func(t *testing.T) {
|
||||
got := struct {
|
||||
Tools server.ToolConfigs `yaml:"tools"`
|
||||
}{}
|
||||
// Parse contents
|
||||
err := yaml.UnmarshalContext(ctx, testutils.FormatYaml(tc.in), &got)
|
||||
if err != nil {
|
||||
t.Fatalf("unable to unmarshal: %s", err)
|
||||
}
|
||||
if diff := cmp.Diff(tc.want, got.Tools); diff != "" {
|
||||
t.Fatalf("incorrect parse: diff %v", diff)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestInvokeNodeCountAndProcessingUnitsValidation(t *testing.T) {
|
||||
t.Parallel()
|
||||
testCases := []struct {
|
||||
name string
|
||||
params parameters.ParamValues
|
||||
wantErr string
|
||||
}{
|
||||
{
|
||||
name: "Both positive",
|
||||
params: parameters.ParamValues{
|
||||
{Name: "nodeCount", Value: 1},
|
||||
{Name: "processingUnits", Value: 1000},
|
||||
},
|
||||
wantErr: "one of nodeCount or processingUnits must be positive, and the other must be 0",
|
||||
},
|
||||
{
|
||||
name: "Both zero",
|
||||
params: parameters.ParamValues{
|
||||
{Name: "nodeCount", Value: 0},
|
||||
{Name: "processingUnits", Value: 0},
|
||||
},
|
||||
wantErr: "one of nodeCount or processingUnits must be positive, and the other must be 0",
|
||||
},
|
||||
}
|
||||
|
||||
for _, tc := range testCases {
|
||||
t.Run(tc.name, func(t *testing.T) {
|
||||
tool := spannercreateinstance.Tool{}
|
||||
_, err := tool.Invoke(context.Background(), nil, tc.params, "")
|
||||
if err == nil || err.Error() != tc.wantErr {
|
||||
t.Errorf("Invoke() error = %v, wantErr %v", err, tc.wantErr)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
@@ -1,178 +0,0 @@
|
||||
// Copyright 2026 Google LLC
|
||||
//
|
||||
// Licensed under the Apache License, Version 2.0 (the "License");
|
||||
// you may not use this file except in compliance with the License.
|
||||
// You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software
|
||||
// distributed under the License is distributed on an "AS IS" BASIS,
|
||||
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
// See the License for the specific language governing permissions and
|
||||
// limitations under the License.
|
||||
|
||||
package spanneradmin
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"io"
|
||||
"net/http"
|
||||
"os"
|
||||
"regexp"
|
||||
"strings"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
instance "cloud.google.com/go/spanner/admin/instance/apiv1"
|
||||
"cloud.google.com/go/spanner/admin/instance/apiv1/instancepb"
|
||||
"github.com/google/uuid"
|
||||
"github.com/googleapis/genai-toolbox/internal/testutils"
|
||||
"github.com/googleapis/genai-toolbox/tests"
|
||||
)
|
||||
|
||||
var (
|
||||
SpannerProject = os.Getenv("SPANNER_PROJECT")
|
||||
)
|
||||
|
||||
func getSpannerAdminVars(t *testing.T) map[string]any {
|
||||
if SpannerProject == "" {
|
||||
t.Fatal("'SPANNER_PROJECT' not set")
|
||||
}
|
||||
|
||||
return map[string]any{
|
||||
"kind": "spanner-admin",
|
||||
"defaultProject": SpannerProject,
|
||||
}
|
||||
}
|
||||
|
||||
func TestSpannerAdminCreateInstance(t *testing.T) {
|
||||
sourceConfig := getSpannerAdminVars(t)
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 20*time.Minute)
|
||||
defer cancel()
|
||||
|
||||
shortUuid := strings.ReplaceAll(uuid.New().String(), "-", "")[:10]
|
||||
instanceId := "test-inst-" + shortUuid
|
||||
|
||||
displayName := "Test Instance " + shortUuid
|
||||
instanceConfig := "regional-us-central1"
|
||||
nodeCount := 1
|
||||
edition := "ENTERPRISE"
|
||||
|
||||
// Setup Admin Client for verification and cleanup
|
||||
adminClient, err := instance.NewInstanceAdminClient(ctx)
|
||||
if err != nil {
|
||||
t.Fatalf("unable to create Spanner instance admin client: %s", err)
|
||||
}
|
||||
defer adminClient.Close()
|
||||
|
||||
// Teardown function
|
||||
defer func() {
|
||||
err := adminClient.DeleteInstance(ctx, &instancepb.DeleteInstanceRequest{
|
||||
Name: fmt.Sprintf("projects/%s/instances/%s", SpannerProject, instanceId),
|
||||
})
|
||||
if err != nil {
|
||||
// If it fails, it might not have been created, log it but don't fail if it's "not found"
|
||||
t.Logf("cleanup: failed to delete instance %s: %s", instanceId, err)
|
||||
} else {
|
||||
t.Logf("cleanup: deleted instance %s", instanceId)
|
||||
}
|
||||
}()
|
||||
|
||||
// Construct Tools Config
|
||||
|
||||
toolsConfig := map[string]any{
|
||||
"sources": map[string]any{
|
||||
"my-spanner-admin": sourceConfig,
|
||||
},
|
||||
"tools": map[string]any{
|
||||
"create-instance-tool": map[string]any{
|
||||
"kind": "spanner-create-instance",
|
||||
"source": "my-spanner-admin",
|
||||
"description": "Creates a Spanner instance.",
|
||||
},
|
||||
},
|
||||
}
|
||||
|
||||
// Start Toolbox Server
|
||||
cmd, cleanup, err := tests.StartCmd(ctx, toolsConfig)
|
||||
if err != nil {
|
||||
t.Fatalf("command initialization returned an error: %s", err)
|
||||
}
|
||||
defer cleanup()
|
||||
|
||||
waitCtx, cancelWait := context.WithTimeout(ctx, 10*time.Second)
|
||||
defer cancelWait()
|
||||
out, err := testutils.WaitForString(waitCtx, regexp.MustCompile(`Server ready to serve`), cmd.Out)
|
||||
if err != nil {
|
||||
t.Logf("toolbox command logs: \n%s", out)
|
||||
t.Fatalf("toolbox didn't start successfully: %s", err)
|
||||
}
|
||||
|
||||
// Prepare Invocation Payload
|
||||
|
||||
payload := map[string]any{
|
||||
"project": SpannerProject,
|
||||
"instanceId": instanceId,
|
||||
"displayName": displayName,
|
||||
"config": instanceConfig,
|
||||
"nodeCount": nodeCount,
|
||||
"edition": edition,
|
||||
"processingUnits": 0,
|
||||
}
|
||||
|
||||
payloadBytes, err := json.Marshal(payload)
|
||||
if err != nil {
|
||||
t.Fatalf("failed to marshal payload: %s", err)
|
||||
}
|
||||
|
||||
// Invoke Tool
|
||||
invokeUrl := "http://127.0.0.1:5000/api/tool/create-instance-tool/invoke"
|
||||
req, err := http.NewRequest(http.MethodPost, invokeUrl, bytes.NewBuffer(payloadBytes))
|
||||
if err != nil {
|
||||
t.Fatalf("unable to create request: %s", err)
|
||||
}
|
||||
req.Header.Add("Content-type", "application/json")
|
||||
|
||||
t.Logf("Invoking create-instance-tool for instance: %s", instanceId)
|
||||
resp, err := http.DefaultClient.Do(req)
|
||||
if err != nil {
|
||||
t.Fatalf("unable to send request: %s", err)
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
|
||||
if resp.StatusCode != http.StatusOK {
|
||||
bodyBytes, _ := io.ReadAll(resp.Body)
|
||||
t.Fatalf("response status code is not 200, got %d: %s", resp.StatusCode, string(bodyBytes))
|
||||
}
|
||||
|
||||
// Check Response
|
||||
var body map[string]interface{}
|
||||
err = json.NewDecoder(resp.Body).Decode(&body)
|
||||
if err != nil {
|
||||
t.Fatalf("error parsing response body")
|
||||
}
|
||||
|
||||
// Verify Instance Exists via Admin Client
|
||||
t.Logf("Verifying instance %s exists...", instanceId)
|
||||
instanceName := fmt.Sprintf("projects/%s/instances/%s", SpannerProject, instanceId)
|
||||
gotInstance, err := adminClient.GetInstance(ctx, &instancepb.GetInstanceRequest{
|
||||
Name: instanceName,
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("failed to get instance from admin client: %s", err)
|
||||
}
|
||||
|
||||
if gotInstance.Name != instanceName {
|
||||
t.Errorf("expected instance name %s, got %s", instanceName, gotInstance.Name)
|
||||
}
|
||||
if gotInstance.DisplayName != displayName {
|
||||
t.Errorf("expected display name %s, got %s", displayName, gotInstance.DisplayName)
|
||||
}
|
||||
if gotInstance.NodeCount != int32(nodeCount) {
|
||||
t.Errorf("expected node count %d, got %d", nodeCount, gotInstance.NodeCount)
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user