mirror of
https://github.com/googleapis/genai-toolbox.git
synced 2026-01-21 05:18:14 -05:00
Compare commits
6 Commits
config-pre
...
dgraph-doc
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
1a989bab10 | ||
|
|
00c3e6d8cb | ||
|
|
d00b6fdf18 | ||
|
|
4d23a3bbf2 | ||
|
|
5e0999ebf5 | ||
|
|
6b02591703 |
@@ -98,6 +98,7 @@ import (
|
|||||||
_ "github.com/googleapis/genai-toolbox/internal/tools/cloudsql/cloudsqlgetinstances"
|
_ "github.com/googleapis/genai-toolbox/internal/tools/cloudsql/cloudsqlgetinstances"
|
||||||
_ "github.com/googleapis/genai-toolbox/internal/tools/cloudsql/cloudsqllistdatabases"
|
_ "github.com/googleapis/genai-toolbox/internal/tools/cloudsql/cloudsqllistdatabases"
|
||||||
_ "github.com/googleapis/genai-toolbox/internal/tools/cloudsql/cloudsqllistinstances"
|
_ "github.com/googleapis/genai-toolbox/internal/tools/cloudsql/cloudsqllistinstances"
|
||||||
|
_ "github.com/googleapis/genai-toolbox/internal/tools/cloudsql/cloudsqlrestorebackup"
|
||||||
_ "github.com/googleapis/genai-toolbox/internal/tools/cloudsql/cloudsqlwaitforoperation"
|
_ "github.com/googleapis/genai-toolbox/internal/tools/cloudsql/cloudsqlwaitforoperation"
|
||||||
_ "github.com/googleapis/genai-toolbox/internal/tools/cloudsqlmssql/cloudsqlmssqlcreateinstance"
|
_ "github.com/googleapis/genai-toolbox/internal/tools/cloudsqlmssql/cloudsqlmssqlcreateinstance"
|
||||||
_ "github.com/googleapis/genai-toolbox/internal/tools/cloudsqlmysql/cloudsqlmysqlcreateinstance"
|
_ "github.com/googleapis/genai-toolbox/internal/tools/cloudsqlmysql/cloudsqlmysqlcreateinstance"
|
||||||
|
|||||||
@@ -1493,7 +1493,7 @@ func TestPrebuiltTools(t *testing.T) {
|
|||||||
wantToolset: server.ToolsetConfigs{
|
wantToolset: server.ToolsetConfigs{
|
||||||
"cloud_sql_postgres_admin_tools": tools.ToolsetConfig{
|
"cloud_sql_postgres_admin_tools": tools.ToolsetConfig{
|
||||||
Name: "cloud_sql_postgres_admin_tools",
|
Name: "cloud_sql_postgres_admin_tools",
|
||||||
ToolNames: []string{"create_instance", "get_instance", "list_instances", "create_database", "list_databases", "create_user", "wait_for_operation", "postgres_upgrade_precheck", "clone_instance", "create_backup"},
|
ToolNames: []string{"create_instance", "get_instance", "list_instances", "create_database", "list_databases", "create_user", "wait_for_operation", "postgres_upgrade_precheck", "clone_instance", "create_backup", "restore_backup"},
|
||||||
},
|
},
|
||||||
},
|
},
|
||||||
},
|
},
|
||||||
@@ -1503,7 +1503,7 @@ func TestPrebuiltTools(t *testing.T) {
|
|||||||
wantToolset: server.ToolsetConfigs{
|
wantToolset: server.ToolsetConfigs{
|
||||||
"cloud_sql_mysql_admin_tools": tools.ToolsetConfig{
|
"cloud_sql_mysql_admin_tools": tools.ToolsetConfig{
|
||||||
Name: "cloud_sql_mysql_admin_tools",
|
Name: "cloud_sql_mysql_admin_tools",
|
||||||
ToolNames: []string{"create_instance", "get_instance", "list_instances", "create_database", "list_databases", "create_user", "wait_for_operation", "clone_instance", "create_backup"},
|
ToolNames: []string{"create_instance", "get_instance", "list_instances", "create_database", "list_databases", "create_user", "wait_for_operation", "clone_instance", "create_backup", "restore_backup"},
|
||||||
},
|
},
|
||||||
},
|
},
|
||||||
},
|
},
|
||||||
@@ -1513,7 +1513,7 @@ func TestPrebuiltTools(t *testing.T) {
|
|||||||
wantToolset: server.ToolsetConfigs{
|
wantToolset: server.ToolsetConfigs{
|
||||||
"cloud_sql_mssql_admin_tools": tools.ToolsetConfig{
|
"cloud_sql_mssql_admin_tools": tools.ToolsetConfig{
|
||||||
Name: "cloud_sql_mssql_admin_tools",
|
Name: "cloud_sql_mssql_admin_tools",
|
||||||
ToolNames: []string{"create_instance", "get_instance", "list_instances", "create_database", "list_databases", "create_user", "wait_for_operation", "clone_instance", "create_backup"},
|
ToolNames: []string{"create_instance", "get_instance", "list_instances", "create_database", "list_databases", "create_user", "wait_for_operation", "clone_instance", "create_backup", "restore_backup"},
|
||||||
},
|
},
|
||||||
},
|
},
|
||||||
},
|
},
|
||||||
|
|||||||
@@ -54,6 +54,7 @@ instance, database and users:
|
|||||||
* `create_instance`
|
* `create_instance`
|
||||||
* `create_user`
|
* `create_user`
|
||||||
* `clone_instance`
|
* `clone_instance`
|
||||||
|
* `restore_backup`
|
||||||
|
|
||||||
## Install MCP Toolbox
|
## Install MCP Toolbox
|
||||||
|
|
||||||
@@ -301,6 +302,7 @@ instances and interacting with your database:
|
|||||||
* **wait_for_operation**: Waits for a Cloud SQL operation to complete.
|
* **wait_for_operation**: Waits for a Cloud SQL operation to complete.
|
||||||
* **clone_instance**: Creates a clone of an existing Cloud SQL for SQL Server instance.
|
* **clone_instance**: Creates a clone of an existing Cloud SQL for SQL Server instance.
|
||||||
* **create_backup**: Creates a backup on a Cloud SQL instance.
|
* **create_backup**: Creates a backup on a Cloud SQL instance.
|
||||||
|
* **restore_backup**: Restores a backup of a Cloud SQL instance.
|
||||||
|
|
||||||
{{< notice note >}}
|
{{< notice note >}}
|
||||||
Prebuilt tools are pre-1.0, so expect some tool changes between versions. LLMs
|
Prebuilt tools are pre-1.0, so expect some tool changes between versions. LLMs
|
||||||
|
|||||||
@@ -54,6 +54,7 @@ database and users:
|
|||||||
* `create_instance`
|
* `create_instance`
|
||||||
* `create_user`
|
* `create_user`
|
||||||
* `clone_instance`
|
* `clone_instance`
|
||||||
|
* `restore_backup`
|
||||||
|
|
||||||
## Install MCP Toolbox
|
## Install MCP Toolbox
|
||||||
|
|
||||||
@@ -301,6 +302,7 @@ instances and interacting with your database:
|
|||||||
* **wait_for_operation**: Waits for a Cloud SQL operation to complete.
|
* **wait_for_operation**: Waits for a Cloud SQL operation to complete.
|
||||||
* **clone_instance**: Creates a clone of an existing Cloud SQL for MySQL instance.
|
* **clone_instance**: Creates a clone of an existing Cloud SQL for MySQL instance.
|
||||||
* **create_backup**: Creates a backup on a Cloud SQL instance.
|
* **create_backup**: Creates a backup on a Cloud SQL instance.
|
||||||
|
* **restore_backup**: Restores a backup of a Cloud SQL instance.
|
||||||
|
|
||||||
{{< notice note >}}
|
{{< notice note >}}
|
||||||
Prebuilt tools are pre-1.0, so expect some tool changes between versions. LLMs
|
Prebuilt tools are pre-1.0, so expect some tool changes between versions. LLMs
|
||||||
|
|||||||
@@ -54,6 +54,7 @@ instance, database and users:
|
|||||||
* `create_instance`
|
* `create_instance`
|
||||||
* `create_user`
|
* `create_user`
|
||||||
* `clone_instance`
|
* `clone_instance`
|
||||||
|
* `restore_backup`
|
||||||
|
|
||||||
## Install MCP Toolbox
|
## Install MCP Toolbox
|
||||||
|
|
||||||
@@ -301,6 +302,7 @@ instances and interacting with your database:
|
|||||||
* **wait_for_operation**: Waits for a Cloud SQL operation to complete.
|
* **wait_for_operation**: Waits for a Cloud SQL operation to complete.
|
||||||
* **clone_instance**: Creates a clone of an existing Cloud SQL for PostgreSQL instance.
|
* **clone_instance**: Creates a clone of an existing Cloud SQL for PostgreSQL instance.
|
||||||
* **create_backup**: Creates a backup on a Cloud SQL instance.
|
* **create_backup**: Creates a backup on a Cloud SQL instance.
|
||||||
|
* **restore_backup**: Restores a backup of a Cloud SQL instance.
|
||||||
|
|
||||||
{{< notice note >}}
|
{{< notice note >}}
|
||||||
Prebuilt tools are pre-1.0, so expect some tool changes between versions. LLMs
|
Prebuilt tools are pre-1.0, so expect some tool changes between versions. LLMs
|
||||||
|
|||||||
@@ -20,6 +20,7 @@ The native SDKs can be combined with MCP clients in many cases.
|
|||||||
|
|
||||||
Toolbox currently supports the following versions of MCP specification:
|
Toolbox currently supports the following versions of MCP specification:
|
||||||
|
|
||||||
|
* [2025-11-25](https://modelcontextprotocol.io/specification/2025-11-25)
|
||||||
* [2025-06-18](https://modelcontextprotocol.io/specification/2025-06-18)
|
* [2025-06-18](https://modelcontextprotocol.io/specification/2025-06-18)
|
||||||
* [2025-03-26](https://modelcontextprotocol.io/specification/2025-03-26)
|
* [2025-03-26](https://modelcontextprotocol.io/specification/2025-03-26)
|
||||||
* [2024-11-05](https://modelcontextprotocol.io/specification/2024-11-05)
|
* [2024-11-05](https://modelcontextprotocol.io/specification/2024-11-05)
|
||||||
|
|||||||
@@ -194,6 +194,7 @@ See [Usage Examples](../reference/cli.md#examples).
|
|||||||
* `create_instance`
|
* `create_instance`
|
||||||
* `create_user`
|
* `create_user`
|
||||||
* `clone_instance`
|
* `clone_instance`
|
||||||
|
* `restore_backup`
|
||||||
|
|
||||||
* **Tools:**
|
* **Tools:**
|
||||||
* `create_instance`: Creates a new Cloud SQL for MySQL instance.
|
* `create_instance`: Creates a new Cloud SQL for MySQL instance.
|
||||||
@@ -205,6 +206,7 @@ See [Usage Examples](../reference/cli.md#examples).
|
|||||||
* `wait_for_operation`: Waits for a Cloud SQL operation to complete.
|
* `wait_for_operation`: Waits for a Cloud SQL operation to complete.
|
||||||
* `clone_instance`: Creates a clone for an existing Cloud SQL for MySQL instance.
|
* `clone_instance`: Creates a clone for an existing Cloud SQL for MySQL instance.
|
||||||
* `create_backup`: Creates a backup on a Cloud SQL instance.
|
* `create_backup`: Creates a backup on a Cloud SQL instance.
|
||||||
|
* `restore_backup`: Restores a backup of a Cloud SQL instance.
|
||||||
|
|
||||||
## Cloud SQL for PostgreSQL
|
## Cloud SQL for PostgreSQL
|
||||||
|
|
||||||
@@ -284,6 +286,7 @@ See [Usage Examples](../reference/cli.md#examples).
|
|||||||
* `create_instance`
|
* `create_instance`
|
||||||
* `create_user`
|
* `create_user`
|
||||||
* `clone_instance`
|
* `clone_instance`
|
||||||
|
* `restore_backup`
|
||||||
* **Tools:**
|
* **Tools:**
|
||||||
* `create_instance`: Creates a new Cloud SQL for PostgreSQL instance.
|
* `create_instance`: Creates a new Cloud SQL for PostgreSQL instance.
|
||||||
* `get_instance`: Gets information about a Cloud SQL instance.
|
* `get_instance`: Gets information about a Cloud SQL instance.
|
||||||
@@ -294,6 +297,7 @@ See [Usage Examples](../reference/cli.md#examples).
|
|||||||
* `wait_for_operation`: Waits for a Cloud SQL operation to complete.
|
* `wait_for_operation`: Waits for a Cloud SQL operation to complete.
|
||||||
* `clone_instance`: Creates a clone for an existing Cloud SQL for PostgreSQL instance.
|
* `clone_instance`: Creates a clone for an existing Cloud SQL for PostgreSQL instance.
|
||||||
* `create_backup`: Creates a backup on a Cloud SQL instance.
|
* `create_backup`: Creates a backup on a Cloud SQL instance.
|
||||||
|
* `restore_backup`: Restores a backup of a Cloud SQL instance.
|
||||||
|
|
||||||
## Cloud SQL for SQL Server
|
## Cloud SQL for SQL Server
|
||||||
|
|
||||||
@@ -347,6 +351,7 @@ See [Usage Examples](../reference/cli.md#examples).
|
|||||||
* `create_instance`
|
* `create_instance`
|
||||||
* `create_user`
|
* `create_user`
|
||||||
* `clone_instance`
|
* `clone_instance`
|
||||||
|
* `restore_backup`
|
||||||
* **Tools:**
|
* **Tools:**
|
||||||
* `create_instance`: Creates a new Cloud SQL for SQL Server instance.
|
* `create_instance`: Creates a new Cloud SQL for SQL Server instance.
|
||||||
* `get_instance`: Gets information about a Cloud SQL instance.
|
* `get_instance`: Gets information about a Cloud SQL instance.
|
||||||
@@ -357,6 +362,7 @@ See [Usage Examples](../reference/cli.md#examples).
|
|||||||
* `wait_for_operation`: Waits for a Cloud SQL operation to complete.
|
* `wait_for_operation`: Waits for a Cloud SQL operation to complete.
|
||||||
* `clone_instance`: Creates a clone for an existing Cloud SQL for SQL Server instance.
|
* `clone_instance`: Creates a clone for an existing Cloud SQL for SQL Server instance.
|
||||||
* `create_backup`: Creates a backup on a Cloud SQL instance.
|
* `create_backup`: Creates a backup on a Cloud SQL instance.
|
||||||
|
* `restore_backup`: Restores a backup of a Cloud SQL instance.
|
||||||
|
|
||||||
## Dataplex
|
## Dataplex
|
||||||
|
|
||||||
|
|||||||
@@ -7,6 +7,17 @@ description: >
|
|||||||
|
|
||||||
---
|
---
|
||||||
|
|
||||||
|
{{< notice note >}}
|
||||||
|
**⚠️ Best Effort Maintenance**
|
||||||
|
|
||||||
|
This integration is maintained on a best-effort basis by the project
|
||||||
|
team/community. While we strive to address issues and provide workarounds when
|
||||||
|
resources are available, there are no guaranteed response times or code fixes.
|
||||||
|
|
||||||
|
The automated integration tests for this module are currently non-functional or
|
||||||
|
failing.
|
||||||
|
{{< /notice >}}
|
||||||
|
|
||||||
## About
|
## About
|
||||||
|
|
||||||
[Dgraph][dgraph-docs] is an open-source graph database. It is designed for
|
[Dgraph][dgraph-docs] is an open-source graph database. It is designed for
|
||||||
|
|||||||
@@ -41,13 +41,13 @@ tools:
|
|||||||
|
|
||||||
### Usage Flow
|
### Usage Flow
|
||||||
|
|
||||||
When using this tool, a `prompt` parameter containing a natural language query is provided to the tool (typically by an agent). The tool then interacts with the Gemini Data Analytics API using the context defined in your configuration.
|
When using this tool, a `query` parameter containing a natural language query is provided to the tool (typically by an agent). The tool then interacts with the Gemini Data Analytics API using the context defined in your configuration.
|
||||||
|
|
||||||
The structure of the response depends on the `generationOptions` configured in your tool definition (e.g., enabling `generateQueryResult` will include the SQL query results).
|
The structure of the response depends on the `generationOptions` configured in your tool definition (e.g., enabling `generateQueryResult` will include the SQL query results).
|
||||||
|
|
||||||
See [Data Analytics API REST documentation](https://clouddocs.devsite.corp.google.com/gemini/docs/conversational-analytics-api/reference/rest/v1alpha/projects.locations/queryData?rep_location=global) for details.
|
See [Data Analytics API REST documentation](https://clouddocs.devsite.corp.google.com/gemini/docs/conversational-analytics-api/reference/rest/v1alpha/projects.locations/queryData?rep_location=global) for details.
|
||||||
|
|
||||||
**Example Input Prompt:**
|
**Example Input Query:**
|
||||||
|
|
||||||
```text
|
```text
|
||||||
How many accounts who have region in Prague are eligible for loans? A3 contains the data of region.
|
How many accounts who have region in Prague are eligible for loans? A3 contains the data of region.
|
||||||
|
|||||||
53
docs/en/resources/tools/cloudsql/cloudsqlrestorebackup.md
Normal file
53
docs/en/resources/tools/cloudsql/cloudsqlrestorebackup.md
Normal file
@@ -0,0 +1,53 @@
|
|||||||
|
---
|
||||||
|
title: cloud-sql-restore-backup
|
||||||
|
type: docs
|
||||||
|
weight: 10
|
||||||
|
description: "Restores a backup of a Cloud SQL instance."
|
||||||
|
---
|
||||||
|
|
||||||
|
The `cloud-sql-restore-backup` tool restores a backup on a Cloud SQL instance using the Cloud SQL Admin API.
|
||||||
|
|
||||||
|
{{< notice info dd>}}
|
||||||
|
This tool uses a `source` of kind `cloud-sql-admin`.
|
||||||
|
{{< /notice >}}
|
||||||
|
|
||||||
|
## Examples
|
||||||
|
|
||||||
|
Basic backup restore
|
||||||
|
|
||||||
|
```yaml
|
||||||
|
tools:
|
||||||
|
backup-restore-basic:
|
||||||
|
kind: cloud-sql-restore-backup
|
||||||
|
source: cloud-sql-admin-source
|
||||||
|
description: "Restores a backup onto the given Cloud SQL instance."
|
||||||
|
```
|
||||||
|
|
||||||
|
## Reference
|
||||||
|
|
||||||
|
### Tool Configuration
|
||||||
|
| **field** | **type** | **required** | **description** |
|
||||||
|
| -------------- | :------: | :----------: | ------------------------------------------------ |
|
||||||
|
| kind | string | true | Must be "cloud-sql-restore-backup". |
|
||||||
|
| source | string | true | The name of the `cloud-sql-admin` source to use. |
|
||||||
|
| description | string | false | A description of the tool. |
|
||||||
|
|
||||||
|
### Tool Inputs
|
||||||
|
|
||||||
|
| **parameter** | **type** | **required** | **description** |
|
||||||
|
| ------------------| :------: | :----------: | -----------------------------------------------------------------------------|
|
||||||
|
| target_project | string | true | The project ID of the instance to restore the backup onto. |
|
||||||
|
| target_instance | string | true | The instance to restore the backup onto. Does not include the project ID. |
|
||||||
|
| backup_id | string | true | The identifier of the backup being restored. |
|
||||||
|
| source_project | string | false | (Optional) The project ID of the instance that the backup belongs to. |
|
||||||
|
| source_instance | string | false | (Optional) Cloud SQL instance ID of the instance that the backup belongs to. |
|
||||||
|
|
||||||
|
## Usage Notes
|
||||||
|
|
||||||
|
- The `backup_id` field can be a BackupRun ID (which will be an int64), backup name, or BackupDR backup name.
|
||||||
|
- If the `backup_id` field contains a BackupRun ID (i.e. an int64), the optional fields `source_project` and `source_instance` must also be provided.
|
||||||
|
|
||||||
|
## See Also
|
||||||
|
- [Cloud SQL Admin API documentation](https://cloud.google.com/sql/docs/mysql/admin-api)
|
||||||
|
- [Toolbox Cloud SQL tools documentation](../cloudsql)
|
||||||
|
- [Cloud SQL Restore API documentation](https://cloud.google.com/sql/docs/mysql/backup-recovery/restoring)
|
||||||
@@ -9,6 +9,17 @@ aliases:
|
|||||||
- /resources/tools/dgraph-dql
|
- /resources/tools/dgraph-dql
|
||||||
---
|
---
|
||||||
|
|
||||||
|
{{< notice note >}}
|
||||||
|
**⚠️ Best Effort Maintenance**
|
||||||
|
|
||||||
|
This integration is maintained on a best-effort basis by the project
|
||||||
|
team/community. While we strive to address issues and provide workarounds when
|
||||||
|
resources are available, there are no guaranteed response times or code fixes.
|
||||||
|
|
||||||
|
The automated integration tests for this module are currently non-functional or
|
||||||
|
failing.
|
||||||
|
{{< /notice >}}
|
||||||
|
|
||||||
## About
|
## About
|
||||||
|
|
||||||
A `dgraph-dql` tool executes a pre-defined DQL statement against a Dgraph
|
A `dgraph-dql` tool executes a pre-defined DQL statement against a Dgraph
|
||||||
|
|||||||
@@ -46,6 +46,9 @@ tools:
|
|||||||
create_backup:
|
create_backup:
|
||||||
kind: cloud-sql-create-backup
|
kind: cloud-sql-create-backup
|
||||||
source: cloud-sql-admin-source
|
source: cloud-sql-admin-source
|
||||||
|
restore_backup:
|
||||||
|
kind: cloud-sql-restore-backup
|
||||||
|
source: cloud-sql-admin-source
|
||||||
|
|
||||||
toolsets:
|
toolsets:
|
||||||
cloud_sql_mssql_admin_tools:
|
cloud_sql_mssql_admin_tools:
|
||||||
@@ -58,3 +61,4 @@ toolsets:
|
|||||||
- wait_for_operation
|
- wait_for_operation
|
||||||
- clone_instance
|
- clone_instance
|
||||||
- create_backup
|
- create_backup
|
||||||
|
- restore_backup
|
||||||
|
|||||||
@@ -46,6 +46,9 @@ tools:
|
|||||||
create_backup:
|
create_backup:
|
||||||
kind: cloud-sql-create-backup
|
kind: cloud-sql-create-backup
|
||||||
source: cloud-sql-admin-source
|
source: cloud-sql-admin-source
|
||||||
|
restore_backup:
|
||||||
|
kind: cloud-sql-restore-backup
|
||||||
|
source: cloud-sql-admin-source
|
||||||
|
|
||||||
toolsets:
|
toolsets:
|
||||||
cloud_sql_mysql_admin_tools:
|
cloud_sql_mysql_admin_tools:
|
||||||
@@ -58,3 +61,4 @@ toolsets:
|
|||||||
- wait_for_operation
|
- wait_for_operation
|
||||||
- clone_instance
|
- clone_instance
|
||||||
- create_backup
|
- create_backup
|
||||||
|
- restore_backup
|
||||||
|
|||||||
@@ -49,6 +49,9 @@ tools:
|
|||||||
create_backup:
|
create_backup:
|
||||||
kind: cloud-sql-create-backup
|
kind: cloud-sql-create-backup
|
||||||
source: cloud-sql-admin-source
|
source: cloud-sql-admin-source
|
||||||
|
restore_backup:
|
||||||
|
kind: cloud-sql-restore-backup
|
||||||
|
source: cloud-sql-admin-source
|
||||||
|
|
||||||
toolsets:
|
toolsets:
|
||||||
cloud_sql_postgres_admin_tools:
|
cloud_sql_postgres_admin_tools:
|
||||||
@@ -62,3 +65,4 @@ toolsets:
|
|||||||
- postgres_upgrade_precheck
|
- postgres_upgrade_precheck
|
||||||
- clone_instance
|
- clone_instance
|
||||||
- create_backup
|
- create_backup
|
||||||
|
- restore_backup
|
||||||
|
|||||||
@@ -27,19 +27,21 @@ import (
|
|||||||
v20241105 "github.com/googleapis/genai-toolbox/internal/server/mcp/v20241105"
|
v20241105 "github.com/googleapis/genai-toolbox/internal/server/mcp/v20241105"
|
||||||
v20250326 "github.com/googleapis/genai-toolbox/internal/server/mcp/v20250326"
|
v20250326 "github.com/googleapis/genai-toolbox/internal/server/mcp/v20250326"
|
||||||
v20250618 "github.com/googleapis/genai-toolbox/internal/server/mcp/v20250618"
|
v20250618 "github.com/googleapis/genai-toolbox/internal/server/mcp/v20250618"
|
||||||
|
v20251125 "github.com/googleapis/genai-toolbox/internal/server/mcp/v20251125"
|
||||||
"github.com/googleapis/genai-toolbox/internal/server/resources"
|
"github.com/googleapis/genai-toolbox/internal/server/resources"
|
||||||
"github.com/googleapis/genai-toolbox/internal/tools"
|
"github.com/googleapis/genai-toolbox/internal/tools"
|
||||||
)
|
)
|
||||||
|
|
||||||
// LATEST_PROTOCOL_VERSION is the latest version of the MCP protocol supported.
|
// LATEST_PROTOCOL_VERSION is the latest version of the MCP protocol supported.
|
||||||
// Update the version used in InitializeResponse when this value is updated.
|
// Update the version used in InitializeResponse when this value is updated.
|
||||||
const LATEST_PROTOCOL_VERSION = v20250618.PROTOCOL_VERSION
|
const LATEST_PROTOCOL_VERSION = v20251125.PROTOCOL_VERSION
|
||||||
|
|
||||||
// SUPPORTED_PROTOCOL_VERSIONS is the MCP protocol versions that are supported.
|
// SUPPORTED_PROTOCOL_VERSIONS is the MCP protocol versions that are supported.
|
||||||
var SUPPORTED_PROTOCOL_VERSIONS = []string{
|
var SUPPORTED_PROTOCOL_VERSIONS = []string{
|
||||||
v20241105.PROTOCOL_VERSION,
|
v20241105.PROTOCOL_VERSION,
|
||||||
v20250326.PROTOCOL_VERSION,
|
v20250326.PROTOCOL_VERSION,
|
||||||
v20250618.PROTOCOL_VERSION,
|
v20250618.PROTOCOL_VERSION,
|
||||||
|
v20251125.PROTOCOL_VERSION,
|
||||||
}
|
}
|
||||||
|
|
||||||
// InitializeResponse runs capability negotiation and protocol version agreement.
|
// InitializeResponse runs capability negotiation and protocol version agreement.
|
||||||
@@ -102,6 +104,8 @@ func NotificationHandler(ctx context.Context, body []byte) error {
|
|||||||
// This is the Operation phase of the lifecycle for MCP client-server connections.
|
// This is the Operation phase of the lifecycle for MCP client-server connections.
|
||||||
func ProcessMethod(ctx context.Context, mcpVersion string, id jsonrpc.RequestId, method string, toolset tools.Toolset, promptset prompts.Promptset, resourceMgr *resources.ResourceManager, body []byte, header http.Header) (any, error) {
|
func ProcessMethod(ctx context.Context, mcpVersion string, id jsonrpc.RequestId, method string, toolset tools.Toolset, promptset prompts.Promptset, resourceMgr *resources.ResourceManager, body []byte, header http.Header) (any, error) {
|
||||||
switch mcpVersion {
|
switch mcpVersion {
|
||||||
|
case v20251125.PROTOCOL_VERSION:
|
||||||
|
return v20251125.ProcessMethod(ctx, id, method, toolset, promptset, resourceMgr, body, header)
|
||||||
case v20250618.PROTOCOL_VERSION:
|
case v20250618.PROTOCOL_VERSION:
|
||||||
return v20250618.ProcessMethod(ctx, id, method, toolset, promptset, resourceMgr, body, header)
|
return v20250618.ProcessMethod(ctx, id, method, toolset, promptset, resourceMgr, body, header)
|
||||||
case v20250326.PROTOCOL_VERSION:
|
case v20250326.PROTOCOL_VERSION:
|
||||||
|
|||||||
326
internal/server/mcp/v20251125/method.go
Normal file
326
internal/server/mcp/v20251125/method.go
Normal file
@@ -0,0 +1,326 @@
|
|||||||
|
// 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 v20251125
|
||||||
|
|
||||||
|
import (
|
||||||
|
"bytes"
|
||||||
|
"context"
|
||||||
|
"encoding/json"
|
||||||
|
"errors"
|
||||||
|
"fmt"
|
||||||
|
"net/http"
|
||||||
|
"strings"
|
||||||
|
|
||||||
|
"github.com/googleapis/genai-toolbox/internal/prompts"
|
||||||
|
"github.com/googleapis/genai-toolbox/internal/server/mcp/jsonrpc"
|
||||||
|
"github.com/googleapis/genai-toolbox/internal/server/resources"
|
||||||
|
"github.com/googleapis/genai-toolbox/internal/tools"
|
||||||
|
"github.com/googleapis/genai-toolbox/internal/util"
|
||||||
|
)
|
||||||
|
|
||||||
|
// ProcessMethod returns a response for the request.
|
||||||
|
func ProcessMethod(ctx context.Context, id jsonrpc.RequestId, method string, toolset tools.Toolset, promptset prompts.Promptset, resourceMgr *resources.ResourceManager, body []byte, header http.Header) (any, error) {
|
||||||
|
switch method {
|
||||||
|
case PING:
|
||||||
|
return pingHandler(id)
|
||||||
|
case TOOLS_LIST:
|
||||||
|
return toolsListHandler(id, toolset, body)
|
||||||
|
case TOOLS_CALL:
|
||||||
|
return toolsCallHandler(ctx, id, resourceMgr, body, header)
|
||||||
|
case PROMPTS_LIST:
|
||||||
|
return promptsListHandler(ctx, id, promptset, body)
|
||||||
|
case PROMPTS_GET:
|
||||||
|
return promptsGetHandler(ctx, id, resourceMgr, body)
|
||||||
|
default:
|
||||||
|
err := fmt.Errorf("invalid method %s", method)
|
||||||
|
return jsonrpc.NewError(id, jsonrpc.METHOD_NOT_FOUND, err.Error(), nil), err
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// pingHandler handles the "ping" method by returning an empty response.
|
||||||
|
func pingHandler(id jsonrpc.RequestId) (any, error) {
|
||||||
|
return jsonrpc.JSONRPCResponse{
|
||||||
|
Jsonrpc: jsonrpc.JSONRPC_VERSION,
|
||||||
|
Id: id,
|
||||||
|
Result: struct{}{},
|
||||||
|
}, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
func toolsListHandler(id jsonrpc.RequestId, toolset tools.Toolset, body []byte) (any, error) {
|
||||||
|
var req ListToolsRequest
|
||||||
|
if err := json.Unmarshal(body, &req); err != nil {
|
||||||
|
err = fmt.Errorf("invalid mcp tools list request: %w", err)
|
||||||
|
return jsonrpc.NewError(id, jsonrpc.INVALID_REQUEST, err.Error(), nil), err
|
||||||
|
}
|
||||||
|
|
||||||
|
result := ListToolsResult{
|
||||||
|
Tools: toolset.McpManifest,
|
||||||
|
}
|
||||||
|
return jsonrpc.JSONRPCResponse{
|
||||||
|
Jsonrpc: jsonrpc.JSONRPC_VERSION,
|
||||||
|
Id: id,
|
||||||
|
Result: result,
|
||||||
|
}, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
// toolsCallHandler generate a response for tools call.
|
||||||
|
func toolsCallHandler(ctx context.Context, id jsonrpc.RequestId, resourceMgr *resources.ResourceManager, body []byte, header http.Header) (any, error) {
|
||||||
|
authServices := resourceMgr.GetAuthServiceMap()
|
||||||
|
|
||||||
|
// retrieve logger from context
|
||||||
|
logger, err := util.LoggerFromContext(ctx)
|
||||||
|
if err != nil {
|
||||||
|
return jsonrpc.NewError(id, jsonrpc.INTERNAL_ERROR, err.Error(), nil), err
|
||||||
|
}
|
||||||
|
|
||||||
|
var req CallToolRequest
|
||||||
|
if err = json.Unmarshal(body, &req); err != nil {
|
||||||
|
err = fmt.Errorf("invalid mcp tools call request: %w", err)
|
||||||
|
return jsonrpc.NewError(id, jsonrpc.INVALID_REQUEST, err.Error(), nil), err
|
||||||
|
}
|
||||||
|
|
||||||
|
toolName := req.Params.Name
|
||||||
|
toolArgument := req.Params.Arguments
|
||||||
|
logger.DebugContext(ctx, fmt.Sprintf("tool name: %s", toolName))
|
||||||
|
tool, ok := resourceMgr.GetTool(toolName)
|
||||||
|
if !ok {
|
||||||
|
err = fmt.Errorf("invalid tool name: tool with name %q does not exist", toolName)
|
||||||
|
return jsonrpc.NewError(id, jsonrpc.INVALID_PARAMS, err.Error(), nil), err
|
||||||
|
}
|
||||||
|
|
||||||
|
// Get access token
|
||||||
|
authTokenHeadername, err := tool.GetAuthTokenHeaderName(resourceMgr)
|
||||||
|
if err != nil {
|
||||||
|
errMsg := fmt.Errorf("error during invocation: %w", err)
|
||||||
|
return jsonrpc.NewError(id, jsonrpc.INTERNAL_ERROR, errMsg.Error(), nil), errMsg
|
||||||
|
}
|
||||||
|
accessToken := tools.AccessToken(header.Get(authTokenHeadername))
|
||||||
|
|
||||||
|
// Check if this specific tool requires the standard authorization header
|
||||||
|
clientAuth, err := tool.RequiresClientAuthorization(resourceMgr)
|
||||||
|
if err != nil {
|
||||||
|
errMsg := fmt.Errorf("error during invocation: %w", err)
|
||||||
|
return jsonrpc.NewError(id, jsonrpc.INTERNAL_ERROR, errMsg.Error(), nil), errMsg
|
||||||
|
}
|
||||||
|
if clientAuth {
|
||||||
|
if accessToken == "" {
|
||||||
|
return jsonrpc.NewError(id, jsonrpc.INVALID_REQUEST, "missing access token in the 'Authorization' header", nil), util.ErrUnauthorized
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// marshal arguments and decode it using decodeJSON instead to prevent loss between floats/int.
|
||||||
|
aMarshal, err := json.Marshal(toolArgument)
|
||||||
|
if err != nil {
|
||||||
|
err = fmt.Errorf("unable to marshal tools argument: %w", err)
|
||||||
|
return jsonrpc.NewError(id, jsonrpc.INTERNAL_ERROR, err.Error(), nil), err
|
||||||
|
}
|
||||||
|
|
||||||
|
var data map[string]any
|
||||||
|
if err = util.DecodeJSON(bytes.NewBuffer(aMarshal), &data); err != nil {
|
||||||
|
err = fmt.Errorf("unable to decode tools argument: %w", err)
|
||||||
|
return jsonrpc.NewError(id, jsonrpc.INTERNAL_ERROR, err.Error(), nil), err
|
||||||
|
}
|
||||||
|
|
||||||
|
// Tool authentication
|
||||||
|
// claimsFromAuth maps the name of the authservice to the claims retrieved from it.
|
||||||
|
claimsFromAuth := make(map[string]map[string]any)
|
||||||
|
|
||||||
|
// if using stdio, header will be nil and auth will not be supported
|
||||||
|
if header != nil {
|
||||||
|
for _, aS := range authServices {
|
||||||
|
claims, err := aS.GetClaimsFromHeader(ctx, header)
|
||||||
|
if err != nil {
|
||||||
|
logger.DebugContext(ctx, err.Error())
|
||||||
|
continue
|
||||||
|
}
|
||||||
|
if claims == nil {
|
||||||
|
// authService not present in header
|
||||||
|
continue
|
||||||
|
}
|
||||||
|
claimsFromAuth[aS.GetName()] = claims
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// Tool authorization check
|
||||||
|
verifiedAuthServices := make([]string, len(claimsFromAuth))
|
||||||
|
i := 0
|
||||||
|
for k := range claimsFromAuth {
|
||||||
|
verifiedAuthServices[i] = k
|
||||||
|
i++
|
||||||
|
}
|
||||||
|
|
||||||
|
// Check if any of the specified auth services is verified
|
||||||
|
isAuthorized := tool.Authorized(verifiedAuthServices)
|
||||||
|
if !isAuthorized {
|
||||||
|
err = fmt.Errorf("unauthorized Tool call: Please make sure your specify correct auth headers: %w", util.ErrUnauthorized)
|
||||||
|
return jsonrpc.NewError(id, jsonrpc.INVALID_REQUEST, err.Error(), nil), err
|
||||||
|
}
|
||||||
|
logger.DebugContext(ctx, "tool invocation authorized")
|
||||||
|
|
||||||
|
params, err := tool.ParseParams(data, claimsFromAuth)
|
||||||
|
if err != nil {
|
||||||
|
err = fmt.Errorf("provided parameters were invalid: %w", err)
|
||||||
|
return jsonrpc.NewError(id, jsonrpc.INVALID_PARAMS, err.Error(), nil), err
|
||||||
|
}
|
||||||
|
logger.DebugContext(ctx, fmt.Sprintf("invocation params: %s", params))
|
||||||
|
|
||||||
|
// run tool invocation and generate response.
|
||||||
|
results, err := tool.Invoke(ctx, resourceMgr, params, accessToken)
|
||||||
|
if err != nil {
|
||||||
|
errStr := err.Error()
|
||||||
|
// Missing authService tokens.
|
||||||
|
if errors.Is(err, util.ErrUnauthorized) {
|
||||||
|
return jsonrpc.NewError(id, jsonrpc.INVALID_REQUEST, err.Error(), nil), err
|
||||||
|
}
|
||||||
|
// Upstream auth error
|
||||||
|
if strings.Contains(errStr, "Error 401") || strings.Contains(errStr, "Error 403") {
|
||||||
|
if clientAuth {
|
||||||
|
// Error with client credentials should pass down to the client
|
||||||
|
return jsonrpc.NewError(id, jsonrpc.INVALID_REQUEST, err.Error(), nil), err
|
||||||
|
}
|
||||||
|
// Auth error with ADC should raise internal 500 error
|
||||||
|
return jsonrpc.NewError(id, jsonrpc.INTERNAL_ERROR, err.Error(), nil), err
|
||||||
|
}
|
||||||
|
text := TextContent{
|
||||||
|
Type: "text",
|
||||||
|
Text: err.Error(),
|
||||||
|
}
|
||||||
|
return jsonrpc.JSONRPCResponse{
|
||||||
|
Jsonrpc: jsonrpc.JSONRPC_VERSION,
|
||||||
|
Id: id,
|
||||||
|
Result: CallToolResult{Content: []TextContent{text}, IsError: true},
|
||||||
|
}, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
content := make([]TextContent, 0)
|
||||||
|
|
||||||
|
sliceRes, ok := results.([]any)
|
||||||
|
if !ok {
|
||||||
|
sliceRes = []any{results}
|
||||||
|
}
|
||||||
|
|
||||||
|
for _, d := range sliceRes {
|
||||||
|
text := TextContent{Type: "text"}
|
||||||
|
dM, err := json.Marshal(d)
|
||||||
|
if err != nil {
|
||||||
|
text.Text = fmt.Sprintf("fail to marshal: %s, result: %s", err, d)
|
||||||
|
} else {
|
||||||
|
text.Text = string(dM)
|
||||||
|
}
|
||||||
|
content = append(content, text)
|
||||||
|
}
|
||||||
|
|
||||||
|
return jsonrpc.JSONRPCResponse{
|
||||||
|
Jsonrpc: jsonrpc.JSONRPC_VERSION,
|
||||||
|
Id: id,
|
||||||
|
Result: CallToolResult{Content: content},
|
||||||
|
}, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
// promptsListHandler handles the "prompts/list" method.
|
||||||
|
func promptsListHandler(ctx context.Context, id jsonrpc.RequestId, promptset prompts.Promptset, body []byte) (any, error) {
|
||||||
|
// retrieve logger from context
|
||||||
|
logger, err := util.LoggerFromContext(ctx)
|
||||||
|
if err != nil {
|
||||||
|
return jsonrpc.NewError(id, jsonrpc.INTERNAL_ERROR, err.Error(), nil), err
|
||||||
|
}
|
||||||
|
logger.DebugContext(ctx, "handling prompts/list request")
|
||||||
|
|
||||||
|
var req ListPromptsRequest
|
||||||
|
if err := json.Unmarshal(body, &req); err != nil {
|
||||||
|
err = fmt.Errorf("invalid mcp prompts list request: %w", err)
|
||||||
|
return jsonrpc.NewError(id, jsonrpc.INVALID_REQUEST, err.Error(), nil), err
|
||||||
|
}
|
||||||
|
|
||||||
|
result := ListPromptsResult{
|
||||||
|
Prompts: promptset.McpManifest,
|
||||||
|
}
|
||||||
|
logger.DebugContext(ctx, fmt.Sprintf("returning %d prompts", len(promptset.McpManifest)))
|
||||||
|
return jsonrpc.JSONRPCResponse{
|
||||||
|
Jsonrpc: jsonrpc.JSONRPC_VERSION,
|
||||||
|
Id: id,
|
||||||
|
Result: result,
|
||||||
|
}, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
// promptsGetHandler handles the "prompts/get" method.
|
||||||
|
func promptsGetHandler(ctx context.Context, id jsonrpc.RequestId, resourceMgr *resources.ResourceManager, body []byte) (any, error) {
|
||||||
|
// retrieve logger from context
|
||||||
|
logger, err := util.LoggerFromContext(ctx)
|
||||||
|
if err != nil {
|
||||||
|
return jsonrpc.NewError(id, jsonrpc.INTERNAL_ERROR, err.Error(), nil), err
|
||||||
|
}
|
||||||
|
logger.DebugContext(ctx, "handling prompts/get request")
|
||||||
|
|
||||||
|
var req GetPromptRequest
|
||||||
|
if err := json.Unmarshal(body, &req); err != nil {
|
||||||
|
err = fmt.Errorf("invalid mcp prompts/get request: %w", err)
|
||||||
|
return jsonrpc.NewError(id, jsonrpc.INVALID_REQUEST, err.Error(), nil), err
|
||||||
|
}
|
||||||
|
|
||||||
|
promptName := req.Params.Name
|
||||||
|
logger.DebugContext(ctx, fmt.Sprintf("prompt name: %s", promptName))
|
||||||
|
prompt, ok := resourceMgr.GetPrompt(promptName)
|
||||||
|
if !ok {
|
||||||
|
err := fmt.Errorf("prompt with name %q does not exist", promptName)
|
||||||
|
return jsonrpc.NewError(id, jsonrpc.INVALID_PARAMS, err.Error(), nil), err
|
||||||
|
}
|
||||||
|
|
||||||
|
// Parse the arguments provided in the request.
|
||||||
|
argValues, err := prompt.ParseArgs(req.Params.Arguments, nil)
|
||||||
|
if err != nil {
|
||||||
|
err = fmt.Errorf("invalid arguments for prompt %q: %w", promptName, err)
|
||||||
|
return jsonrpc.NewError(id, jsonrpc.INVALID_PARAMS, err.Error(), nil), err
|
||||||
|
}
|
||||||
|
logger.DebugContext(ctx, fmt.Sprintf("parsed args: %v", argValues))
|
||||||
|
|
||||||
|
// Substitute the argument values into the prompt's messages.
|
||||||
|
substituted, err := prompt.SubstituteParams(argValues)
|
||||||
|
if err != nil {
|
||||||
|
err = fmt.Errorf("error substituting params for prompt %q: %w", promptName, err)
|
||||||
|
return jsonrpc.NewError(id, jsonrpc.INTERNAL_ERROR, err.Error(), nil), err
|
||||||
|
}
|
||||||
|
|
||||||
|
// Cast the result to the expected []prompts.Message type.
|
||||||
|
substitutedMessages, ok := substituted.([]prompts.Message)
|
||||||
|
if !ok {
|
||||||
|
err = fmt.Errorf("internal error: SubstituteParams returned unexpected type")
|
||||||
|
return jsonrpc.NewError(id, jsonrpc.INTERNAL_ERROR, err.Error(), nil), err
|
||||||
|
}
|
||||||
|
logger.DebugContext(ctx, "substituted params successfully")
|
||||||
|
|
||||||
|
// Format the response messages into the required structure.
|
||||||
|
promptMessages := make([]PromptMessage, len(substitutedMessages))
|
||||||
|
for i, msg := range substitutedMessages {
|
||||||
|
promptMessages[i] = PromptMessage{
|
||||||
|
Role: msg.Role,
|
||||||
|
Content: TextContent{
|
||||||
|
Type: "text",
|
||||||
|
Text: msg.Content,
|
||||||
|
},
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
result := GetPromptResult{
|
||||||
|
Description: prompt.Manifest().Description,
|
||||||
|
Messages: promptMessages,
|
||||||
|
}
|
||||||
|
|
||||||
|
return jsonrpc.JSONRPCResponse{
|
||||||
|
Jsonrpc: jsonrpc.JSONRPC_VERSION,
|
||||||
|
Id: id,
|
||||||
|
Result: result,
|
||||||
|
}, nil
|
||||||
|
}
|
||||||
219
internal/server/mcp/v20251125/types.go
Normal file
219
internal/server/mcp/v20251125/types.go
Normal file
@@ -0,0 +1,219 @@
|
|||||||
|
// 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 v20251125
|
||||||
|
|
||||||
|
import (
|
||||||
|
"github.com/googleapis/genai-toolbox/internal/prompts"
|
||||||
|
"github.com/googleapis/genai-toolbox/internal/server/mcp/jsonrpc"
|
||||||
|
"github.com/googleapis/genai-toolbox/internal/tools"
|
||||||
|
)
|
||||||
|
|
||||||
|
// SERVER_NAME is the server name used in Implementation.
|
||||||
|
const SERVER_NAME = "Toolbox"
|
||||||
|
|
||||||
|
// PROTOCOL_VERSION is the version of the MCP protocol in this package.
|
||||||
|
const PROTOCOL_VERSION = "2025-11-25"
|
||||||
|
|
||||||
|
// methods that are supported.
|
||||||
|
const (
|
||||||
|
PING = "ping"
|
||||||
|
TOOLS_LIST = "tools/list"
|
||||||
|
TOOLS_CALL = "tools/call"
|
||||||
|
PROMPTS_LIST = "prompts/list"
|
||||||
|
PROMPTS_GET = "prompts/get"
|
||||||
|
)
|
||||||
|
|
||||||
|
/* Empty result */
|
||||||
|
|
||||||
|
// EmptyResult represents a response that indicates success but carries no data.
|
||||||
|
type EmptyResult jsonrpc.Result
|
||||||
|
|
||||||
|
/* Pagination */
|
||||||
|
|
||||||
|
// Cursor is an opaque token used to represent a cursor for pagination.
|
||||||
|
type Cursor string
|
||||||
|
|
||||||
|
type PaginatedRequest struct {
|
||||||
|
jsonrpc.Request
|
||||||
|
Params struct {
|
||||||
|
// An opaque token representing the current pagination position.
|
||||||
|
// If provided, the server should return results starting after this cursor.
|
||||||
|
Cursor Cursor `json:"cursor,omitempty"`
|
||||||
|
} `json:"params,omitempty"`
|
||||||
|
}
|
||||||
|
|
||||||
|
type PaginatedResult struct {
|
||||||
|
jsonrpc.Result
|
||||||
|
// An opaque token representing the pagination position after the last returned result.
|
||||||
|
// If present, there may be more results available.
|
||||||
|
NextCursor Cursor `json:"nextCursor,omitempty"`
|
||||||
|
}
|
||||||
|
|
||||||
|
/* Tools */
|
||||||
|
|
||||||
|
// Sent from the client to request a list of tools the server has.
|
||||||
|
type ListToolsRequest struct {
|
||||||
|
PaginatedRequest
|
||||||
|
}
|
||||||
|
|
||||||
|
// The server's response to a tools/list request from the client.
|
||||||
|
type ListToolsResult struct {
|
||||||
|
PaginatedResult
|
||||||
|
Tools []tools.McpManifest `json:"tools"`
|
||||||
|
}
|
||||||
|
|
||||||
|
// Used by the client to invoke a tool provided by the server.
|
||||||
|
type CallToolRequest struct {
|
||||||
|
jsonrpc.Request
|
||||||
|
Params struct {
|
||||||
|
Name string `json:"name"`
|
||||||
|
Arguments map[string]any `json:"arguments,omitempty"`
|
||||||
|
} `json:"params,omitempty"`
|
||||||
|
}
|
||||||
|
|
||||||
|
// The sender or recipient of messages and data in a conversation.
|
||||||
|
type Role string
|
||||||
|
|
||||||
|
const (
|
||||||
|
RoleUser Role = "user"
|
||||||
|
RoleAssistant Role = "assistant"
|
||||||
|
)
|
||||||
|
|
||||||
|
// Base for objects that include optional annotations for the client.
|
||||||
|
// The client can use annotations to inform how objects are used or displayed
|
||||||
|
type Annotated struct {
|
||||||
|
Annotations *struct {
|
||||||
|
// Describes who the intended customer of this object or data is.
|
||||||
|
// It can include multiple entries to indicate content useful for multiple
|
||||||
|
// audiences (e.g., `["user", "assistant"]`).
|
||||||
|
Audience []Role `json:"audience,omitempty"`
|
||||||
|
// Describes how important this data is for operating the server.
|
||||||
|
//
|
||||||
|
// A value of 1 means "most important," and indicates that the data is
|
||||||
|
// effectively required, while 0 means "least important," and indicates that
|
||||||
|
// the data is entirely optional.
|
||||||
|
//
|
||||||
|
// @TJS-type number
|
||||||
|
// @minimum 0
|
||||||
|
// @maximum 1
|
||||||
|
Priority float64 `json:"priority,omitempty"`
|
||||||
|
} `json:"annotations,omitempty"`
|
||||||
|
}
|
||||||
|
|
||||||
|
// TextContent represents text provided to or from an LLM.
|
||||||
|
type TextContent struct {
|
||||||
|
Annotated
|
||||||
|
Type string `json:"type"`
|
||||||
|
// The text content of the message.
|
||||||
|
Text string `json:"text"`
|
||||||
|
}
|
||||||
|
|
||||||
|
// The server's response to a tool call.
|
||||||
|
//
|
||||||
|
// Any errors that originate from the tool SHOULD be reported inside the result
|
||||||
|
// object, with `isError` set to true, _not_ as an MCP protocol-level error
|
||||||
|
// response. Otherwise, the LLM would not be able to see that an error occurred
|
||||||
|
// and self-correct.
|
||||||
|
//
|
||||||
|
// However, any errors in _finding_ the tool, an error indicating that the
|
||||||
|
// server does not support tool calls, or any other exceptional conditions,
|
||||||
|
// should be reported as an MCP error response.
|
||||||
|
type CallToolResult struct {
|
||||||
|
jsonrpc.Result
|
||||||
|
// Could be either a TextContent, ImageContent, or EmbeddedResources
|
||||||
|
// For Toolbox, we will only be sending TextContent
|
||||||
|
Content []TextContent `json:"content"`
|
||||||
|
// Whether the tool call ended in an error.
|
||||||
|
// If not set, this is assumed to be false (the call was successful).
|
||||||
|
//
|
||||||
|
// Any errors that originate from the tool SHOULD be reported inside the result
|
||||||
|
// object, with `isError` set to true, _not_ as an MCP protocol-level error
|
||||||
|
// response. Otherwise, the LLM would not be able to see that an error occurred
|
||||||
|
// and self-correct.
|
||||||
|
//
|
||||||
|
// However, any errors in _finding_ the tool, an error indicating that the
|
||||||
|
// server does not support tool calls, or any other exceptional conditions,
|
||||||
|
// should be reported as an MCP error response.
|
||||||
|
IsError bool `json:"isError,omitempty"`
|
||||||
|
// An optional JSON object that represents the structured result of the tool call.
|
||||||
|
StructuredContent map[string]any `json:"structuredContent,omitempty"`
|
||||||
|
}
|
||||||
|
|
||||||
|
// Additional properties describing a Tool to clients.
|
||||||
|
//
|
||||||
|
// NOTE: all properties in ToolAnnotations are **hints**.
|
||||||
|
// They are not guaranteed to provide a faithful description of
|
||||||
|
// tool behavior (including descriptive properties like `title`).
|
||||||
|
//
|
||||||
|
// Clients should never make tool use decisions based on ToolAnnotations
|
||||||
|
// received from untrusted servers.
|
||||||
|
type ToolAnnotations struct {
|
||||||
|
// A human-readable title for the tool.
|
||||||
|
Title string `json:"title,omitempty"`
|
||||||
|
// If true, the tool does not modify its environment.
|
||||||
|
// Default: false
|
||||||
|
ReadOnlyHint bool `json:"readOnlyHint,omitempty"`
|
||||||
|
// If true, the tool may perform destructive updates to its environment.
|
||||||
|
// If false, the tool performs only additive updates.
|
||||||
|
// (This property is meaningful only when `readOnlyHint == false`)
|
||||||
|
// Default: true
|
||||||
|
DestructiveHint bool `json:"destructiveHint,omitempty"`
|
||||||
|
// If true, calling the tool repeatedly with the same arguments
|
||||||
|
// will have no additional effect on the its environment.
|
||||||
|
// (This property is meaningful only when `readOnlyHint == false`)
|
||||||
|
// Default: false
|
||||||
|
IdempotentHint bool `json:"idempotentHint,omitempty"`
|
||||||
|
// If true, this tool may interact with an "open world" of external
|
||||||
|
// entities. If false, the tool's domain of interaction is closed.
|
||||||
|
// For example, the world of a web search tool is open, whereas that
|
||||||
|
// of a memory tool is not.
|
||||||
|
// Default: true
|
||||||
|
OpenWorldHint bool `json:"openWorldHint,omitempty"`
|
||||||
|
}
|
||||||
|
|
||||||
|
/* Prompts */
|
||||||
|
|
||||||
|
// Sent from the client to request a list of prompts the server has.
|
||||||
|
type ListPromptsRequest struct {
|
||||||
|
PaginatedRequest
|
||||||
|
}
|
||||||
|
|
||||||
|
// The server's response to a prompts/list request from the client.
|
||||||
|
type ListPromptsResult struct {
|
||||||
|
PaginatedResult
|
||||||
|
Prompts []prompts.McpManifest `json:"prompts"`
|
||||||
|
}
|
||||||
|
|
||||||
|
// Used by the client to get a prompt provided by the server.
|
||||||
|
type GetPromptRequest struct {
|
||||||
|
jsonrpc.Request
|
||||||
|
Params struct {
|
||||||
|
Name string `json:"name"`
|
||||||
|
Arguments map[string]any `json:"arguments,omitempty"`
|
||||||
|
} `json:"params"`
|
||||||
|
}
|
||||||
|
|
||||||
|
// The server's response to a prompts/get request from the client.
|
||||||
|
type GetPromptResult struct {
|
||||||
|
jsonrpc.Result
|
||||||
|
Description string `json:"description,omitempty"`
|
||||||
|
Messages []PromptMessage `json:"messages"`
|
||||||
|
}
|
||||||
|
|
||||||
|
// Describes a message returned as part of a prompt.
|
||||||
|
type PromptMessage struct {
|
||||||
|
Role string `json:"role"`
|
||||||
|
Content TextContent `json:"content"`
|
||||||
|
}
|
||||||
@@ -37,6 +37,7 @@ const jsonrpcVersion = "2.0"
|
|||||||
const protocolVersion20241105 = "2024-11-05"
|
const protocolVersion20241105 = "2024-11-05"
|
||||||
const protocolVersion20250326 = "2025-03-26"
|
const protocolVersion20250326 = "2025-03-26"
|
||||||
const protocolVersion20250618 = "2025-06-18"
|
const protocolVersion20250618 = "2025-06-18"
|
||||||
|
const protocolVersion20251125 = "2025-11-25"
|
||||||
const serverName = "Toolbox"
|
const serverName = "Toolbox"
|
||||||
|
|
||||||
var basicInputSchema = map[string]any{
|
var basicInputSchema = map[string]any{
|
||||||
@@ -485,6 +486,23 @@ func TestMcpEndpoint(t *testing.T) {
|
|||||||
},
|
},
|
||||||
},
|
},
|
||||||
},
|
},
|
||||||
|
{
|
||||||
|
name: "version 2025-11-25",
|
||||||
|
protocol: protocolVersion20251125,
|
||||||
|
idHeader: false,
|
||||||
|
initWant: map[string]any{
|
||||||
|
"jsonrpc": "2.0",
|
||||||
|
"id": "mcp-initialize",
|
||||||
|
"result": map[string]any{
|
||||||
|
"protocolVersion": "2025-11-25",
|
||||||
|
"capabilities": map[string]any{
|
||||||
|
"tools": map[string]any{"listChanged": false},
|
||||||
|
"prompts": map[string]any{"listChanged": false},
|
||||||
|
},
|
||||||
|
"serverInfo": map[string]any{"name": serverName, "version": fakeVersionString},
|
||||||
|
},
|
||||||
|
},
|
||||||
|
},
|
||||||
}
|
}
|
||||||
for _, vtc := range versTestCases {
|
for _, vtc := range versTestCases {
|
||||||
t.Run(vtc.name, func(t *testing.T) {
|
t.Run(vtc.name, func(t *testing.T) {
|
||||||
@@ -494,8 +512,7 @@ func TestMcpEndpoint(t *testing.T) {
|
|||||||
if sessionId != "" {
|
if sessionId != "" {
|
||||||
header["Mcp-Session-Id"] = sessionId
|
header["Mcp-Session-Id"] = sessionId
|
||||||
}
|
}
|
||||||
|
if vtc.protocol != protocolVersion20241105 && vtc.protocol != protocolVersion20250326 {
|
||||||
if vtc.protocol == protocolVersion20250618 {
|
|
||||||
header["MCP-Protocol-Version"] = vtc.protocol
|
header["MCP-Protocol-Version"] = vtc.protocol
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|||||||
@@ -304,10 +304,14 @@ func hostCheck(allowedHosts map[string]struct{}) func(http.Handler) http.Handler
|
|||||||
return func(next http.Handler) http.Handler {
|
return func(next http.Handler) http.Handler {
|
||||||
return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||||
_, hasWildcard := allowedHosts["*"]
|
_, hasWildcard := allowedHosts["*"]
|
||||||
_, hostIsAllowed := allowedHosts[r.Host]
|
hostname := r.Host
|
||||||
|
if host, _, err := net.SplitHostPort(r.Host); err == nil {
|
||||||
|
hostname = host
|
||||||
|
}
|
||||||
|
_, hostIsAllowed := allowedHosts[hostname]
|
||||||
if !hasWildcard && !hostIsAllowed {
|
if !hasWildcard && !hostIsAllowed {
|
||||||
// Return 400 Bad Request or 403 Forbidden to block the attack
|
// Return 403 Forbidden to block the attack
|
||||||
http.Error(w, "Invalid Host header", http.StatusBadRequest)
|
http.Error(w, "Invalid Host header", http.StatusForbidden)
|
||||||
return
|
return
|
||||||
}
|
}
|
||||||
next.ServeHTTP(w, r)
|
next.ServeHTTP(w, r)
|
||||||
@@ -406,7 +410,11 @@ func NewServer(ctx context.Context, cfg ServerConfig) (*Server, error) {
|
|||||||
}
|
}
|
||||||
allowedHostsMap := make(map[string]struct{}, len(cfg.AllowedHosts))
|
allowedHostsMap := make(map[string]struct{}, len(cfg.AllowedHosts))
|
||||||
for _, h := range cfg.AllowedHosts {
|
for _, h := range cfg.AllowedHosts {
|
||||||
allowedHostsMap[h] = struct{}{}
|
hostname := h
|
||||||
|
if host, _, err := net.SplitHostPort(h); err == nil {
|
||||||
|
hostname = host
|
||||||
|
}
|
||||||
|
allowedHostsMap[hostname] = struct{}{}
|
||||||
}
|
}
|
||||||
r.Use(hostCheck(allowedHostsMap))
|
r.Use(hostCheck(allowedHostsMap))
|
||||||
|
|
||||||
|
|||||||
@@ -19,6 +19,7 @@ import (
|
|||||||
"fmt"
|
"fmt"
|
||||||
"net/http"
|
"net/http"
|
||||||
"regexp"
|
"regexp"
|
||||||
|
"strconv"
|
||||||
"strings"
|
"strings"
|
||||||
"text/template"
|
"text/template"
|
||||||
"time"
|
"time"
|
||||||
@@ -36,7 +37,10 @@ import (
|
|||||||
|
|
||||||
const SourceKind string = "cloud-sql-admin"
|
const SourceKind string = "cloud-sql-admin"
|
||||||
|
|
||||||
var targetLinkRegex = regexp.MustCompile(`/projects/([^/]+)/instances/([^/]+)/databases/([^/]+)`)
|
var (
|
||||||
|
targetLinkRegex = regexp.MustCompile(`/projects/([^/]+)/instances/([^/]+)/databases/([^/]+)`)
|
||||||
|
backupDRRegex = regexp.MustCompile(`^projects/([^/]+)/locations/([^/]+)/backupVaults/([^/]+)/dataSources/([^/]+)/backups/([^/]+)$`)
|
||||||
|
)
|
||||||
|
|
||||||
// validate interface
|
// validate interface
|
||||||
var _ sources.SourceConfig = Config{}
|
var _ sources.SourceConfig = Config{}
|
||||||
@@ -374,6 +378,48 @@ func (s *Source) InsertBackupRun(ctx context.Context, project, instance, locatio
|
|||||||
return resp, nil
|
return resp, nil
|
||||||
}
|
}
|
||||||
|
|
||||||
|
func (s *Source) RestoreBackup(ctx context.Context, targetProject, targetInstance, sourceProject, sourceInstance, backupID, accessToken string) (any, error) {
|
||||||
|
request := &sqladmin.InstancesRestoreBackupRequest{}
|
||||||
|
|
||||||
|
// There are 3 scenarios for the backup identifier:
|
||||||
|
// 1. The identifier is an int64 containing the timestamp of the BackupRun.
|
||||||
|
// This is used to restore standard backups, and the RestoreBackupContext
|
||||||
|
// field should be populated with the backup ID and source instance info.
|
||||||
|
// 2. The identifier is a string of the format
|
||||||
|
// 'projects/{project-id}/locations/{location}/backupVaults/{backupvault}/dataSources/{datasource}/backups/{backup-uid}'.
|
||||||
|
// This is used to restore BackupDR backups, and the BackupdrBackup field
|
||||||
|
// should be populated.
|
||||||
|
// 3. The identifer is a string of the format
|
||||||
|
// 'projects/{project-id}/backups/{backup-uid}'. In this case, the Backup
|
||||||
|
// field should be populated.
|
||||||
|
if backupRunID, err := strconv.ParseInt(backupID, 10, 64); err == nil {
|
||||||
|
if sourceProject == "" || targetInstance == "" {
|
||||||
|
return nil, fmt.Errorf("source project and instance are required when restoring via backup ID")
|
||||||
|
}
|
||||||
|
request.RestoreBackupContext = &sqladmin.RestoreBackupContext{
|
||||||
|
Project: sourceProject,
|
||||||
|
InstanceId: sourceInstance,
|
||||||
|
BackupRunId: backupRunID,
|
||||||
|
}
|
||||||
|
} else if backupDRRegex.MatchString(backupID) {
|
||||||
|
request.BackupdrBackup = backupID
|
||||||
|
} else {
|
||||||
|
request.Backup = backupID
|
||||||
|
}
|
||||||
|
|
||||||
|
service, err := s.GetService(ctx, string(accessToken))
|
||||||
|
if err != nil {
|
||||||
|
return nil, err
|
||||||
|
}
|
||||||
|
|
||||||
|
resp, err := service.Instances.RestoreBackup(targetProject, targetInstance, request).Do()
|
||||||
|
if err != nil {
|
||||||
|
return nil, fmt.Errorf("error restoring backup: %w", err)
|
||||||
|
}
|
||||||
|
|
||||||
|
return resp, nil
|
||||||
|
}
|
||||||
|
|
||||||
func generateCloudSQLConnectionMessage(ctx context.Context, source *Source, logger log.Logger, opResponse map[string]any, connectionMessageTemplate string) (string, bool) {
|
func generateCloudSQLConnectionMessage(ctx context.Context, source *Source, logger log.Logger, opResponse map[string]any, connectionMessageTemplate string) (string, bool) {
|
||||||
operationType, ok := opResponse["operationType"].(string)
|
operationType, ok := opResponse["operationType"].(string)
|
||||||
if !ok || operationType != "CREATE_DATABASE" {
|
if !ok || operationType != "CREATE_DATABASE" {
|
||||||
|
|||||||
@@ -28,6 +28,21 @@ import (
|
|||||||
|
|
||||||
const kind string = "cloud-gemini-data-analytics-query"
|
const kind string = "cloud-gemini-data-analytics-query"
|
||||||
|
|
||||||
|
// Guidance is the tool guidance string.
|
||||||
|
const Guidance = `Tool guidance:
|
||||||
|
Inputs:
|
||||||
|
1. query: A natural language formulation of a database query.
|
||||||
|
Outputs: (all optional)
|
||||||
|
1. disambiguation_question: Clarification questions or comments where the tool needs the users' input.
|
||||||
|
2. generated_query: The generated query for the user query.
|
||||||
|
3. intent_explanation: An explanation for why the tool produced ` + "`generated_query`" + `.
|
||||||
|
4. query_result: The result of executing ` + "`generated_query`" + `.
|
||||||
|
5. natural_language_answer: The natural language answer that summarizes the ` + "`query`" + ` and ` + "`query_result`" + `.
|
||||||
|
|
||||||
|
Usage guidance:
|
||||||
|
1. If ` + "`disambiguation_question`" + ` is produced, then solicit the needed inputs from the user and try the tool with a new ` + "`query`" + ` that has the needed clarification.
|
||||||
|
2. If ` + "`natural_language_answer`" + ` is produced, use ` + "`intent_explanation`" + ` and ` + "`generated_query`" + ` to see if you need to clarify any assumptions for the user.`
|
||||||
|
|
||||||
func init() {
|
func init() {
|
||||||
if !tools.Register(kind, newConfig) {
|
if !tools.Register(kind, newConfig) {
|
||||||
panic(fmt.Sprintf("tool kind %q already registered", kind))
|
panic(fmt.Sprintf("tool kind %q already registered", kind))
|
||||||
@@ -68,11 +83,18 @@ func (cfg Config) ToolConfigKind() string {
|
|||||||
|
|
||||||
func (cfg Config) Initialize(srcs map[string]sources.Source) (tools.Tool, error) {
|
func (cfg Config) Initialize(srcs map[string]sources.Source) (tools.Tool, error) {
|
||||||
// Define the parameters for the Gemini Data Analytics Query API
|
// Define the parameters for the Gemini Data Analytics Query API
|
||||||
// The prompt is the only input parameter.
|
// The query is the only input parameter.
|
||||||
allParameters := parameters.Parameters{
|
allParameters := parameters.Parameters{
|
||||||
parameters.NewStringParameterWithRequired("prompt", "The natural language question to ask.", true),
|
parameters.NewStringParameterWithRequired("query", "A natural language formulation of a database query.", true),
|
||||||
}
|
}
|
||||||
|
// The input and outputs are for tool guidance, usage guidance is for multi-turn interaction.
|
||||||
|
guidance := Guidance
|
||||||
|
|
||||||
|
if cfg.Description != "" {
|
||||||
|
cfg.Description += "\n\n" + guidance
|
||||||
|
} else {
|
||||||
|
cfg.Description = guidance
|
||||||
|
}
|
||||||
mcpManifest := tools.GetMcpManifest(cfg.Name, cfg.Description, cfg.AuthRequired, allParameters, nil)
|
mcpManifest := tools.GetMcpManifest(cfg.Name, cfg.Description, cfg.AuthRequired, allParameters, nil)
|
||||||
|
|
||||||
return Tool{
|
return Tool{
|
||||||
@@ -105,9 +127,9 @@ func (t Tool) Invoke(ctx context.Context, resourceMgr tools.SourceProvider, para
|
|||||||
}
|
}
|
||||||
|
|
||||||
paramsMap := params.AsMap()
|
paramsMap := params.AsMap()
|
||||||
prompt, ok := paramsMap["prompt"].(string)
|
query, ok := paramsMap["query"].(string)
|
||||||
if !ok {
|
if !ok {
|
||||||
return nil, fmt.Errorf("prompt parameter not found or not a string")
|
return nil, fmt.Errorf("query parameter not found or not a string")
|
||||||
}
|
}
|
||||||
|
|
||||||
// Parse the access token if provided
|
// Parse the access token if provided
|
||||||
@@ -125,7 +147,7 @@ func (t Tool) Invoke(ctx context.Context, resourceMgr tools.SourceProvider, para
|
|||||||
|
|
||||||
payload := &QueryDataRequest{
|
payload := &QueryDataRequest{
|
||||||
Parent: payloadParent,
|
Parent: payloadParent,
|
||||||
Prompt: prompt,
|
Prompt: query,
|
||||||
Context: t.Context,
|
Context: t.Context,
|
||||||
GenerationOptions: t.GenerationOptions,
|
GenerationOptions: t.GenerationOptions,
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -328,9 +328,9 @@ func TestInvoke(t *testing.T) {
|
|||||||
t.Fatalf("failed to initialize tool: %v", err)
|
t.Fatalf("failed to initialize tool: %v", err)
|
||||||
}
|
}
|
||||||
|
|
||||||
// Prepare parameters for invocation - ONLY prompt
|
// Prepare parameters for invocation - ONLY query
|
||||||
params := parameters.ParamValues{
|
params := parameters.ParamValues{
|
||||||
{Name: "prompt", Value: "How many accounts who have region in Prague are eligible for loans?"},
|
{Name: "query", Value: "How many accounts who have region in Prague are eligible for loans?"},
|
||||||
}
|
}
|
||||||
|
|
||||||
resourceMgr := resources.NewResourceManager(srcs, nil, nil, nil, nil, nil, nil)
|
resourceMgr := resources.NewResourceManager(srcs, nil, nil, nil, nil, nil, nil)
|
||||||
|
|||||||
@@ -0,0 +1,183 @@
|
|||||||
|
// 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 cloudsqlrestorebackup
|
||||||
|
|
||||||
|
import (
|
||||||
|
"context"
|
||||||
|
"fmt"
|
||||||
|
|
||||||
|
"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"
|
||||||
|
"google.golang.org/api/sqladmin/v1"
|
||||||
|
)
|
||||||
|
|
||||||
|
const kind string = "cloud-sql-restore-backup"
|
||||||
|
|
||||||
|
var _ tools.ToolConfig = Config{}
|
||||||
|
|
||||||
|
type compatibleSource interface {
|
||||||
|
GetDefaultProject() string
|
||||||
|
GetService(context.Context, string) (*sqladmin.Service, error)
|
||||||
|
UseClientAuthorization() bool
|
||||||
|
RestoreBackup(ctx context.Context, targetProject, targetInstance, sourceProject, sourceInstance, backupID, accessToken string) (any, error)
|
||||||
|
}
|
||||||
|
|
||||||
|
// Config defines the configuration for the restore-backup 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"`
|
||||||
|
}
|
||||||
|
|
||||||
|
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
|
||||||
|
}
|
||||||
|
|
||||||
|
// 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 targetProjectParam parameters.Parameter
|
||||||
|
if project != "" {
|
||||||
|
targetProjectParam = parameters.NewStringParameterWithDefault("target_project", project, "The GCP project ID. This is pre-configured; do not ask for it unless the user explicitly provides a different one.")
|
||||||
|
} else {
|
||||||
|
targetProjectParam = parameters.NewStringParameter("target_project", "The project ID")
|
||||||
|
}
|
||||||
|
|
||||||
|
allParameters := parameters.Parameters{
|
||||||
|
targetProjectParam,
|
||||||
|
parameters.NewStringParameter("target_instance", "Cloud SQL instance ID of the target instance. This does not include the project ID."),
|
||||||
|
parameters.NewStringParameter("backup_id", "Identifier of the backup being restored. Can be a BackupRun ID, backup name, or BackupDR backup name. Use the full backup ID as provided, do not try to parse it"),
|
||||||
|
parameters.NewStringParameterWithRequired("source_project", "GCP project ID of the instance that the backup belongs to. Only required if the backup_id is a BackupRun ID.", false),
|
||||||
|
parameters.NewStringParameterWithRequired("source_instance", "Cloud SQL instance ID of the instance that the backup belongs to. Only required if the backup_id is a BackupRun ID.", false),
|
||||||
|
}
|
||||||
|
paramManifest := allParameters.Manifest()
|
||||||
|
|
||||||
|
description := cfg.Description
|
||||||
|
if description == "" {
|
||||||
|
description = "Restores a backup on a Cloud SQL 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 restore-backup tool.
|
||||||
|
type Tool struct {
|
||||||
|
Config
|
||||||
|
AllParams parameters.Parameters `yaml:"allParams"`
|
||||||
|
manifest tools.Manifest
|
||||||
|
mcpManifest tools.McpManifest
|
||||||
|
}
|
||||||
|
|
||||||
|
func (t Tool) ToConfig() tools.ToolConfig {
|
||||||
|
return t.Config
|
||||||
|
}
|
||||||
|
|
||||||
|
func (t Tool) Invoke(ctx context.Context, resourceMgr tools.SourceProvider, params parameters.ParamValues, accessToken tools.AccessToken) (any, error) {
|
||||||
|
source, err := tools.GetCompatibleSource[compatibleSource](resourceMgr, t.Source, t.Name, t.Kind)
|
||||||
|
if err != nil {
|
||||||
|
return nil, err
|
||||||
|
}
|
||||||
|
paramsMap := params.AsMap()
|
||||||
|
|
||||||
|
targetProject, ok := paramsMap["target_project"].(string)
|
||||||
|
if !ok {
|
||||||
|
return nil, fmt.Errorf("error casting 'target_project' parameter: %v", paramsMap["target_project"])
|
||||||
|
}
|
||||||
|
targetInstance, ok := paramsMap["target_instance"].(string)
|
||||||
|
if !ok {
|
||||||
|
return nil, fmt.Errorf("error casting 'target_instance' parameter: %v", paramsMap["target_instance"])
|
||||||
|
}
|
||||||
|
backupID, ok := paramsMap["backup_id"].(string)
|
||||||
|
if !ok {
|
||||||
|
return nil, fmt.Errorf("error casting 'backup_id' parameter: %v", paramsMap["backup_id"])
|
||||||
|
}
|
||||||
|
sourceProject, _ := paramsMap["source_project"].(string)
|
||||||
|
sourceInstance, _ := paramsMap["source_instance"].(string)
|
||||||
|
|
||||||
|
return source.RestoreBackup(ctx, targetProject, targetInstance, sourceProject, sourceInstance, backupID, string(accessToken))
|
||||||
|
}
|
||||||
|
|
||||||
|
// 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
|
||||||
|
}
|
||||||
@@ -0,0 +1,71 @@
|
|||||||
|
// 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 cloudsqlrestorebackup_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/testutils"
|
||||||
|
"github.com/googleapis/genai-toolbox/internal/tools/cloudsql/cloudsqlrestorebackup"
|
||||||
|
)
|
||||||
|
|
||||||
|
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:
|
||||||
|
restore-backup-tool:
|
||||||
|
kind: cloud-sql-restore-backup
|
||||||
|
description: a test description
|
||||||
|
source: a-source
|
||||||
|
`,
|
||||||
|
want: server.ToolConfigs{
|
||||||
|
"restore-backup-tool": cloudsqlrestorebackup.Config{
|
||||||
|
Name: "restore-backup-tool",
|
||||||
|
Kind: "cloud-sql-restore-backup",
|
||||||
|
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"`
|
||||||
|
}{}
|
||||||
|
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)
|
||||||
|
}
|
||||||
|
})
|
||||||
|
}
|
||||||
|
}
|
||||||
88
server.json
88
server.json
@@ -31,6 +31,18 @@
|
|||||||
"default": "tools.yaml",
|
"default": "tools.yaml",
|
||||||
"isRequired": false
|
"isRequired": false
|
||||||
},
|
},
|
||||||
|
{
|
||||||
|
"type": "named",
|
||||||
|
"name": "--tools-files",
|
||||||
|
"description": "Multiple file paths specifying tool configurations. Files will be merged. Cannot be used with –-tools-file or –-tools-folder.",
|
||||||
|
"isRequired": false
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"type": "named",
|
||||||
|
"name": "--tools-folder",
|
||||||
|
"description": "Directory path containing YAML tool configuration files. All .yaml and .yml files in the directory will be loaded and merged. Cannot be used with –-tools-file or –-tools-files.",
|
||||||
|
"isRequired": false
|
||||||
|
},
|
||||||
{
|
{
|
||||||
"type": "named",
|
"type": "named",
|
||||||
"name": "--address",
|
"name": "--address",
|
||||||
@@ -70,6 +82,82 @@
|
|||||||
"warn",
|
"warn",
|
||||||
"error"
|
"error"
|
||||||
]
|
]
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"type": "named",
|
||||||
|
"name": "--logging-format",
|
||||||
|
"description": "Specify logging format to use.",
|
||||||
|
"default": "standard",
|
||||||
|
"choices": [
|
||||||
|
"standard",
|
||||||
|
"json"
|
||||||
|
]
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"type": "named",
|
||||||
|
"name": "--disable-reload",
|
||||||
|
"description": "Disables dynamic reloading of tools file.",
|
||||||
|
"format": "boolean",
|
||||||
|
"isRequired": false
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"type": "named",
|
||||||
|
"name": "--prebuilt",
|
||||||
|
"description": "Use a prebuilt tool configuration by source type.",
|
||||||
|
"isRequired": false
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"type": "named",
|
||||||
|
"name": "--stdio",
|
||||||
|
"description": "Listens via MCP STDIO instead of acting as a remote HTTP server.",
|
||||||
|
"format": "boolean",
|
||||||
|
"isRequired": false
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"type": "named",
|
||||||
|
"name": "--telemetry-gcp",
|
||||||
|
"description": "Enable exporting directly to Google Cloud Monitoring.",
|
||||||
|
"format": "boolean",
|
||||||
|
"isRequired": false
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"type": "named",
|
||||||
|
"name": "--telemetry-otlp",
|
||||||
|
"description": "Enable exporting using OpenTelemetry Protocol (OTLP) to the specified endpoint (e.g. 'http://127.0.0.1:4318').",
|
||||||
|
"isRequired": false
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"type": "named",
|
||||||
|
"name": "--telemetry-service-name",
|
||||||
|
"description": "Sets the value of the service.name resource attribute for telemetry data.",
|
||||||
|
"default": "toolbox",
|
||||||
|
"isRequired": false
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"type": "named",
|
||||||
|
"name": "--ui",
|
||||||
|
"description": "Launches the Toolbox UI web server.",
|
||||||
|
"format": "boolean",
|
||||||
|
"isRequired": false
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"type": "named",
|
||||||
|
"name": "--allowed-origins",
|
||||||
|
"description": "Specifies a list of origins permitted to access this server.",
|
||||||
|
"default": "*",
|
||||||
|
"isRequired": false
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"type": "named",
|
||||||
|
"name": "--help",
|
||||||
|
"description": "Show help for toolbox",
|
||||||
|
"isRequired": false
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"type": "named",
|
||||||
|
"name": "--version",
|
||||||
|
"description": "Show version for toolbox",
|
||||||
|
"isRequired": false
|
||||||
}
|
}
|
||||||
]
|
]
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -139,12 +139,12 @@ func TestCloudGdaToolEndpoints(t *testing.T) {
|
|||||||
// 1. RunToolGetTestByName
|
// 1. RunToolGetTestByName
|
||||||
expectedManifest := map[string]any{
|
expectedManifest := map[string]any{
|
||||||
toolName: map[string]any{
|
toolName: map[string]any{
|
||||||
"description": "Test GDA Tool",
|
"description": "Test GDA Tool\n\n" + cloudgda.Guidance,
|
||||||
"parameters": []any{
|
"parameters": []any{
|
||||||
map[string]any{
|
map[string]any{
|
||||||
"name": "prompt",
|
"name": "query",
|
||||||
"type": "string",
|
"type": "string",
|
||||||
"description": "The natural language question to ask.",
|
"description": "A natural language formulation of a database query.",
|
||||||
"required": true,
|
"required": true,
|
||||||
"authSources": []any{},
|
"authSources": []any{},
|
||||||
},
|
},
|
||||||
@@ -155,7 +155,7 @@ func TestCloudGdaToolEndpoints(t *testing.T) {
|
|||||||
tests.RunToolGetTestByName(t, toolName, expectedManifest)
|
tests.RunToolGetTestByName(t, toolName, expectedManifest)
|
||||||
|
|
||||||
// 2. RunToolInvokeParametersTest
|
// 2. RunToolInvokeParametersTest
|
||||||
params := []byte(`{"prompt": "test question"}`)
|
params := []byte(`{"query": "test question"}`)
|
||||||
tests.RunToolInvokeParametersTest(t, toolName, params, "\"queryResult\":\"SELECT * FROM table;\"")
|
tests.RunToolInvokeParametersTest(t, toolName, params, "\"queryResult\":\"SELECT * FROM table;\"")
|
||||||
|
|
||||||
// 3. Manual MCP Tool Call Test
|
// 3. Manual MCP Tool Call Test
|
||||||
@@ -172,7 +172,7 @@ func TestCloudGdaToolEndpoints(t *testing.T) {
|
|||||||
Params: map[string]any{
|
Params: map[string]any{
|
||||||
"name": toolName,
|
"name": toolName,
|
||||||
"arguments": map[string]any{
|
"arguments": map[string]any{
|
||||||
"prompt": "test question",
|
"query": "test question",
|
||||||
},
|
},
|
||||||
},
|
},
|
||||||
}
|
}
|
||||||
|
|||||||
267
tests/cloudsql/cloud_sql_restore_backup_test.go
Normal file
267
tests/cloudsql/cloud_sql_restore_backup_test.go
Normal file
@@ -0,0 +1,267 @@
|
|||||||
|
// 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 cloudsql
|
||||||
|
|
||||||
|
import (
|
||||||
|
"bytes"
|
||||||
|
"context"
|
||||||
|
"encoding/json"
|
||||||
|
"fmt"
|
||||||
|
"io"
|
||||||
|
"net/http"
|
||||||
|
"net/http/httptest"
|
||||||
|
"net/url"
|
||||||
|
"reflect"
|
||||||
|
"regexp"
|
||||||
|
"strings"
|
||||||
|
"testing"
|
||||||
|
"time"
|
||||||
|
|
||||||
|
"github.com/google/go-cmp/cmp"
|
||||||
|
"github.com/googleapis/genai-toolbox/internal/testutils"
|
||||||
|
"github.com/googleapis/genai-toolbox/tests"
|
||||||
|
"google.golang.org/api/sqladmin/v1"
|
||||||
|
)
|
||||||
|
|
||||||
|
var (
|
||||||
|
restoreBackupToolKind = "cloud-sql-restore-backup"
|
||||||
|
)
|
||||||
|
|
||||||
|
type restoreBackupTransport struct {
|
||||||
|
transport http.RoundTripper
|
||||||
|
url *url.URL
|
||||||
|
}
|
||||||
|
|
||||||
|
func (t *restoreBackupTransport) RoundTrip(req *http.Request) (*http.Response, error) {
|
||||||
|
if strings.HasPrefix(req.URL.String(), "https://sqladmin.googleapis.com") {
|
||||||
|
req.URL.Scheme = t.url.Scheme
|
||||||
|
req.URL.Host = t.url.Host
|
||||||
|
}
|
||||||
|
return t.transport.RoundTrip(req)
|
||||||
|
}
|
||||||
|
|
||||||
|
type masterRestoreBackupHandler struct {
|
||||||
|
t *testing.T
|
||||||
|
}
|
||||||
|
|
||||||
|
func (h *masterRestoreBackupHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
|
||||||
|
if !strings.Contains(r.UserAgent(), "genai-toolbox/") {
|
||||||
|
h.t.Errorf("User-Agent header not found")
|
||||||
|
}
|
||||||
|
var body sqladmin.InstancesRestoreBackupRequest
|
||||||
|
if err := json.NewDecoder(r.Body).Decode(&body); err != nil {
|
||||||
|
h.t.Fatalf("failed to decode request body: %v", err)
|
||||||
|
} else {
|
||||||
|
h.t.Logf("Received request body: %+v", body)
|
||||||
|
}
|
||||||
|
|
||||||
|
var expectedBody sqladmin.InstancesRestoreBackupRequest
|
||||||
|
var response any
|
||||||
|
var statusCode int
|
||||||
|
|
||||||
|
switch {
|
||||||
|
case body.Backup != "":
|
||||||
|
expectedBody = sqladmin.InstancesRestoreBackupRequest{
|
||||||
|
Backup: "projects/p1/backups/test-uid",
|
||||||
|
}
|
||||||
|
response = map[string]any{"name": "op1", "status": "PENDING"}
|
||||||
|
statusCode = http.StatusOK
|
||||||
|
case body.BackupdrBackup != "":
|
||||||
|
expectedBody = sqladmin.InstancesRestoreBackupRequest{
|
||||||
|
BackupdrBackup: "projects/p1/locations/us-central1/backupVaults/test-vault/dataSources/test-ds/backups/test-uid",
|
||||||
|
}
|
||||||
|
response = map[string]any{"name": "op1", "status": "PENDING"}
|
||||||
|
statusCode = http.StatusOK
|
||||||
|
case body.RestoreBackupContext != nil:
|
||||||
|
expectedBody = sqladmin.InstancesRestoreBackupRequest{
|
||||||
|
RestoreBackupContext: &sqladmin.RestoreBackupContext{
|
||||||
|
Project: "p1",
|
||||||
|
InstanceId: "source",
|
||||||
|
BackupRunId: 12345,
|
||||||
|
},
|
||||||
|
}
|
||||||
|
response = map[string]any{"name": "op1", "status": "PENDING"}
|
||||||
|
statusCode = http.StatusOK
|
||||||
|
default:
|
||||||
|
http.Error(w, fmt.Sprintf("unhandled restore request body: %v", body), http.StatusInternalServerError)
|
||||||
|
return
|
||||||
|
}
|
||||||
|
|
||||||
|
if diff := cmp.Diff(expectedBody, body); diff != "" {
|
||||||
|
h.t.Errorf("unexpected request body (-want +got):\n%s", diff)
|
||||||
|
}
|
||||||
|
|
||||||
|
w.Header().Set("Content-Type", "application/json")
|
||||||
|
w.WriteHeader(statusCode)
|
||||||
|
if err := json.NewEncoder(w).Encode(response); err != nil {
|
||||||
|
http.Error(w, err.Error(), http.StatusInternalServerError)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestRestoreBackupToolEndpoints(t *testing.T) {
|
||||||
|
ctx, cancel := context.WithTimeout(context.Background(), time.Minute)
|
||||||
|
defer cancel()
|
||||||
|
|
||||||
|
handler := &masterRestoreBackupHandler{t: t}
|
||||||
|
server := httptest.NewServer(handler)
|
||||||
|
defer server.Close()
|
||||||
|
|
||||||
|
serverURL, err := url.Parse(server.URL)
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("failed to parse server URL: %v", err)
|
||||||
|
}
|
||||||
|
|
||||||
|
originalTransport := http.DefaultClient.Transport
|
||||||
|
if originalTransport == nil {
|
||||||
|
originalTransport = http.DefaultTransport
|
||||||
|
}
|
||||||
|
http.DefaultClient.Transport = &restoreBackupTransport{
|
||||||
|
transport: originalTransport,
|
||||||
|
url: serverURL,
|
||||||
|
}
|
||||||
|
t.Cleanup(func() {
|
||||||
|
http.DefaultClient.Transport = originalTransport
|
||||||
|
})
|
||||||
|
|
||||||
|
var args []string
|
||||||
|
toolsFile := getRestoreBackupToolsConfig()
|
||||||
|
cmd, cleanup, err := tests.StartCmd(ctx, toolsFile, args...)
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("command initialization returned an error: %s", err)
|
||||||
|
}
|
||||||
|
defer cleanup()
|
||||||
|
|
||||||
|
waitCtx, cancel := context.WithTimeout(ctx, 30*time.Second)
|
||||||
|
defer cancel()
|
||||||
|
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)
|
||||||
|
}
|
||||||
|
|
||||||
|
tcs := []struct {
|
||||||
|
name string
|
||||||
|
toolName string
|
||||||
|
body string
|
||||||
|
want string
|
||||||
|
expectError bool
|
||||||
|
errorStatus int
|
||||||
|
}{
|
||||||
|
{
|
||||||
|
name: "successful restore with standard backup",
|
||||||
|
toolName: "restore-backup",
|
||||||
|
body: `{"target_project": "p1", "target_instance": "instance-standard", "backup_id": "12345", "source_project": "p1", "source_instance": "source"}`,
|
||||||
|
want: `{"name":"op1","status":"PENDING"}`,
|
||||||
|
},
|
||||||
|
{
|
||||||
|
name: "successful restore with project level backup",
|
||||||
|
toolName: "restore-backup",
|
||||||
|
body: `{"target_project": "p1", "target_instance": "instance-project-level", "backup_id": "projects/p1/backups/test-uid"}`,
|
||||||
|
want: `{"name":"op1","status":"PENDING"}`,
|
||||||
|
},
|
||||||
|
{
|
||||||
|
name: "successful restore with BackupDR backup",
|
||||||
|
toolName: "restore-backup",
|
||||||
|
body: `{"target_project": "p1", "target_instance": "instance-project-level", "backup_id": "projects/p1/locations/us-central1/backupVaults/test-vault/dataSources/test-ds/backups/test-uid"}`,
|
||||||
|
want: `{"name":"op1","status":"PENDING"}`,
|
||||||
|
},
|
||||||
|
{
|
||||||
|
name: "missing source instance info for standard backup",
|
||||||
|
toolName: "restore-backup",
|
||||||
|
body: `{"target_project": "p1", "target_instance": "instance-project-level", "backup_id": "12345"}`,
|
||||||
|
expectError: true,
|
||||||
|
errorStatus: http.StatusBadRequest,
|
||||||
|
},
|
||||||
|
{
|
||||||
|
name: "missing backup identifier",
|
||||||
|
toolName: "restore-backup",
|
||||||
|
body: `{"target_project": "p1", "target_instance": "instance-project-level"}`,
|
||||||
|
expectError: true,
|
||||||
|
errorStatus: http.StatusBadRequest,
|
||||||
|
},
|
||||||
|
{
|
||||||
|
name: "missing target instance info",
|
||||||
|
toolName: "restore-backup",
|
||||||
|
body: `{"backup_id": "12345"}`,
|
||||||
|
expectError: true,
|
||||||
|
errorStatus: http.StatusBadRequest,
|
||||||
|
},
|
||||||
|
}
|
||||||
|
|
||||||
|
for _, tc := range tcs {
|
||||||
|
tc := tc
|
||||||
|
t.Run(tc.name, func(t *testing.T) {
|
||||||
|
api := fmt.Sprintf("http://127.0.0.1:5000/api/tool/%s/invoke", tc.toolName)
|
||||||
|
req, err := http.NewRequest(http.MethodPost, api, bytes.NewBufferString(tc.body))
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("unable to create request: %s", err)
|
||||||
|
}
|
||||||
|
req.Header.Add("Content-type", "application/json")
|
||||||
|
resp, err := http.DefaultClient.Do(req)
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("unable to send request: %s", err)
|
||||||
|
}
|
||||||
|
defer resp.Body.Close()
|
||||||
|
|
||||||
|
if tc.expectError {
|
||||||
|
if resp.StatusCode != tc.errorStatus {
|
||||||
|
bodyBytes, _ := io.ReadAll(resp.Body)
|
||||||
|
t.Fatalf("expected status %d but got %d: %s", tc.errorStatus, resp.StatusCode, string(bodyBytes))
|
||||||
|
}
|
||||||
|
return
|
||||||
|
}
|
||||||
|
|
||||||
|
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))
|
||||||
|
}
|
||||||
|
|
||||||
|
var result struct {
|
||||||
|
Result string `json:"result"`
|
||||||
|
}
|
||||||
|
if err := json.NewDecoder(resp.Body).Decode(&result); err != nil {
|
||||||
|
t.Fatalf("failed to decode response: %v", err)
|
||||||
|
}
|
||||||
|
|
||||||
|
var got, want map[string]any
|
||||||
|
if err := json.Unmarshal([]byte(result.Result), &got); err != nil {
|
||||||
|
t.Fatalf("failed to unmarshal result: %v", err)
|
||||||
|
}
|
||||||
|
if err := json.Unmarshal([]byte(tc.want), &want); err != nil {
|
||||||
|
t.Fatalf("failed to unmarshal want: %v", err)
|
||||||
|
}
|
||||||
|
|
||||||
|
if !reflect.DeepEqual(got, want) {
|
||||||
|
t.Fatalf("unexpected result: got %+v, want %+v", got, want)
|
||||||
|
}
|
||||||
|
})
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func getRestoreBackupToolsConfig() map[string]any {
|
||||||
|
return map[string]any{
|
||||||
|
"sources": map[string]any{
|
||||||
|
"my-cloud-sql-source": map[string]any{
|
||||||
|
"kind": "cloud-sql-admin",
|
||||||
|
},
|
||||||
|
},
|
||||||
|
"tools": map[string]any{
|
||||||
|
"restore-backup": map[string]any{
|
||||||
|
"kind": restoreBackupToolKind,
|
||||||
|
"source": "my-cloud-sql-source",
|
||||||
|
},
|
||||||
|
},
|
||||||
|
}
|
||||||
|
}
|
||||||
Reference in New Issue
Block a user