mirror of
https://github.com/googleapis/genai-toolbox.git
synced 2026-01-29 01:08:01 -05:00
Compare commits
23 Commits
processing
...
spanner-cr
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
caa33d41b2 | ||
|
|
668ca5da8a | ||
|
|
a7d0f19716 | ||
|
|
de7c65eb1c | ||
|
|
589059ed3f | ||
|
|
6d43488ddf | ||
|
|
3a317a3455 | ||
|
|
af62ddf9c1 | ||
|
|
1475b4d092 | ||
|
|
c67b0cf0fc | ||
|
|
511c4651f3 | ||
|
|
62095feadb | ||
|
|
f9f34b1005 | ||
|
|
4170fe309f | ||
|
|
a9edd5e32d | ||
|
|
09c979d8db | ||
|
|
7d79f4909a | ||
|
|
330dd843bf | ||
|
|
e831f71421 | ||
|
|
4f62db499d | ||
|
|
e6de2170cf | ||
|
|
5b7f5a3039 | ||
|
|
302b564ca1 |
@@ -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"]
|
||||
|
||||
@@ -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"
|
||||
|
||||
59
docs/SPANNERADMIN_README.md
Normal file
59
docs/SPANNERADMIN_README.md
Normal file
@@ -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).
|
||||
43
docs/en/resources/sources/spanner-admin.md
Normal file
43
docs/en/resources/sources/spanner-admin.md
Normal file
@@ -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`. |
|
||||
@@ -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. |
|
||||
@@ -1,52 +0,0 @@
|
||||
---
|
||||
title: "Pre and Post processing"
|
||||
type: docs
|
||||
weight: 1
|
||||
description: >
|
||||
Pre and Post processing in GenAI applications.
|
||||
---
|
||||
|
||||
Pre and post processing allow developers to intercept and modify interactions between the agent and its tools or the user. This capability is essential for building robust, secure, and compliant agents.
|
||||
|
||||
## Types of Processing
|
||||
|
||||
### Pre-processing
|
||||
|
||||
Pre-processing occurs before a tool is executed or an agent processes a message. Key types include:
|
||||
|
||||
- **Input Sanitization & Redaction**: Detecting and masking sensitive information (like PII) in user queries or tool arguments to prevent it from being logged or sent to unauthorized systems.
|
||||
- **Business Logic Validation**: Verifying that the proposed action complies with business rules (e.g., ensuring a requested hotel stay does not exceed 14 days, or checking if a user has sufficient permission).
|
||||
- **Security Guardrails**: Analyzing inputs for potential prompt injection attacks or malicious payloads.
|
||||
|
||||
### Post-processing
|
||||
|
||||
Post-processing occurs after a tool has executed or the model has generated a response. Key types include:
|
||||
|
||||
- **Response Enrichment**: Injecting additional data into the tool output that wasn't part of the raw API response (e.g., calculating loyalty points earned based on the booking value).
|
||||
- **Output Formatting**: Transforming raw data (like JSON or XML) into a more human-readable or model-friendly format to improve the agent's understanding.
|
||||
- **Compliance Auditing**: Logging the final outcome of transactions, including the original request and the result, to a secure audit trail.
|
||||
|
||||
## Processing Scopes
|
||||
|
||||
Processing logic can be applied at different levels of the application:
|
||||
|
||||
### Tool Level
|
||||
|
||||
Wraps individual tool executions. This is best for logic specific to a single tool or a set of tools.
|
||||
|
||||
- **Scope**: Intercepts the raw inputs (arguments) to a tool and its outputs.
|
||||
- **Use Cases**: Argument validation, output formatting, specific privacy rules for sensitive tools.
|
||||
|
||||
### Model Level
|
||||
|
||||
Intercepts individual calls to the Large Language Model (LLM).
|
||||
|
||||
- **Scope**: Intercepts the list of messages (prompt) sent to the model and the generation (response) received.
|
||||
- **Use Cases**: Global PII redaction (across all tools/chat), prompt engineering/injection, token usage tracking, and hallucination detection.
|
||||
|
||||
### Agent Level
|
||||
|
||||
Wraps the high-level agent execution loop (e.g., a "turn" in the conversation).
|
||||
|
||||
- **Scope**: Intercepts the initial user input and the final agent response, enveloping one or more model calls and tool executions.
|
||||
- **Use Cases**: User authentication, rate limiting, session management, and end-to-end audit logging.
|
||||
@@ -1,5 +0,0 @@
|
||||
Final Client Response:
|
||||
AI:
|
||||
Booking Confirmed!
|
||||
Loyalty Points
|
||||
POLICY CHECK: Intercepting 'book-hotel'
|
||||
@@ -1,31 +0,0 @@
|
||||
---
|
||||
title: "(Python) Pre and post processing"
|
||||
type: docs
|
||||
weight: 4
|
||||
description: >
|
||||
How to add pre and post processing to your Python toolbox applications.
|
||||
---
|
||||
|
||||
## Prerequisites
|
||||
|
||||
This tutorial assumes that you have set up a basic toolbox application as described in the [local quickstart](../../getting-started/local_quickstart).
|
||||
|
||||
This guide demonstrates how to implement these patterns in your Toolbox applications.
|
||||
|
||||
## Python
|
||||
|
||||
{{< tabpane persist=header >}}
|
||||
{{% tab header="ADK" text=true %}}
|
||||
Coming soon.
|
||||
{{% /tab %}}
|
||||
{{% tab header="Langchain" text=true %}}
|
||||
The following example demonstrates how to use `ToolboxClient` with LangChain's middleware to implement pre and post processing for tool calls.
|
||||
|
||||
```py
|
||||
{{< include "python/langchain/agent.py" >}}
|
||||
```
|
||||
|
||||
For more information, see the [LangChain Middleware documentation](https://docs.langchain.com/oss/python/langchain/middleware/custom#wrap-style-hooks).
|
||||
You can also add model-level (`wrap_model`) and agent-level (`before_agent`, `after_agent`) hooks to intercept messages at different stages of the execution loop. See the [LangChain Middleware documentation](https://docs.langchain.com/oss/python/langchain/middleware/custom#wrap-style-hooks) for details on these additional hook types.
|
||||
{{% /tab %}}
|
||||
{{< /tabpane >}}
|
||||
@@ -1,4 +0,0 @@
|
||||
# This file makes the 'pre_post_processing/python' directory a Python package.
|
||||
|
||||
# You can include any package-level initialization logic here if needed.
|
||||
# For now, this file is empty.
|
||||
@@ -1,58 +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.
|
||||
|
||||
import asyncio
|
||||
import importlib
|
||||
import os
|
||||
from pathlib import Path
|
||||
|
||||
import pytest
|
||||
|
||||
ORCH_NAME = os.environ.get("ORCH_NAME")
|
||||
module_path = f"python.{ORCH_NAME}.agent"
|
||||
agent = importlib.import_module(module_path)
|
||||
|
||||
|
||||
@pytest.fixture(scope="module")
|
||||
def golden_keywords():
|
||||
"""Loads expected keywords from the golden.txt file."""
|
||||
golden_file_path = Path(__file__).resolve().parent.parent / "golden.txt"
|
||||
if not golden_file_path.exists():
|
||||
pytest.fail(f"Golden file not found: {golden_file_path}")
|
||||
try:
|
||||
with open(golden_file_path, "r") as f:
|
||||
return [line.strip() for line in f.readlines() if line.strip()]
|
||||
except Exception as e:
|
||||
pytest.fail(f"Could not read golden.txt: {e}")
|
||||
|
||||
|
||||
# --- Execution Tests ---
|
||||
class TestExecution:
|
||||
"""Test framework execution and output validation."""
|
||||
|
||||
@pytest.fixture(scope="function")
|
||||
def script_output(self, capsys):
|
||||
"""Run the agent function and return its output."""
|
||||
asyncio.run(agent.main())
|
||||
return capsys.readouterr()
|
||||
|
||||
def test_script_runs_without_errors(self, script_output):
|
||||
"""Test that the script runs and produces no stderr."""
|
||||
assert script_output.err == "", f"Script produced stderr: {script_output.err}"
|
||||
|
||||
def test_keywords_in_output(self, script_output, golden_keywords):
|
||||
"""Test that expected keywords are present in the script's output."""
|
||||
output = script_output.out
|
||||
missing_keywords = [kw for kw in golden_keywords if kw not in output]
|
||||
assert not missing_keywords, f"Missing keywords in output: {missing_keywords}"
|
||||
@@ -1,111 +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.
|
||||
|
||||
import asyncio
|
||||
from datetime import datetime
|
||||
|
||||
from langchain.agents import create_agent
|
||||
from langchain.agents.middleware import wrap_tool_call
|
||||
from langchain_core.messages import ToolMessage
|
||||
from langchain_google_vertexai import ChatVertexAI
|
||||
from toolbox_langchain import ToolboxClient
|
||||
|
||||
system_prompt = """
|
||||
You're a helpful hotel assistant. You handle hotel searching, booking and
|
||||
cancellations. When the user searches for a hotel, mention it's name, id,
|
||||
location and price tier. Always mention hotel ids while performing any
|
||||
searches. This is very important for any operations. For any bookings or
|
||||
cancellations, please provide the appropriate confirmation. Be sure to
|
||||
update checkin or checkout dates if mentioned by the user.
|
||||
Don't ask for confirmations from the user.
|
||||
"""
|
||||
|
||||
|
||||
# Pre processing
|
||||
@wrap_tool_call
|
||||
async def enforce_business_rules(request, handler):
|
||||
"""
|
||||
Business Logic Validation:
|
||||
Enforces max stay duration (e.g., max 14 days).
|
||||
"""
|
||||
tool_call = request.tool_call
|
||||
name = tool_call["name"]
|
||||
args = tool_call["args"]
|
||||
|
||||
print(f"POLICY CHECK: Intercepting '{name}'")
|
||||
|
||||
if name == "update-hotel":
|
||||
if "checkin_date" in args and "checkout_date" in args:
|
||||
try:
|
||||
start = datetime.fromisoformat(args["checkin_date"])
|
||||
end = datetime.fromisoformat(args["checkout_date"])
|
||||
duration = (end - start).days
|
||||
|
||||
if duration > 14:
|
||||
print("BLOCKED: Stay too long")
|
||||
return ToolMessage(
|
||||
content="Error: Maximum stay duration is 14 days.",
|
||||
tool_call_id=tool_call["id"],
|
||||
)
|
||||
except ValueError:
|
||||
pass # Ignore invalid date formats
|
||||
|
||||
return await handler(request)
|
||||
|
||||
|
||||
# Post processing
|
||||
@wrap_tool_call
|
||||
async def enrich_response(request, handler):
|
||||
"""
|
||||
Post-Processing & Enrichment:
|
||||
Adds loyalty points information to successful bookings.
|
||||
Standardizes output format.
|
||||
"""
|
||||
result = await handler(request)
|
||||
|
||||
if isinstance(result, ToolMessage):
|
||||
content = str(result.content)
|
||||
tool_name = request.tool_call["name"]
|
||||
|
||||
if tool_name == "book-hotel" and "Error" not in content:
|
||||
loyalty_bonus = 500
|
||||
result.content = f"Booking Confirmed! \n You earned {loyalty_bonus} Loyalty Points with this stay.\n\nSystem Details: {content}"
|
||||
|
||||
return result
|
||||
|
||||
|
||||
async def main():
|
||||
async with ToolboxClient("http://127.0.0.1:5000") as client:
|
||||
tools = await client.aload_toolset("my-toolset")
|
||||
model = ChatVertexAI(model="gemini-2.5-flash")
|
||||
agent = create_agent(
|
||||
system_prompt=system_prompt,
|
||||
model=model,
|
||||
tools=tools,
|
||||
middleware=[enforce_business_rules, enrich_response],
|
||||
)
|
||||
|
||||
user_input = "Book hotel with id 3."
|
||||
response = await agent.ainvoke(
|
||||
{"messages": [{"role": "user", "content": user_input}]}
|
||||
)
|
||||
|
||||
print("-" * 50)
|
||||
print("Final Client Response:")
|
||||
last_ai_msg = response["messages"][-1].content
|
||||
print(f"AI: {last_ai_msg}")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
asyncio.run(main())
|
||||
@@ -1,2 +0,0 @@
|
||||
langchain==1.2.6
|
||||
toolbox-langchain==0.5.7
|
||||
@@ -50,6 +50,7 @@ var expectedToolSources = []string{
|
||||
"serverless-spark",
|
||||
"singlestore",
|
||||
"snowflake",
|
||||
"spanner-admin",
|
||||
"spanner-postgres",
|
||||
"spanner",
|
||||
"sqlite",
|
||||
|
||||
27
internal/prebuiltconfigs/tools/spanner-admin.yaml
Normal file
27
internal/prebuiltconfigs/tools/spanner-admin.yaml
Normal file
@@ -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
|
||||
120
internal/sources/spanneradmin/spanneradmin.go
Normal file
120
internal/sources/spanneradmin/spanneradmin.go
Normal file
@@ -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
|
||||
}
|
||||
122
internal/sources/spanneradmin/spanneradmin_test.go
Normal file
122
internal/sources/spanneradmin/spanneradmin_test.go
Normal file
@@ -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)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
@@ -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
|
||||
}
|
||||
@@ -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)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
178
tests/spanneradmin/spanneradmin_integration_test.go
Normal file
178
tests/spanneradmin/spanneradmin_integration_test.go
Normal file
@@ -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)
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user