diff --git a/cmd/root.go b/cmd/root.go index 6cc1970bf7..c5fd1ca0ff 100644 --- a/cmd/root.go +++ b/cmd/root.go @@ -203,7 +203,7 @@ func NewCommand(opts ...Option) *Command { flags.BoolVar(&cmd.cfg.TelemetryGCP, "telemetry-gcp", false, "Enable exporting directly to Google Cloud Monitoring.") flags.StringVar(&cmd.cfg.TelemetryOTLP, "telemetry-otlp", "", "Enable exporting using OpenTelemetry Protocol (OTLP) to the specified endpoint (e.g. 'http://127.0.0.1:4318')") flags.StringVar(&cmd.cfg.TelemetryServiceName, "telemetry-service-name", "toolbox", "Sets the value of the service.name resource attribute for telemetry data.") - flags.StringVar(&cmd.prebuiltConfig, "prebuilt", "", "Use a prebuilt tool configuration by source type. Cannot be used with --tools-file. Allowed: 'alloydb-postgres', 'bigquery', 'cloud-sql-mysql', 'cloud-sql-postgres', 'cloud-sql-mssql', 'firestore', 'postgres', 'spanner', 'spanner-postgres'.") + flags.StringVar(&cmd.prebuiltConfig, "prebuilt", "", "Use a prebuilt tool configuration by source type. Cannot be used with --tools-file. Allowed: 'alloydb-postgres-admin', alloydb-postgres', 'bigquery', 'cloud-sql-mysql', 'cloud-sql-postgres', 'cloud-sql-mssql', 'firestore', 'postgres', 'spanner', 'spanner-postgres'.") flags.BoolVar(&cmd.cfg.Stdio, "stdio", false, "Listens via MCP STDIO instead of acting as a remote HTTP server.") flags.BoolVar(&cmd.cfg.DisableReload, "disable-reload", false, "Disables dynamic reloading of tools file.") diff --git a/cmd/root_test.go b/cmd/root_test.go index b968270f79..ab3e8caba4 100644 --- a/cmd/root_test.go +++ b/cmd/root_test.go @@ -1161,6 +1161,7 @@ func TestSingleEdit(t *testing.T) { } func TestPrebuiltTools(t *testing.T) { + alloydb_admin_config, _ := prebuiltconfigs.Get("alloydb-postgres-admin") alloydb_config, _ := prebuiltconfigs.Get("alloydb-postgres") bigquery_config, _ := prebuiltconfigs.Get("bigquery") cloudsqlpg_config, _ := prebuiltconfigs.Get("cloud-sql-postgres") @@ -1180,6 +1181,16 @@ func TestPrebuiltTools(t *testing.T) { in []byte wantToolset server.ToolsetConfigs }{ + { + name: "alloydb postgres admin prebuilt tools", + in: alloydb_admin_config, + wantToolset: server.ToolsetConfigs{ + "alloydb-postgres-admin-tools": tools.ToolsetConfig{ + Name: "alloydb-postgres-admin-tools", + ToolNames: []string{"alloydb-create-cluster", "alloydb-operations-get", "alloydb-create-instance"}, + }, + }, + }, { name: "alloydb prebuilt tools", in: alloydb_config, diff --git a/internal/prebuiltconfigs/prebuiltconfigs_test.go b/internal/prebuiltconfigs/prebuiltconfigs_test.go index 74fdae2ed7..499a4b5646 100644 --- a/internal/prebuiltconfigs/prebuiltconfigs_test.go +++ b/internal/prebuiltconfigs/prebuiltconfigs_test.go @@ -23,6 +23,7 @@ import ( func TestLoadPrebuiltToolYAMLs(t *testing.T) { test_name := "test load prebuilt configs" expectedKeys := []string{ + "alloydb-postgres-admin", "alloydb-postgres", "bigquery", "cloud-sql-mssql", @@ -65,6 +66,7 @@ func TestLoadPrebuiltToolYAMLs(t *testing.T) { } func TestGetPrebuiltTool(t *testing.T) { + alloydb_admin_config, _ := Get("alloydb-postgres-admin") alloydb_config, _ := Get("alloydb-postgres") bigquery_config, _ := Get("bigquery") cloudsqlpg_config, _ := Get("cloud-sql-postgres") @@ -74,6 +76,9 @@ func TestGetPrebuiltTool(t *testing.T) { postgresconfig, _ := Get("postgres") spanner_config, _ := Get("spanner") spannerpg_config, _ := Get("spanner-postgres") + if len(alloydb_admin_config) <= 0 { + t.Fatalf("unexpected error: could not fetch alloydb prebuilt tools yaml") + } if len(alloydb_config) <= 0 { t.Fatalf("unexpected error: could not fetch alloydb prebuilt tools yaml") } diff --git a/internal/prebuiltconfigs/tools/alloydb-postgres-admin.yaml b/internal/prebuiltconfigs/tools/alloydb-postgres-admin.yaml new file mode 100644 index 0000000000..1e4878c11c --- /dev/null +++ b/internal/prebuiltconfigs/tools/alloydb-postgres-admin.yaml @@ -0,0 +1,123 @@ +sources: + alloydb-api-source: + kind: http + baseUrl: https://alloydb.googleapis.com + headers: + Authorization: Bearer ${API_KEY} + Content-Type: application/json +tools: + alloydb-create-cluster: + kind: http + source: alloydb-api-source + method: POST + path: /v1/projects/{{.projectId}}/locations/{{.locationId}}/clusters + description: "Create a new AlloyDB cluster. This is a long-running operation, but the API call returns quickly. This will return operation id to be used by get operations tool. Take all parameters from user in one go." + pathParams: + - name: projectId + type: string + description: "The dynamic path parameter for project id provided by user." + - name: locationId + type: string + description: "The dynamic path parameter for location. The default value is us-central1. If quota is exhausted then use other regions." + default: us-central1 + queryParams: + - name: clusterId + type: string + description: "A unique ID for the AlloyDB cluster." + requestBody: | + { + "networkConfig": { + "network": "projects/{{.project}}/global/networks/{{.network}}" + }, + "initialUser": { + "password": "{{.password}}", + "user": "{{.user}}" + } + } + bodyParams: + - name: project + type: string + description: "The dynamic path parameter for project id." + - name: network + type: string + description: "The name of the VPC network to connect the cluster to (e.g., 'default')." + default: default + - name: password + type: string + description: "A secure password for the initial 'postgres' user or the custom user provided." + - name: user + type: string + description: "The name for the initial superuser. If not provided, it defaults to 'postgres'. The initial database will always be named 'postgres'." + alloydb-operations-get: + kind: alloydb-wait-for-operation + source: alloydb-api-source + description: "This will poll on operations API until the operation is done. For checking operation status we need projectId, locationID and operationId. Once instance is created give follow up steps on how to use the variables to bring data plane MCP server up in local and remote setup." + delay: 1s + maxDelay: 4m + multiplier: 2 + maxRetries: 10 + alloydb-create-instance: + kind: http + source: alloydb-api-source + method: POST + path: /v1/projects/{{.projectId}}/locations/{{.locationId}}/clusters/{{.clusterId}}/instances + description: "Creates a new AlloyDB instance (PRIMARY, READ_POOL, or SECONDARY) within a cluster. This is a long-running operation. Take all parameters from user in one go. This will return operation id to be used by get operations tool." + pathParams: + - name: projectId + type: string + description: "The GCP project ID." + - name: locationId + type: string + description: "The location of the cluster (e.g., 'us-central1')." + default: us-central1 + - name: clusterId + type: string + description: "The ID of the cluster to create the instance in." + queryParams: + - name: instanceId + type: string + description: "A unique ID for the new AlloyDB instance." + requestBody: | + { + "instanceType": "{{.instanceType}}", + {{- if .displayName }} + "displayName": "{{.displayName}}", + {{- end }} + {{- if eq .instanceType "READ_POOL" }} + "readPoolConfig": { + "nodeCount": {{.nodeCount}} + }, + {{- end }} + {{- if eq .instanceType "SECONDARY" }} + "secondaryConfig": { + "primaryClusterName": "{{.primaryClusterName}}" + }, + {{- end }} + "networkConfig": { + "enablePublicIp": true + }, + "databaseFlags": { + "password.enforce_complexity": "on" + } + } + bodyParams: + - name: instanceType + type: string + description: "The type of instance to create. Required. Valid values are: PRIMARY, READ_POOL, SECONDARY." + - name: displayName + type: string + description: "An optional, user-friendly name for the instance." + - name: nodeCount + type: integer + description: "The number of nodes in the read pool. Required only if instanceType is READ_POOL. Default is 1." + default: 1 + - name: primaryClusterName + type: string + description: "The full resource name of the primary cluster for a SECONDARY instance. Required only if instanceType is SECONDARY. Otherwise don't ask" + default: "" + +toolsets: + alloydb-postgres-admin-tools: + - alloydb-create-cluster + - alloydb-operations-get + - alloydb-create-instance