Skip to content

Commit 677254e

Browse files
authored
feat(tools/alloydb-get-user): Add get-user tool for alloydb (#1436)
## Description --- This pull request introduces a new custom tool kind `alloydb-get-user` that retrieves detailed information for a specific AlloyDB user. ### Example Configuration ```yaml tools: get_user: kind: alloydb-get-user source: alloydb-admin-source description: Use this tool to retrieve detailed information for a specific AlloyDB user. ``` ### Example Request ``` curl -X POST http://127.0.0.1:5000/api/tool/get_user/invoke \ -H "Content-Type: application/json" \ -d '{ "projectId": "example-project", "locationId": "us-central1", "clusterId": "my-alloydb-cluster", "userId": "my-alloydb-user", }' ``` ## PR Checklist --- > Thank you for opening a Pull Request! Before submitting your PR, there are a > few things you can do to make sure it goes smoothly: - [x] Make sure you reviewed [CONTRIBUTING.md](https://github.com/googleapis/genai-toolbox/blob/main/CONTRIBUTING.md) - [ ] Make sure to open an issue as a [bug/issue](https://github.com/googleapis/genai-toolbox/issues/new/choose) before writing your code! That way we can discuss the change, evaluate designs, and agree on the general idea - [x] Ensure the tests and linter pass - [x] Code coverage does not decrease (if any source code was changed) - [x] Appropriate docs were updated (if necessary) - [x] Make sure to add `!` if this involve a breaking change 🛠️ Fixes #<issue_number_goes_here>
1 parent f2d9e3b commit 677254e

File tree

7 files changed

+412
-1
lines changed

7 files changed

+412
-1
lines changed

cmd/root.go

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -44,6 +44,7 @@ import (
4444
// Import tool packages for side effect of registration
4545
_ "github.com/googleapis/genai-toolbox/internal/tools/alloydb/alloydbgetcluster"
4646
_ "github.com/googleapis/genai-toolbox/internal/tools/alloydb/alloydbgetinstance"
47+
_ "github.com/googleapis/genai-toolbox/internal/tools/alloydb/alloydbgetuser"
4748
_ "github.com/googleapis/genai-toolbox/internal/tools/alloydb/alloydblistclusters"
4849
_ "github.com/googleapis/genai-toolbox/internal/tools/alloydb/alloydblistinstances"
4950
_ "github.com/googleapis/genai-toolbox/internal/tools/alloydb/alloydblistusers"

cmd/root_test.go

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1342,7 +1342,7 @@ func TestPrebuiltTools(t *testing.T) {
13421342
wantToolset: server.ToolsetConfigs{
13431343
"alloydb_postgres_admin_tools": tools.ToolsetConfig{
13441344
Name: "alloydb_postgres_admin_tools",
1345-
ToolNames: []string{"create_cluster", "wait_for_operation", "create_instance", "list_clusters", "list_instances", "list_users", "create_user", "get_cluster", "get_instance"},
1345+
ToolNames: []string{"create_cluster", "wait_for_operation", "create_instance", "list_clusters", "list_instances", "list_users", "create_user", "get_cluster", "get_instance", "get_user"},
13461346
},
13471347
},
13481348
},
Lines changed: 38 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,38 @@
1+
---
2+
title: "alloydb-get-user"
3+
type: docs
4+
weight: 1
5+
description: >
6+
The "alloydb-get-user" tool retrieves details for a specific AlloyDB user.
7+
aliases:
8+
- /resources/tools/alloydb-get-user
9+
---
10+
11+
## About
12+
13+
The `alloydb-get-user` tool retrieves detailed information for a single, specified AlloyDB user. It is compatible with [alloydb-admin](../../sources/alloydb-admin.md) source.
14+
15+
| Parameter | Type | Description | Required |
16+
| :--------- | :----- | :--------------------------------------------------------------------------------------- | :------- |
17+
| `project` | string | The GCP project ID to get user for. | Yes |
18+
| `location` | string | The location of the cluster (e.g., 'us-central1'). | Yes |
19+
| `cluster` | string | The ID of the cluster to retrieve the user from. | Yes |
20+
| `user` | string | The ID of the user to retrieve. | Yes |
21+
> **Note**
22+
> This tool authenticates using the credentials configured in its [alloydb-admin](../../sources/alloydb-admin.md) source which can be either [Application Default Credentials](https://cloud.google.com/docs/authentication/application-default-credentials) or client-side OAuth.
23+
24+
## Example
25+
26+
```yaml
27+
tools:
28+
get_specific_user:
29+
kind: alloydb-get-user
30+
source: my-alloydb-admin-source
31+
description: Use this tool to retrieve details for a specific AlloyDB user.
32+
```
33+
## Reference
34+
| **field** | **type** | **required** | **description** |
35+
|-------------|:------------------------------------------:|:------------:|--------------------------------------------------------------------------------------------------|
36+
| kind | string | true | Must be alloydb-get-user. | |
37+
| source | string | true | The name of an `alloydb-admin` source. |
38+
| description | string | true | Description of the tool that is passed to the agent. |

internal/prebuiltconfigs/tools/alloydb-postgres-admin.yaml

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -197,6 +197,10 @@ tools:
197197
kind: alloydb-get-instance
198198
source: alloydb-admin-source
199199
description: "Retrieves details of a specific AlloyDB instance."
200+
get_user:
201+
kind: alloydb-get-user
202+
source: alloydb-admin-source
203+
description: "Retrieves details of a specific AlloyDB user."
200204

201205
toolsets:
202206
alloydb_postgres_admin_tools:
@@ -209,3 +213,4 @@ toolsets:
209213
- create_user
210214
- get_cluster
211215
- get_instance
216+
- get_user
Lines changed: 171 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,171 @@
1+
2+
// Copyright 2025 Google LLC
3+
//
4+
// Licensed under the Apache License, Version 2.0 (the "License");
5+
// you may not use this file except in compliance with the License.
6+
// You may obtain a copy of the License at
7+
//
8+
// http://www.apache.org/licenses/LICENSE-2.0
9+
//
10+
// Unless required by applicable law or agreed to in writing, software
11+
// distributed under the License is distributed on an "AS IS" BASIS,
12+
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13+
// See the License for the specific language governing permissions and
14+
// limitations under the License.
15+
16+
package alloydbgetuser
17+
18+
import (
19+
"context"
20+
"fmt"
21+
22+
yaml "github.com/goccy/go-yaml"
23+
"github.com/googleapis/genai-toolbox/internal/sources"
24+
alloydbadmin "github.com/googleapis/genai-toolbox/internal/sources/alloydbadmin"
25+
"github.com/googleapis/genai-toolbox/internal/tools"
26+
)
27+
28+
const kind string = "alloydb-get-user"
29+
30+
func init() {
31+
if !tools.Register(kind, newConfig) {
32+
panic(fmt.Sprintf("tool kind %q already registered", kind))
33+
}
34+
}
35+
36+
func newConfig(ctx context.Context, name string, decoder *yaml.Decoder) (tools.ToolConfig, error) {
37+
actual := Config{Name: name}
38+
if err := decoder.DecodeContext(ctx, &actual); err != nil {
39+
return nil, err
40+
}
41+
return actual, nil
42+
}
43+
44+
// Configuration for the get-user tool.
45+
type Config struct {
46+
Name string `yaml:"name" validate:"required"`
47+
Kind string `yaml:"kind" validate:"required"`
48+
Source string `yaml:"source" validate:"required"`
49+
Description string `yaml:"description" validate:"required"`
50+
AuthRequired []string `yaml:"authRequired"`
51+
BaseURL string `yaml:"baseURL"`
52+
}
53+
54+
// validate interface
55+
var _ tools.ToolConfig = Config{}
56+
57+
// ToolConfigKind returns the kind of the tool.
58+
func (cfg Config) ToolConfigKind() string {
59+
return kind
60+
}
61+
62+
// Initialize initializes the tool from the configuration.
63+
func (cfg Config) Initialize(srcs map[string]sources.Source) (tools.Tool, error) {
64+
rawS, ok := srcs[cfg.Source]
65+
if !ok {
66+
return nil, fmt.Errorf("source %q not found", cfg.Source)
67+
}
68+
69+
s, ok := rawS.(*alloydbadmin.Source)
70+
if !ok {
71+
return nil, fmt.Errorf("invalid source for %q tool: source kind must be `%s`", kind, alloydbadmin.SourceKind)
72+
}
73+
74+
allParameters := tools.Parameters{
75+
tools.NewStringParameter("project", "The GCP project ID."),
76+
tools.NewStringParameter("location", "The location of the cluster (e.g., 'us-central1')."),
77+
tools.NewStringParameter("cluster", "The ID of the cluster."),
78+
tools.NewStringParameter("user", "The ID of the user."),
79+
}
80+
paramManifest := allParameters.Manifest()
81+
82+
inputSchema := allParameters.McpManifest()
83+
inputSchema.Required = []string{"project", "location", "cluster", "user"}
84+
85+
mcpManifest := tools.McpManifest{
86+
Name: cfg.Name,
87+
Description: cfg.Description,
88+
InputSchema: inputSchema,
89+
}
90+
91+
return Tool{
92+
Name: cfg.Name,
93+
Kind: kind,
94+
Source: s,
95+
AllParams: allParameters,
96+
manifest: tools.Manifest{Description: cfg.Description, Parameters: paramManifest},
97+
mcpManifest: mcpManifest,
98+
}, nil
99+
}
100+
101+
// Tool represents the get-user tool.
102+
type Tool struct {
103+
Name string `yaml:"name"`
104+
Kind string `yaml:"kind"`
105+
106+
Source *alloydbadmin.Source
107+
AllParams tools.Parameters
108+
109+
manifest tools.Manifest
110+
mcpManifest tools.McpManifest
111+
}
112+
113+
// Invoke executes the tool's logic.
114+
func (t Tool) Invoke(ctx context.Context, params tools.ParamValues, accessToken tools.AccessToken) (any, error) {
115+
paramsMap := params.AsMap()
116+
117+
project, ok := paramsMap["project"].(string)
118+
if !ok {
119+
return nil, fmt.Errorf("invalid or missing 'project' parameter; expected a string")
120+
}
121+
location, ok := paramsMap["location"].(string)
122+
if !ok {
123+
return nil, fmt.Errorf("invalid 'location' parameter; expected a string")
124+
}
125+
cluster, ok := paramsMap["cluster"].(string)
126+
if !ok {
127+
return nil, fmt.Errorf("invalid 'cluster' parameter; expected a string")
128+
}
129+
user, ok := paramsMap["user"].(string)
130+
if !ok {
131+
return nil, fmt.Errorf("invalid 'user' parameter; expected a string")
132+
}
133+
134+
service, err := t.Source.GetService(ctx, string(accessToken))
135+
if err != nil {
136+
return nil, err
137+
}
138+
139+
urlString := fmt.Sprintf("projects/%s/locations/%s/clusters/%s/users/%s", project, location, cluster, user)
140+
141+
resp, err := service.Projects.Locations.Clusters.Users.Get(urlString).Do()
142+
if err != nil {
143+
return nil, fmt.Errorf("error getting AlloyDB user: %w", err)
144+
}
145+
146+
return resp, nil
147+
}
148+
149+
// ParseParams parses the parameters for the tool.
150+
func (t Tool) ParseParams(data map[string]any, claims map[string]map[string]any) (tools.ParamValues, error) {
151+
return tools.ParseParams(t.AllParams, data, claims)
152+
}
153+
154+
// Manifest returns the tool's manifest.
155+
func (t Tool) Manifest() tools.Manifest {
156+
return t.manifest
157+
}
158+
159+
// McpManifest returns the tool's MCP manifest.
160+
func (t Tool) McpManifest() tools.McpManifest {
161+
return t.mcpManifest
162+
}
163+
164+
// Authorized checks if the tool is authorized.
165+
func (t Tool) Authorized(verifiedAuthServices []string) bool {
166+
return true
167+
}
168+
169+
func (t Tool) RequiresClientAuthorization() bool {
170+
return t.Source.UseClientAuthorization()
171+
}
Lines changed: 95 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,95 @@
1+
2+
// Copyright 2025 Google LLC
3+
//
4+
// Licensed under the Apache License, Version 2.0 (the "License");
5+
// you may not use this file except in compliance with the License.
6+
// You may obtain a copy of the License at
7+
//
8+
// http://www.apache.org/licenses/LICENSE-2.0
9+
//
10+
// Unless required by applicable law or agreed to in writing, software
11+
// distributed under the License is distributed on an "AS IS" BASIS,
12+
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13+
// See the License for the specific language governing permissions and
14+
// limitations under the License.
15+
16+
package alloydbgetuser_test
17+
18+
import (
19+
"testing"
20+
21+
yaml "github.com/goccy/go-yaml"
22+
"github.com/google/go-cmp/cmp"
23+
"github.com/googleapis/genai-toolbox/internal/server"
24+
"github.com/googleapis/genai-toolbox/internal/testutils"
25+
alloydbgetuser "github.com/googleapis/genai-toolbox/internal/tools/alloydb/alloydbgetuser"
26+
)
27+
28+
func TestParseFromYaml(t *testing.T) {
29+
ctx, err := testutils.ContextWithNewLogger()
30+
if err != nil {
31+
t.Fatalf("unexpected error: %s", err)
32+
}
33+
tcs := []struct {
34+
desc string
35+
in string
36+
want server.ToolConfigs
37+
}{
38+
{
39+
desc: "basic example",
40+
in: `
41+
tools:
42+
get-my-user:
43+
kind: alloydb-get-user
44+
source: my-alloydb-admin-source
45+
description: some description
46+
`,
47+
want: server.ToolConfigs{
48+
"get-my-user": alloydbgetuser.Config{
49+
Name: "get-my-user",
50+
Kind: "alloydb-get-user",
51+
Source: "my-alloydb-admin-source",
52+
Description: "some description",
53+
AuthRequired: []string{},
54+
},
55+
},
56+
},
57+
{
58+
desc: "with auth required",
59+
in: `
60+
tools:
61+
get-my-user-auth:
62+
kind: alloydb-get-user
63+
source: my-alloydb-admin-source
64+
description: some description
65+
authRequired:
66+
- my-google-auth-service
67+
- other-auth-service
68+
`,
69+
want: server.ToolConfigs{
70+
"get-my-user-auth": alloydbgetuser.Config{
71+
Name: "get-my-user-auth",
72+
Kind: "alloydb-get-user",
73+
Source: "my-alloydb-admin-source",
74+
Description: "some description",
75+
AuthRequired: []string{"my-google-auth-service", "other-auth-service"},
76+
},
77+
},
78+
},
79+
}
80+
for _, tc := range tcs {
81+
t.Run(tc.desc, func(t *testing.T) {
82+
got := struct {
83+
Tools server.ToolConfigs `yaml:"tools"`
84+
}{}
85+
// Parse contents
86+
err := yaml.UnmarshalContext(ctx, testutils.FormatYaml(tc.in), &got)
87+
if err != nil {
88+
t.Fatalf("unable to unmarshal: %s", err)
89+
}
90+
if diff := cmp.Diff(tc.want, got.Tools); diff != "" {
91+
t.Fatalf("incorrect parse: diff %v", diff)
92+
}
93+
})
94+
}
95+
}

0 commit comments

Comments
 (0)