mirror of
https://github.com/googleapis/genai-toolbox.git
synced 2026-01-08 15:14:00 -05:00
feat(tools/postgres-list-schemas): add new postgres-list-schemas tool (#1741)
## Description Add a read-only PostgreSQL custom list_schemas tool, that returns the schemas present in the database excluding system and temporary schemas. Returns the schema name, schema owner, grants, number of functions, number of tables, and number of views within each schema. <img width="1985" height="1043" alt="Screenshot 2025-10-20 at 7 45 45 PM" src="https://github.com/user-attachments/assets/8c4f0bb8-587c-489a-8795-efa79e92b06f" /> <img width="3372" height="1694" alt="3NpZG7W6h3XGsM7" src="https://github.com/user-attachments/assets/370b5440-cc48-4c4e-82ea-4fd508cbcf2b" /> > Should include a concise description of the changes (bug or feature), it's > impact, along with a summary of the solution ## 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) - [x] 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> Co-authored-by: Yuan Teoh <45984206+Yuan325@users.noreply.github.com>
This commit is contained in:
@@ -160,6 +160,7 @@ import (
|
||||
_ "github.com/googleapis/genai-toolbox/internal/tools/postgres/postgreslistactivequeries"
|
||||
_ "github.com/googleapis/genai-toolbox/internal/tools/postgres/postgreslistavailableextensions"
|
||||
_ "github.com/googleapis/genai-toolbox/internal/tools/postgres/postgreslistinstalledextensions"
|
||||
_ "github.com/googleapis/genai-toolbox/internal/tools/postgres/postgreslistschemas"
|
||||
_ "github.com/googleapis/genai-toolbox/internal/tools/postgres/postgreslisttables"
|
||||
_ "github.com/googleapis/genai-toolbox/internal/tools/postgres/postgreslistviews"
|
||||
_ "github.com/googleapis/genai-toolbox/internal/tools/postgres/postgressql"
|
||||
|
||||
@@ -1404,7 +1404,7 @@ func TestPrebuiltTools(t *testing.T) {
|
||||
wantToolset: server.ToolsetConfigs{
|
||||
"alloydb_postgres_database_tools": tools.ToolsetConfig{
|
||||
Name: "alloydb_postgres_database_tools",
|
||||
ToolNames: []string{"execute_sql", "list_tables", "list_active_queries", "list_available_extensions", "list_installed_extensions", "list_autovacuum_configurations", "list_memory_configurations", "list_top_bloated_tables", "list_replication_slots", "list_invalid_indexes", "get_query_plan", "list_views"},
|
||||
ToolNames: []string{"execute_sql", "list_tables", "list_active_queries", "list_available_extensions", "list_installed_extensions", "list_autovacuum_configurations", "list_memory_configurations", "list_top_bloated_tables", "list_replication_slots", "list_invalid_indexes", "get_query_plan", "list_views", "list_schemas"},
|
||||
},
|
||||
},
|
||||
},
|
||||
@@ -1434,7 +1434,7 @@ func TestPrebuiltTools(t *testing.T) {
|
||||
wantToolset: server.ToolsetConfigs{
|
||||
"cloud_sql_postgres_database_tools": tools.ToolsetConfig{
|
||||
Name: "cloud_sql_postgres_database_tools",
|
||||
ToolNames: []string{"execute_sql", "list_tables", "list_active_queries", "list_available_extensions", "list_installed_extensions", "list_autovacuum_configurations", "list_memory_configurations", "list_top_bloated_tables", "list_replication_slots", "list_invalid_indexes", "get_query_plan", "list_views"},
|
||||
ToolNames: []string{"execute_sql", "list_tables", "list_active_queries", "list_available_extensions", "list_installed_extensions", "list_autovacuum_configurations", "list_memory_configurations", "list_top_bloated_tables", "list_replication_slots", "list_invalid_indexes", "get_query_plan", "list_views", "list_schemas"},
|
||||
},
|
||||
},
|
||||
},
|
||||
@@ -1534,7 +1534,7 @@ func TestPrebuiltTools(t *testing.T) {
|
||||
wantToolset: server.ToolsetConfigs{
|
||||
"postgres_database_tools": tools.ToolsetConfig{
|
||||
Name: "postgres_database_tools",
|
||||
ToolNames: []string{"execute_sql", "list_tables", "list_active_queries", "list_available_extensions", "list_installed_extensions", "list_autovacuum_configurations", "list_memory_configurations", "list_top_bloated_tables", "list_replication_slots", "list_invalid_indexes", "get_query_plan", "list_views"},
|
||||
ToolNames: []string{"execute_sql", "list_tables", "list_active_queries", "list_available_extensions", "list_installed_extensions", "list_autovacuum_configurations", "list_memory_configurations", "list_top_bloated_tables", "list_replication_slots", "list_invalid_indexes", "get_query_plan", "list_views", "list_schemas"},
|
||||
},
|
||||
},
|
||||
},
|
||||
|
||||
@@ -45,6 +45,7 @@ details on how to connect your AI tools (IDEs) to databases via Toolbox and MCP.
|
||||
* `get_query_plan`: Generate the execution plan of a statement.
|
||||
* `list_views`: Lists views in the database from pg_views with a default
|
||||
limit of 50 rows. Returns schemaname, viewname and the ownername.
|
||||
* `list_schemas`: Lists schemas in the database.
|
||||
|
||||
## AlloyDB Postgres Admin
|
||||
|
||||
@@ -214,6 +215,7 @@ details on how to connect your AI tools (IDEs) to databases via Toolbox and MCP.
|
||||
* `get_query_plan`: Generate the execution plan of a statement.
|
||||
* `list_views`: Lists views in the database from pg_views with a default
|
||||
limit of 50 rows. Returns schemaname, viewname and the ownername.
|
||||
* `list_schemas`: Lists schemas in the database.
|
||||
|
||||
## Cloud SQL for PostgreSQL Observability
|
||||
|
||||
@@ -509,6 +511,7 @@ details on how to connect your AI tools (IDEs) to databases via Toolbox and MCP.
|
||||
* `get_query_plan`: Generate the execution plan of a statement.
|
||||
* `list_views`: Lists views in the database from pg_views with a default
|
||||
limit of 50 rows. Returns schemaname, viewname and the ownername.
|
||||
* `list_schemas`: Lists schemas in the database.
|
||||
|
||||
## Google Cloud Serverless for Apache Spark
|
||||
|
||||
|
||||
@@ -48,6 +48,9 @@ cluster][alloydb-free-trial].
|
||||
- [`postgres-list-views`](../tools/postgres/postgres-list-views.md)
|
||||
List views in an AlloyDB for PostgreSQL database.
|
||||
|
||||
- [`postgres-list-schemas`](../tools/postgres/postgres-list-schemas.md)
|
||||
List schemas in an AlloyDB for PostgreSQL database.
|
||||
|
||||
### Pre-built Configurations
|
||||
|
||||
- [AlloyDB using MCP](https://googleapis.github.io/genai-toolbox/how-to/connect-ide/alloydb_pg_mcp/)
|
||||
|
||||
@@ -44,6 +44,9 @@ to a database by following these instructions][csql-pg-quickstart].
|
||||
- [`postgres-list-views`](../tools/postgres/postgres-list-views.md)
|
||||
List views in a PostgreSQL database.
|
||||
|
||||
- [`postgres-list-schemas`](../tools/postgres/postgres-list-schemas.md)
|
||||
List schemas in a PostgreSQL database.
|
||||
|
||||
### Pre-built Configurations
|
||||
|
||||
- [Cloud SQL for Postgres using
|
||||
|
||||
@@ -38,6 +38,9 @@ reputation for reliability, feature robustness, and performance.
|
||||
- [`postgres-list-views`](../tools/postgres/postgres-list-views.md)
|
||||
List views in a PostgreSQL database.
|
||||
|
||||
- [`postgres-list-schemas`](../tools/postgres/postgres-list-views.md)
|
||||
List schemas in a PostgreSQL database.
|
||||
|
||||
### Pre-built Configurations
|
||||
|
||||
- [PostgreSQL using MCP](https://googleapis.github.io/genai-toolbox/how-to/connect-ide/postgres_mcp/)
|
||||
|
||||
55
docs/en/resources/tools/postgres/postgres-list-schemas.md
Normal file
55
docs/en/resources/tools/postgres/postgres-list-schemas.md
Normal file
@@ -0,0 +1,55 @@
|
||||
---
|
||||
title: "postgres-list-schemas"
|
||||
type: docs
|
||||
weight: 1
|
||||
description: >
|
||||
The "postgres-list-schemas" tool lists user-defined schemas in a database.
|
||||
aliases:
|
||||
- /resources/tools/postgres-list-schemas
|
||||
---
|
||||
|
||||
## About
|
||||
|
||||
The `postgres-list-schemas` tool retrieves information about schemas in a database excluding system
|
||||
and temporary schemas. It's compatible with any of the following sources:
|
||||
|
||||
- [alloydb-postgres](../../sources/alloydb-pg.md)
|
||||
- [cloud-sql-postgres](../../sources/cloud-sql-pg.md)
|
||||
- [postgres](../../sources/postgres.md)
|
||||
|
||||
`postgres-list-schemas` lists detailed information as JSON for each schema. The tool takes the following
|
||||
input parameters:
|
||||
|
||||
- `schema_name` (optional): A pattern to filter schema names using SQL LIKE operator.
|
||||
If omitted, all user-defined schemas are returned.
|
||||
|
||||
## Example
|
||||
|
||||
```yaml
|
||||
tools:
|
||||
list_schemas:
|
||||
kind: postgres-list-schemas
|
||||
source: postgres-source
|
||||
description: "Lists all schemas in the database ordered by schema name and excluding system and temporary schemas. It returns the schema name, schema owner, grants, number of functions, number of tables and number of views within each schema."
|
||||
```
|
||||
|
||||
The response is a json array with the following elements:
|
||||
|
||||
```json
|
||||
{
|
||||
"schema_name": "name of the schema.",
|
||||
"owner": "role that owns the schema",
|
||||
"grants": "A JSON object detailing the privileges (e.g., USAGE, CREATE) granted to different roles or PUBLIC on the schema.",
|
||||
"tables": "The total count of tables within the schema",
|
||||
"views": "The total count of views within the schema",
|
||||
"functions": "The total count of functions",
|
||||
}
|
||||
```
|
||||
|
||||
## Reference
|
||||
|
||||
| **field** | **type** | **required** | **description** |
|
||||
|-------------|:--------:|:------------:|----------------------------------------------------|
|
||||
| kind | string | true | Must be "postgres-list-schemas". |
|
||||
| source | string | true | Name of the source the SQL should execute on. |
|
||||
| description | string | false | Description of the tool that is passed to the LLM. |
|
||||
@@ -159,6 +159,11 @@ tools:
|
||||
list_views:
|
||||
kind: postgres-list-views
|
||||
source: alloydb-pg-source
|
||||
|
||||
list_schemas:
|
||||
kind: postgres-list-schemas
|
||||
source: alloydb-pg-source
|
||||
|
||||
toolsets:
|
||||
alloydb_postgres_database_tools:
|
||||
- execute_sql
|
||||
@@ -173,3 +178,4 @@ toolsets:
|
||||
- list_invalid_indexes
|
||||
- get_query_plan
|
||||
- list_views
|
||||
- list_schemas
|
||||
|
||||
@@ -159,6 +159,10 @@ tools:
|
||||
kind: postgres-list-views
|
||||
source: cloudsql-pg-source
|
||||
|
||||
list_schemas:
|
||||
kind: postgres-list-schemas
|
||||
source: cloudsql-pg-source
|
||||
|
||||
toolsets:
|
||||
cloud_sql_postgres_database_tools:
|
||||
- execute_sql
|
||||
@@ -173,3 +177,4 @@ toolsets:
|
||||
- list_invalid_indexes
|
||||
- get_query_plan
|
||||
- list_views
|
||||
- list_schemas
|
||||
|
||||
@@ -158,6 +158,10 @@ tools:
|
||||
kind: postgres-list-views
|
||||
source: postgresql-source
|
||||
|
||||
list_schemas:
|
||||
kind: postgres-list-schemas
|
||||
source: postgresql-source
|
||||
|
||||
toolsets:
|
||||
postgres_database_tools:
|
||||
- execute_sql
|
||||
@@ -172,3 +176,4 @@ toolsets:
|
||||
- list_invalid_indexes
|
||||
- get_query_plan
|
||||
- list_views
|
||||
- list_schemas
|
||||
|
||||
@@ -0,0 +1,219 @@
|
||||
// Copyright 2025 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 postgreslistschemas
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
|
||||
yaml "github.com/goccy/go-yaml"
|
||||
"github.com/googleapis/genai-toolbox/internal/sources"
|
||||
"github.com/googleapis/genai-toolbox/internal/sources/alloydbpg"
|
||||
"github.com/googleapis/genai-toolbox/internal/sources/cloudsqlpg"
|
||||
"github.com/googleapis/genai-toolbox/internal/sources/postgres"
|
||||
"github.com/googleapis/genai-toolbox/internal/tools"
|
||||
"github.com/jackc/pgx/v5/pgxpool"
|
||||
)
|
||||
|
||||
const kind string = "postgres-list-schemas"
|
||||
|
||||
const listSchemasStatement = `
|
||||
WITH
|
||||
schema_grants AS (
|
||||
SELECT schema_oid, jsonb_object_agg(grantee, privileges) AS grants
|
||||
FROM
|
||||
(
|
||||
SELECT
|
||||
n.oid AS schema_oid,
|
||||
CASE
|
||||
WHEN p.grantee = 0 THEN 'PUBLIC'
|
||||
ELSE pg_catalog.pg_get_userbyid(p.grantee)
|
||||
END
|
||||
AS grantee,
|
||||
jsonb_agg(p.privilege_type ORDER BY p.privilege_type) AS privileges
|
||||
FROM pg_catalog.pg_namespace n, aclexplode(n.nspacl) p
|
||||
WHERE n.nspacl IS NOT NULL
|
||||
GROUP BY n.oid, grantee
|
||||
) permissions_by_grantee
|
||||
GROUP BY schema_oid
|
||||
),
|
||||
all_schemas AS (
|
||||
SELECT
|
||||
n.nspname AS schema_name,
|
||||
pg_catalog.pg_get_userbyid(n.nspowner) AS owner,
|
||||
COALESCE(sg.grants, '{}'::jsonb) AS grants,
|
||||
(
|
||||
SELECT COUNT(*)
|
||||
FROM pg_catalog.pg_class c
|
||||
WHERE c.relnamespace = n.oid AND c.relkind = 'r'
|
||||
) AS tables,
|
||||
(
|
||||
SELECT COUNT(*)
|
||||
FROM pg_catalog.pg_class c
|
||||
WHERE c.relnamespace = n.oid AND c.relkind = 'v'
|
||||
) AS views,
|
||||
(SELECT COUNT(*) FROM pg_catalog.pg_proc p WHERE p.pronamespace = n.oid)
|
||||
AS functions
|
||||
FROM pg_catalog.pg_namespace n
|
||||
LEFT JOIN schema_grants sg
|
||||
ON n.oid = sg.schema_oid
|
||||
)
|
||||
SELECT *
|
||||
FROM all_schemas
|
||||
-- Exclude system schemas and temporary schemas created per session.
|
||||
WHERE
|
||||
schema_name NOT IN ('pg_catalog', 'information_schema', 'pg_toast')
|
||||
AND schema_name NOT LIKE 'pg_temp_%'
|
||||
AND ($1::text IS NULL OR schema_name LIKE '%' || $1::text || '%')
|
||||
ORDER BY schema_name;
|
||||
`
|
||||
|
||||
func init() {
|
||||
if !tools.Register(kind, newConfig) {
|
||||
panic(fmt.Sprintf("tool kind %q already registered", kind))
|
||||
}
|
||||
}
|
||||
|
||||
func newConfig(ctx context.Context, name string, decoder *yaml.Decoder) (tools.ToolConfig, error) {
|
||||
actual := Config{Name: name}
|
||||
if err := decoder.DecodeContext(ctx, &actual); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return actual, nil
|
||||
}
|
||||
|
||||
type compatibleSource interface {
|
||||
PostgresPool() *pgxpool.Pool
|
||||
}
|
||||
|
||||
// validate compatible sources are still compatible
|
||||
var _ compatibleSource = &alloydbpg.Source{}
|
||||
var _ compatibleSource = &cloudsqlpg.Source{}
|
||||
var _ compatibleSource = &postgres.Source{}
|
||||
|
||||
var compatibleSources = [...]string{alloydbpg.SourceKind, cloudsqlpg.SourceKind, postgres.SourceKind}
|
||||
|
||||
type Config struct {
|
||||
Name string `yaml:"name" validate:"required"`
|
||||
Kind string `yaml:"kind" validate:"required"`
|
||||
Source string `yaml:"source" validate:"required"`
|
||||
Description string `yaml:"description"`
|
||||
AuthRequired []string `yaml:"authRequired"`
|
||||
}
|
||||
|
||||
// validate interface
|
||||
var _ tools.ToolConfig = Config{}
|
||||
|
||||
func (cfg Config) ToolConfigKind() string {
|
||||
return kind
|
||||
}
|
||||
|
||||
func (cfg Config) Initialize(srcs map[string]sources.Source) (tools.Tool, error) {
|
||||
// verify source exists
|
||||
rawS, ok := srcs[cfg.Source]
|
||||
if !ok {
|
||||
return nil, fmt.Errorf("no source named %q configured", cfg.Source)
|
||||
}
|
||||
|
||||
// verify the source is compatible
|
||||
s, ok := rawS.(compatibleSource)
|
||||
if !ok {
|
||||
return nil, fmt.Errorf("invalid source for %q tool: source kind must be one of %q", kind, compatibleSources)
|
||||
}
|
||||
|
||||
allParameters := tools.Parameters{
|
||||
tools.NewStringParameterWithDefault("schema_name", "", "Optional: A specific schema name pattern to search for."),
|
||||
}
|
||||
description := cfg.Description
|
||||
if description == "" {
|
||||
description = "Lists all schemas in the database ordered by schema name and excluding system and temporary schemas. It returns the schema name, schema owner, grants, number of functions, number of tables and number of views within each schema."
|
||||
}
|
||||
mcpManifest := tools.GetMcpManifest(cfg.Name, description, cfg.AuthRequired, allParameters)
|
||||
|
||||
// finish tool setup
|
||||
return Tool{
|
||||
name: cfg.Name,
|
||||
kind: kind,
|
||||
authRequired: cfg.AuthRequired,
|
||||
allParams: allParameters,
|
||||
pool: s.PostgresPool(),
|
||||
manifest: tools.Manifest{
|
||||
Description: cfg.Description,
|
||||
Parameters: allParameters.Manifest(),
|
||||
AuthRequired: cfg.AuthRequired,
|
||||
},
|
||||
mcpManifest: mcpManifest,
|
||||
}, nil
|
||||
}
|
||||
|
||||
// validate interface
|
||||
var _ tools.Tool = Tool{}
|
||||
|
||||
type Tool struct {
|
||||
name string `yaml:"name"`
|
||||
kind string `yaml:"kind"`
|
||||
authRequired []string `yaml:"authRequired"`
|
||||
allParams tools.Parameters `yaml:"allParams"`
|
||||
pool *pgxpool.Pool
|
||||
manifest tools.Manifest
|
||||
mcpManifest tools.McpManifest
|
||||
}
|
||||
|
||||
func (t Tool) Invoke(ctx context.Context, params tools.ParamValues, accessToken tools.AccessToken) (any, error) {
|
||||
sliceParams := params.AsSlice()
|
||||
|
||||
results, err := t.pool.Query(ctx, listSchemasStatement, sliceParams...)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("unable to execute query: %w", err)
|
||||
}
|
||||
defer results.Close()
|
||||
|
||||
fields := results.FieldDescriptions()
|
||||
var out []map[string]any
|
||||
|
||||
for results.Next() {
|
||||
values, err := results.Values()
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("unable to parse row: %w", err)
|
||||
}
|
||||
rowMap := make(map[string]any)
|
||||
for i, field := range fields {
|
||||
rowMap[string(field.Name)] = values[i]
|
||||
}
|
||||
out = append(out, rowMap)
|
||||
}
|
||||
|
||||
return out, nil
|
||||
}
|
||||
|
||||
func (t Tool) ParseParams(data map[string]any, claims map[string]map[string]any) (tools.ParamValues, error) {
|
||||
return tools.ParseParams(t.allParams, data, claims)
|
||||
}
|
||||
|
||||
func (t Tool) Manifest() tools.Manifest {
|
||||
return t.manifest
|
||||
}
|
||||
|
||||
func (t Tool) McpManifest() tools.McpManifest {
|
||||
return t.mcpManifest
|
||||
}
|
||||
|
||||
func (t Tool) Authorized(verifiedAuthServices []string) bool {
|
||||
return tools.IsAuthorized(t.authRequired, verifiedAuthServices)
|
||||
}
|
||||
|
||||
func (t Tool) RequiresClientAuthorization() bool {
|
||||
return false
|
||||
}
|
||||
@@ -0,0 +1,95 @@
|
||||
// Copyright 2025 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 postgreslistschemas_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/postgres/postgreslistschemas"
|
||||
)
|
||||
|
||||
func TestParseFromYamlPostgreslistSchemas(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:
|
||||
example_tool:
|
||||
kind: postgres-list-schemas
|
||||
source: my-postgres-instance
|
||||
description: some description
|
||||
authRequired:
|
||||
- my-google-auth-service
|
||||
- other-auth-service
|
||||
`,
|
||||
want: server.ToolConfigs{
|
||||
"example_tool": postgreslistschemas.Config{
|
||||
Name: "example_tool",
|
||||
Kind: "postgres-list-schemas",
|
||||
Source: "my-postgres-instance",
|
||||
Description: "some description",
|
||||
AuthRequired: []string{"my-google-auth-service", "other-auth-service"},
|
||||
},
|
||||
},
|
||||
},
|
||||
{
|
||||
desc: "basic example",
|
||||
in: `
|
||||
tools:
|
||||
example_tool:
|
||||
kind: postgres-list-schemas
|
||||
source: my-postgres-instance
|
||||
description: some description
|
||||
`,
|
||||
want: server.ToolConfigs{
|
||||
"example_tool": postgreslistschemas.Config{
|
||||
Name: "example_tool",
|
||||
Kind: "postgres-list-schemas",
|
||||
Source: "my-postgres-instance",
|
||||
Description: "some description",
|
||||
AuthRequired: []string{},
|
||||
},
|
||||
},
|
||||
},
|
||||
}
|
||||
for _, tc := range tcs {
|
||||
t.Run(tc.desc, func(t *testing.T) {
|
||||
got := struct {
|
||||
Tools server.ToolConfigs `yaml:"tools"`
|
||||
}{}
|
||||
// Parse contents
|
||||
err := yaml.UnmarshalContext(ctx, testutils.FormatYaml(tc.in), &got)
|
||||
if err != nil {
|
||||
t.Fatalf("unable to unmarshal: %s", err)
|
||||
}
|
||||
if diff := cmp.Diff(tc.want, got.Tools); diff != "" {
|
||||
t.Fatalf("incorrect parse: diff %v", diff)
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
}
|
||||
@@ -150,6 +150,8 @@ func TestAlloyDBPgToolEndpoints(t *testing.T) {
|
||||
tmplSelectCombined, tmplSelectFilterCombined := tests.GetPostgresSQLTmplToolStatement()
|
||||
toolsFile = tests.AddTemplateParamConfig(t, toolsFile, AlloyDBPostgresToolKind, tmplSelectCombined, tmplSelectFilterCombined, "")
|
||||
|
||||
toolsFile = tests.AddPostgresPrebuiltConfig(t, toolsFile)
|
||||
|
||||
cmd, cleanup, err := tests.StartCmd(ctx, toolsFile, args...)
|
||||
if err != nil {
|
||||
t.Fatalf("command initialization returned an error: %s", err)
|
||||
@@ -173,6 +175,7 @@ func TestAlloyDBPgToolEndpoints(t *testing.T) {
|
||||
tests.RunMCPToolCallMethod(t, failInvocationWant, mcpSelect1Want)
|
||||
tests.RunExecuteSqlToolInvokeTest(t, createTableStatement, select1Want)
|
||||
tests.RunToolInvokeWithTemplateParameters(t, tableNameTemplateParam)
|
||||
tests.RunPostgresListSchemasTest(t, ctx, pool)
|
||||
}
|
||||
|
||||
// Test connection with different IP type
|
||||
|
||||
@@ -135,6 +135,7 @@ func TestCloudSQLPgSimpleToolEndpoints(t *testing.T) {
|
||||
tmplSelectCombined, tmplSelectFilterCombined := tests.GetPostgresSQLTmplToolStatement()
|
||||
toolsFile = tests.AddTemplateParamConfig(t, toolsFile, CloudSQLPostgresToolKind, tmplSelectCombined, tmplSelectFilterCombined, "")
|
||||
|
||||
toolsFile = tests.AddPostgresPrebuiltConfig(t, toolsFile)
|
||||
cmd, cleanup, err := tests.StartCmd(ctx, toolsFile, args...)
|
||||
if err != nil {
|
||||
t.Fatalf("command initialization returned an error: %s", err)
|
||||
@@ -158,6 +159,7 @@ func TestCloudSQLPgSimpleToolEndpoints(t *testing.T) {
|
||||
tests.RunMCPToolCallMethod(t, mcpMyFailToolWant, mcpSelect1Want)
|
||||
tests.RunExecuteSqlToolInvokeTest(t, createTableStatement, select1Want)
|
||||
tests.RunToolInvokeWithTemplateParameters(t, tableNameTemplateParam)
|
||||
tests.RunPostgresListSchemasTest(t, ctx, pool)
|
||||
}
|
||||
|
||||
// Test connection with different IP type
|
||||
|
||||
@@ -33,6 +33,10 @@ import (
|
||||
"github.com/jackc/pgx/v5/pgxpool"
|
||||
)
|
||||
|
||||
var (
|
||||
PostgresListSchemasToolKind = "postgres-list-schemas"
|
||||
)
|
||||
|
||||
// GetToolsConfig returns a mock tools config file
|
||||
func GetToolsConfig(sourceConfig map[string]any, toolKind, paramToolStatement, idParamToolStmt, nameParamToolStmt, arrayToolStatement, authToolStatement string) map[string]any {
|
||||
// Write config into a file and pass it to command
|
||||
@@ -190,6 +194,19 @@ func AddExecuteSqlConfig(t *testing.T, config map[string]any, toolKind string) m
|
||||
return config
|
||||
}
|
||||
|
||||
func AddPostgresPrebuiltConfig(t *testing.T, config map[string]any) map[string]any {
|
||||
tools, ok := config["tools"].(map[string]any)
|
||||
if !ok {
|
||||
t.Fatalf("unable to get tools from config")
|
||||
}
|
||||
tools["list_schemas"] = map[string]any{
|
||||
"kind": PostgresListSchemasToolKind,
|
||||
"source": "my-instance",
|
||||
}
|
||||
config["tools"] = tools
|
||||
return config
|
||||
}
|
||||
|
||||
func AddTemplateParamConfig(t *testing.T, config map[string]any, toolKind, tmplSelectCombined, tmplSelectFilterCombined string, tmplSelectAll string) map[string]any {
|
||||
toolsMap, ok := config["tools"].(map[string]any)
|
||||
if !ok {
|
||||
|
||||
@@ -168,6 +168,7 @@ func TestPostgres(t *testing.T) {
|
||||
toolsFile = tests.AddTemplateParamConfig(t, toolsFile, PostgresToolKind, tmplSelectCombined, tmplSelectFilterCombined, "")
|
||||
|
||||
toolsFile = addPrebuiltToolConfig(t, toolsFile)
|
||||
toolsFile = tests.AddPostgresPrebuiltConfig(t, toolsFile)
|
||||
|
||||
cmd, cleanup, err := tests.StartCmd(ctx, toolsFile, args...)
|
||||
if err != nil {
|
||||
@@ -196,6 +197,7 @@ func TestPostgres(t *testing.T) {
|
||||
// Run specific Postgres tool tests
|
||||
runPostgresListTablesTest(t, tableNameParam, tableNameAuth)
|
||||
runPostgresListViewsTest(t, ctx, pool, tableNameParam)
|
||||
tests.RunPostgresListSchemasTest(t, ctx, pool)
|
||||
runPostgresListActiveQueriesTest(t, ctx, pool)
|
||||
runPostgresListAvailableExtensionsTest(t)
|
||||
runPostgresListInstalledExtensionsTest(t)
|
||||
|
||||
@@ -31,8 +31,10 @@ import (
|
||||
|
||||
"github.com/google/go-cmp/cmp"
|
||||
"github.com/google/go-cmp/cmp/cmpopts"
|
||||
"github.com/google/uuid"
|
||||
"github.com/googleapis/genai-toolbox/internal/server/mcp/jsonrpc"
|
||||
"github.com/googleapis/genai-toolbox/internal/sources"
|
||||
"github.com/jackc/pgx/v5/pgxpool"
|
||||
)
|
||||
|
||||
// RunToolGet runs the tool get endpoint
|
||||
@@ -1139,6 +1141,94 @@ func RunMCPToolCallMethod(t *testing.T, myFailToolWant, select1Want string, opti
|
||||
}
|
||||
}
|
||||
|
||||
func setupPostgresSchemas(t *testing.T, ctx context.Context, pool *pgxpool.Pool, schemaName string) func() {
|
||||
createSchemaStmt := fmt.Sprintf("CREATE SCHEMA %s", schemaName)
|
||||
_, err := pool.Exec(ctx, createSchemaStmt)
|
||||
if err != nil {
|
||||
t.Fatalf("failed to create schema: %v", err)
|
||||
}
|
||||
|
||||
return func() {
|
||||
dropSchemaStmt := fmt.Sprintf("DROP SCHEMA %s", schemaName)
|
||||
_, err := pool.Exec(ctx, dropSchemaStmt)
|
||||
if err != nil {
|
||||
t.Fatalf("failed to drop schema: %v", err)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func RunPostgresListSchemasTest(t *testing.T, ctx context.Context, pool *pgxpool.Pool) {
|
||||
schemaName := "test_schema_" + strings.ReplaceAll(uuid.New().String(), "-", "")
|
||||
cleanup := setupPostgresSchemas(t, ctx, pool, schemaName)
|
||||
defer cleanup()
|
||||
|
||||
wantSchema := map[string]any{"functions": float64(0), "grants": map[string]any{}, "owner": "postgres", "schema_name": schemaName, "tables": float64(0), "views": float64(0)}
|
||||
|
||||
invokeTcs := []struct {
|
||||
name string
|
||||
requestBody io.Reader
|
||||
wantStatusCode int
|
||||
want []map[string]any
|
||||
}{
|
||||
{
|
||||
name: "invoke list_schemas with schema_name",
|
||||
requestBody: bytes.NewBuffer([]byte(fmt.Sprintf(`{"schema_name": "%s"}`, schemaName))),
|
||||
wantStatusCode: http.StatusOK,
|
||||
want: []map[string]any{wantSchema},
|
||||
},
|
||||
{
|
||||
name: "invoke list_schemas with non-existent schema",
|
||||
requestBody: bytes.NewBuffer([]byte(`{"schema_name": "non_existent_schema"}`)),
|
||||
wantStatusCode: http.StatusOK,
|
||||
want: nil,
|
||||
},
|
||||
}
|
||||
for _, tc := range invokeTcs {
|
||||
t.Run(tc.name, func(t *testing.T) {
|
||||
const api = "http://127.0.0.1:5000/api/tool/list_schemas/invoke"
|
||||
req, err := http.NewRequest(http.MethodPost, api, tc.requestBody)
|
||||
if err != nil {
|
||||
t.Fatalf("unable to create request: %v", err)
|
||||
}
|
||||
req.Header.Add("Content-type", "application/json")
|
||||
resp, err := http.DefaultClient.Do(req)
|
||||
if err != nil {
|
||||
t.Fatalf("unable to send request: %v", err)
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
|
||||
if resp.StatusCode != tc.wantStatusCode {
|
||||
body, _ := io.ReadAll(resp.Body)
|
||||
t.Fatalf("wrong status code: got %d, want %d, body: %s", resp.StatusCode, tc.wantStatusCode, string(body))
|
||||
}
|
||||
if tc.wantStatusCode != http.StatusOK {
|
||||
return
|
||||
}
|
||||
|
||||
var bodyWrapper struct {
|
||||
Result json.RawMessage `json:"result"`
|
||||
}
|
||||
if err := json.NewDecoder(resp.Body).Decode(&bodyWrapper); err != nil {
|
||||
t.Fatalf("error decoding response wrapper: %v", err)
|
||||
}
|
||||
|
||||
var resultString string
|
||||
if err := json.Unmarshal(bodyWrapper.Result, &resultString); err != nil {
|
||||
resultString = string(bodyWrapper.Result)
|
||||
}
|
||||
|
||||
var got []map[string]any
|
||||
if err := json.Unmarshal([]byte(resultString), &got); err != nil {
|
||||
t.Fatalf("failed to unmarshal nested result string: %v", err)
|
||||
}
|
||||
|
||||
if diff := cmp.Diff(tc.want, got); diff != "" {
|
||||
t.Errorf("Unexpected result (-want +got):\n%s", diff)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
// RunMySQLListTablesTest run tests against the mysql-list-tables tool
|
||||
func RunMySQLListTablesTest(t *testing.T, databaseName, tableNameParam, tableNameAuth string) {
|
||||
type tableInfo struct {
|
||||
|
||||
Reference in New Issue
Block a user