feat(tools): add list_indexes, list_sequences tools for postgres (#1765)

## Description

Adds the following tools for Postgres:
(1) list_indexes:  Lists available user indexes in a given database. 
(2) list_sequences: Lists all the sequences in the database ordered by
sequence name.

list_indexes:
<img width="1708" height="816" alt="Screenshot 2025-10-21 at 10 24
30 PM"
src="https://github.com/user-attachments/assets/d1888f53-6013-4beb-b0dc-b94f3d66135e"
/>

<img width="1914" height="1017" alt="Screenshot 2025-10-22 at 8 23
16 AM"
src="https://github.com/user-attachments/assets/7c185ef4-4550-4bb2-8051-0eca8f40dc7b"
/>

list_sequences:

<img width="2539" height="902" alt="Screenshot 2025-10-22 at 8 04 25 AM"
src="https://github.com/user-attachments/assets/f22cfaad-412b-4df2-99f3-ee813f538ff7"
/>

<img width="1082" height="610" alt="Screenshot 2025-10-22 at 8 06 07 AM"
src="https://github.com/user-attachments/assets/102ed3ca-33f6-409f-9ba1-363c9e15d5d3"
/>



> 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 #1738
This commit is contained in:
Srividya Reddy
2025-11-14 02:47:03 +05:30
committed by GitHub
parent 22b5aca395
commit 897c63dcea
20 changed files with 1052 additions and 4 deletions

View File

@@ -181,8 +181,10 @@ import (
_ "github.com/googleapis/genai-toolbox/internal/tools/postgres/postgresexecutesql"
_ "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/postgreslistindexes"
_ "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/postgreslistsequences"
_ "github.com/googleapis/genai-toolbox/internal/tools/postgres/postgreslisttables"
_ "github.com/googleapis/genai-toolbox/internal/tools/postgres/postgreslisttriggers"
_ "github.com/googleapis/genai-toolbox/internal/tools/postgres/postgreslistviews"

View File

@@ -1477,7 +1477,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", "list_schemas", "database_overview", "list_triggers"},
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", "database_overview", "list_triggers", "list_indexes", "list_sequences"},
},
},
},
@@ -1507,7 +1507,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", "list_schemas", "database_overview", "list_triggers"},
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", "database_overview", "list_triggers", "list_indexes", "list_sequences"},
},
},
},
@@ -1607,7 +1607,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", "list_schemas", "database_overview", "list_triggers"},
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", "database_overview", "list_triggers", "list_indexes", "list_sequences"},
},
},
},

View File

@@ -48,6 +48,8 @@ details on how to connect your AI tools (IDEs) to databases via Toolbox and MCP.
* `list_schemas`: Lists schemas in the database.
* `database_overview`: Fetches the current state of the PostgreSQL server.
* `list_triggers`: Lists triggers in the database.
* `list_indexes`: List available user indexes in a PostgreSQL database.
* `list_sequences`: List sequences in a PostgreSQL database.
## AlloyDB Postgres Admin
@@ -220,6 +222,8 @@ details on how to connect your AI tools (IDEs) to databases via Toolbox and MCP.
* `list_schemas`: Lists schemas in the database.
* `database_overview`: Fetches the current state of the PostgreSQL server.
* `list_triggers`: Lists triggers in the database.
* `list_indexes`: List available user indexes in a PostgreSQL database.
* `list_sequences`: List sequences in a PostgreSQL database.
## Cloud SQL for PostgreSQL Observability
@@ -519,6 +523,8 @@ details on how to connect your AI tools (IDEs) to databases via Toolbox and MCP.
* `list_schemas`: Lists schemas in the database.
* `database_overview`: Fetches the current state of the PostgreSQL server.
* `list_triggers`: Lists triggers in the database.
* `list_indexes`: List available user indexes in a PostgreSQL database.
* `list_sequences`: List sequences in a PostgreSQL database.
## Google Cloud Serverless for Apache Spark

View File

@@ -57,6 +57,12 @@ cluster][alloydb-free-trial].
- [`postgres-list-triggers`](../tools/postgres/postgres-list-triggers.md)
List triggers in an AlloyDB for PostgreSQL database.
- [`postgres-list-indexes`](../tools/postgres/postgres-list-indexes.md)
List available user indexes in a PostgreSQL database.
- [`postgres-list-sequences`](../tools/postgres/postgres-list-sequences.md)
List sequences in a PostgreSQL database.
### Pre-built Configurations
- [AlloyDB using MCP](https://googleapis.github.io/genai-toolbox/how-to/connect-ide/alloydb_pg_mcp/)

View File

@@ -53,6 +53,12 @@ to a database by following these instructions][csql-pg-quickstart].
- [`postgres-list-triggers`](../tools/postgres/postgres-list-triggers.md)
List triggers in a PostgreSQL database.
- [`postgres-list-indexes`](../tools/postgres/postgres-list-indexes.md)
List available user indexes in a PostgreSQL database.
- [`postgres-list-sequences`](../tools/postgres/postgres-list-sequences.md)
List sequences in a PostgreSQL database.
### Pre-built Configurations
- [Cloud SQL for Postgres using

View File

@@ -47,6 +47,12 @@ reputation for reliability, feature robustness, and performance.
- [`postgres-list-triggers`](../tools/postgres/postgres-list-triggers.md)
List triggers in a PostgreSQL database.
- [`postgres-list-indexes`](../tools/postgres/postgres-list-indexes.md)
List available user indexes in a PostgreSQL database.
- [`postgres-list-sequences`](../tools/postgres/postgres-list-sequences.md)
List sequences in a PostgreSQL database.
### Pre-built Configurations
- [PostgreSQL using MCP](https://googleapis.github.io/genai-toolbox/how-to/connect-ide/postgres_mcp/)

View File

@@ -0,0 +1,67 @@
---
title: "postgres-list-indexes"
type: docs
weight: 1
description: >
The "postgres-list-indexes" tool lists indexes in a Postgres database.
aliases:
- /resources/tools/postgres-list-indexes
---
## About
The `postgres-list-indexes` tool lists available user indexes in the database excluding those in `pg_catalog` and `information_schema`. 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-indexes` lists detailed information as JSON for indexes. The tool takes the following input parameters:
- `table_name` (optional): A text to filter results by table name. The input is used within a LIKE clause. Default: `""`
- `index_name` (optional): A text to filter results by index name. The input is used within a LIKE clause. Default: `""`
- `schema_name` (optional): A text to filter results by schema name. The input is used within a LIKE clause. Default: `""`
- `limit` (optional): The maximum number of rows to return. Default: `50`.
## Example
```yaml
tools:
list_indexes:
kind: postgres-list-indexes
source: postgres-source
description: |
Lists available user indexes in the database, excluding system schemas (pg_catalog,
information_schema). For each index, the following properties are returned:
schema name, table name, index name, index type (access method), a boolean
indicating if it's a unique index, a boolean indicating if it's for a primary key,
the index definition, index size in bytes, the number of index scans, the number of
index tuples read, the number of table tuples fetched via index scans, and a boolean
indicating if the index has been used at least once.
```
The response is a json array with the following elements:
```json
{
"schema_name": "schema name",
"table_name": "table name",
"index_name": "index name",
"index_type": "index access method (e.g btree, hash, gin)",
"is_unique": "boolean indicating if the index is unique",
"is_primary": "boolean indicating if the index is for a primary key",
"index_definition": "index definition statement",
"index_size_bytes": "index size in bytes",
"index_scans": "Number of index scans initiated on this index",
"tuples_read": "Number of index entries returned by scans on this index",
"tuples_fetched": "Number of live table rows fetched by simple index scans using this index",
"is_used": "boolean indicating if the index has been scanned at least once"
}
```
## Reference
| **field** | **type** | **required** | **description** |
|-------------|:--------:|:-------------:|------------------------------------------------------|
| kind | string | true | Must be "postgres-list-indexes". |
| source | string | true | Name of the source the SQL should execute on. |
| description | string | false | Description of the tool that is passed to the agent. |

View File

@@ -0,0 +1,61 @@
---
title: "postgres-list-sequences"
type: docs
weight: 1
description: >
The "postgres-list-sequences" tool lists sequences in a Postgres database.
aliases:
- /resources/tools/postgres-list-sequences
---
## About
The `postgres-list-sequences` tool retrieves information about sequences in a Postgres database. 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-sequences` lists detailed information as JSON for all sequences. The tool takes the following input parameters:
- `sequencename` (optional): A text to filter results by sequence name. The input is used within a LIKE clause. Default: `""`
- `schemaname` (optional): A text to filter results by schema name. The input is used within a LIKE clause. Default: `""`
- `limit` (optional): The maximum number of rows to return. Default: `50`.
## Example
```yaml
tools:
list_indexes:
kind: postgres-list-sequences
source: postgres-source
description: |
Lists all the sequences in the database ordered by sequence name.
Returns sequence name, schema name, sequence owner, data type of the
sequence, starting value, minimum value, maximum value of the sequence,
the value by which the sequence is incremented, and the last value
generated by generated by the sequence in the current session.
```
The response is a json array with the following elements:
```json
{
"sequencename": "sequence name",
"schemaname": "schema name",
"sequenceowner": "owner of the sequence",
"data_type": "data type of the sequence",
"start_value": "starting value of the sequence",
"min_value": "minimum value of the sequence",
"max_value": "maximum value of the sequence",
"increment_by": "increment value of the sequence",
"last_value": "last value of the sequence"
}
```
## Reference
| **field** | **type** | **required** | **description** |
|-------------|:--------:|:-------------:|------------------------------------------------------|
| kind | string | true | Must be "postgres-list-sequences". |
| source | string | true | Name of the source the SQL should execute on. |
| description | string | false | Description of the tool that is passed to the agent. |

View File

@@ -163,6 +163,14 @@ tools:
list_schemas:
kind: postgres-list-schemas
source: alloydb-pg-source
list_indexes:
kind: postgres-list-indexes
source: alloydb-pg-source
list_sequences:
kind: postgres-list-sequences
source: alloydb-pg-source
database_overview:
kind: postgres-database-overview
@@ -189,3 +197,5 @@ toolsets:
- list_schemas
- database_overview
- list_triggers
- list_indexes
- list_sequences

View File

@@ -166,10 +166,18 @@ tools:
database_overview:
kind: postgres-database-overview
source: cloudsql-pg-source
list_triggers:
kind: postgres-list-triggers
source: cloudsql-pg-source
list_indexes:
kind: postgres-list-indexes
source: cloudsql-pg-source
list_sequences:
kind: postgres-list-sequences
source: cloudsql-pg-source
toolsets:
cloud_sql_postgres_database_tools:
@@ -188,3 +196,5 @@ toolsets:
- list_schemas
- database_overview
- list_triggers
- list_indexes
- list_sequences

View File

@@ -168,6 +168,14 @@ tools:
list_triggers:
kind: postgres-list-triggers
source: postgresql-source
list_indexes:
kind: postgres-list-indexes
source: postgresql-source
list_sequences:
kind: postgres-list-sequences
source: postgresql-source
toolsets:
@@ -187,3 +195,5 @@ toolsets:
- list_schemas
- database_overview
- list_triggers
- list_indexes
- list_sequences

View File

@@ -0,0 +1,215 @@
// 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 postgreslistindexes
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-indexes"
const listIndexesStatement = `
WITH IndexDetails AS (
SELECT
s.schemaname AS schema_name,
t.relname AS table_name,
i.relname AS index_name,
am.amname AS index_type,
ix.indisunique AS is_unique,
ix.indisprimary AS is_primary,
pg_get_indexdef(i.oid) AS index_definition,
pg_relation_size(i.oid) AS index_size_bytes,
s.idx_scan AS index_scans,
s.idx_tup_read AS tuples_read,
s.idx_tup_fetch AS tuples_fetched,
CASE
WHEN s.idx_scan > 0 THEN true
ELSE false
END AS is_used
FROM pg_catalog.pg_class t
JOIN pg_catalog.pg_index ix
ON t.oid = ix.indrelid
JOIN pg_catalog.pg_class i
ON i.oid = ix.indexrelid
JOIN pg_catalog.pg_am am
ON i.relam = am.oid
JOIN pg_catalog.pg_stat_all_indexes s
ON i.oid = s.indexrelid
WHERE
t.relkind = 'r'
AND s.schemaname NOT IN ('pg_catalog', 'information_schema')
)
SELECT *
FROM IndexDetails
WHERE
($1::text IS NULL OR schema_name LIKE '%' || $1 || '%')
AND ($2::text IS NULL OR table_name LIKE '%' || $2 || '%')
AND ($3::text IS NULL OR index_name LIKE '%' || $3 || '%')
ORDER BY
schema_name,
table_name,
index_name
LIMIT COALESCE($4::int, 50);
`
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 text to filter results by schema name. The input is used within a LIKE clause."),
tools.NewStringParameterWithDefault("table_name", "", "Optional: a text to filter results by table name. The input is used within a LIKE clause."),
tools.NewStringParameterWithDefault("index_name", "", "Optional: a text to filter results by index name. The input is used within a LIKE clause."),
tools.NewIntParameterWithDefault("limit", 50, "Optional: The maximum number of rows to return. Default is 50"),
}
description := cfg.Description
if description == "" {
description = "Lists available user indexes in the database, excluding system schemas (pg_catalog, information_schema). For each index, the following properties are returned: schema name, table name, index name, index type (access method), a boolean indicating if it's a unique index, a boolean indicating if it's for a primary key, the index definition, index size in bytes, the number of index scans, the number of index tuples read, the number of table tuples fetched via index scans, and a boolean indicating if the index has been used at least once."
}
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, listIndexesStatement, 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
}

View File

@@ -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 postgreslistindexes_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/postgreslistindexes"
)
func TestParseFromYamlPostgresListIndexes(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-indexes
source: my-postgres-instance
description: some description
authRequired:
- my-google-auth-service
- other-auth-service
`,
want: server.ToolConfigs{
"example_tool": postgreslistindexes.Config{
Name: "example_tool",
Kind: "postgres-list-indexes",
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-indexes
source: my-postgres-instance
description: some description
`,
want: server.ToolConfigs{
"example_tool": postgreslistindexes.Config{
Name: "example_tool",
Kind: "postgres-list-indexes",
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)
}
})
}
}

View File

@@ -0,0 +1,190 @@
// 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 postgreslistsequences
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-sequences"
const listSequencesStatement = `
SELECT
sequencename,
schemaname,
sequenceowner,
data_type,
start_value,
min_value,
max_value,
increment_by,
last_value
FROM pg_sequences
WHERE
($1::text IS NULL OR schemaname LIKE '%' || $1 || '%')
AND ($2::text IS NULL OR sequencename LIKE '%' || $2 || '%')
ORDER BY schemaname, sequencename
LIMIT COALESCE($3::int, 50);
`
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("schemaname", "", "Optional: A specific schema name pattern to search for."),
tools.NewStringParameterWithDefault("sequencename", "", "Optional: A specific sequence name pattern to search for."),
tools.NewIntParameterWithDefault("limit", 50, "Optional: The maximum number of rows to return. Default is 50"),
}
description := cfg.Description
if description == "" {
description = "Lists sequences in the database. Returns sequence name, schema name, sequence owner, data type of the sequence, starting value, minimum value, maximum value of the sequence, the value by which the sequence is incremented, and the last value generated by the sequence in the current session"
}
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, listSequencesStatement, 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
}

View File

@@ -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 postgreslistsequences_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/postgreslistsequences"
)
func TestParseFromYamlPostgresListSequences(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-sequences
source: my-postgres-instance
description: some description
authRequired:
- my-google-auth-service
- other-auth-service
`,
want: server.ToolConfigs{
"example_tool": postgreslistsequences.Config{
Name: "example_tool",
Kind: "postgres-list-sequences",
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-sequences
source: my-postgres-instance
description: some description
`,
want: server.ToolConfigs{
"example_tool": postgreslistsequences.Config{
Name: "example_tool",
Kind: "postgres-list-sequences",
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)
}
})
}
}

View File

@@ -188,6 +188,8 @@ func TestAlloyDBPgToolEndpoints(t *testing.T) {
tests.RunPostgresListInstalledExtensionsTest(t)
tests.RunPostgresDatabaseOverviewTest(t, ctx, pool)
tests.RunPostgresListTriggersTest(t, ctx, pool)
tests.RunPostgresListIndexesTest(t, ctx, pool)
tests.RunPostgresListSequencesTest(t, ctx, pool)
}
// Test connection with different IP type

View File

@@ -172,6 +172,8 @@ func TestCloudSQLPgSimpleToolEndpoints(t *testing.T) {
tests.RunPostgresListInstalledExtensionsTest(t)
tests.RunPostgresDatabaseOverviewTest(t, ctx, pool)
tests.RunPostgresListTriggersTest(t, ctx, pool)
tests.RunPostgresListIndexesTest(t, ctx, pool)
tests.RunPostgresListSequencesTest(t, ctx, pool)
}
// Test connection with different IP type

View File

@@ -200,6 +200,8 @@ func AddPostgresPrebuiltConfig(t *testing.T, config map[string]any) map[string]a
PostgresListViewsToolKind = "postgres-list-views"
PostgresDatabaseOverviewToolKind = "postgres-database-overview"
PostgresListTriggersToolKind = "postgres-list-triggers"
PostgresListIndexesToolKind = "postgres-list-indexes"
PostgresListSequencesToolKind = "postgres-list-sequences"
)
tools, ok := config["tools"].(map[string]any)
@@ -248,6 +250,16 @@ func AddPostgresPrebuiltConfig(t *testing.T, config map[string]any) map[string]a
"kind": PostgresListTriggersToolKind,
"source": "my-instance",
}
tools["list_indexes"] = map[string]any{
"kind": PostgresListIndexesToolKind,
"source": "my-instance",
}
tools["list_sequences"] = map[string]any{
"kind": PostgresListSequencesToolKind,
"source": "my-instance",
}
config["tools"] = tools
return config
}

View File

@@ -151,4 +151,6 @@ func TestPostgres(t *testing.T) {
tests.RunPostgresListInstalledExtensionsTest(t)
tests.RunPostgresDatabaseOverviewTest(t, ctx, pool)
tests.RunPostgresListTriggersTest(t, ctx, pool)
tests.RunPostgresListIndexesTest(t, ctx, pool)
tests.RunPostgresListSequencesTest(t, ctx, pool)
}

View File

@@ -1826,6 +1826,257 @@ func RunPostgresListInstalledExtensionsTest(t *testing.T) {
}
}
func setupPostgresIndex(t *testing.T, ctx context.Context, pool *pgxpool.Pool, schemaName string, tableName string) func(t *testing.T) {
t.Helper()
createSchemaStmt := fmt.Sprintf("CREATE SCHEMA IF NOT EXISTS %s;", schemaName)
if _, err := pool.Exec(ctx, createSchemaStmt); err != nil {
t.Fatalf("unable to create schema %s: %v", schemaName, err)
}
fullTableName := fmt.Sprintf("%s.%s", schemaName, tableName)
createTableStmt := fmt.Sprintf("CREATE TABLE %s (id SERIAL PRIMARY KEY, name TEXT, email TEXT);", fullTableName)
if _, err := pool.Exec(ctx, createTableStmt); err != nil {
t.Fatalf("unable to create table %s: %v", fullTableName, err)
}
// Create a unique index on email
index1Stmt := fmt.Sprintf("CREATE UNIQUE INDEX %s_email_idx ON %s (email);", tableName, fullTableName)
if _, err := pool.Exec(ctx, index1Stmt); err != nil {
t.Fatalf("unable to create index %s_email_idx: %v", tableName, err)
}
// Create a non-unique index on name
index2Stmt := fmt.Sprintf("CREATE INDEX %s_name_idx ON %s (name);", tableName, fullTableName)
if _, err := pool.Exec(ctx, index2Stmt); err != nil {
t.Fatalf("unable to create index %s_name_idx: %v", tableName, err)
}
return func(t *testing.T) {
t.Helper()
if _, err := pool.Exec(ctx, fmt.Sprintf("DROP SCHEMA IF EXISTS %s CASCADE;", schemaName)); err != nil {
t.Errorf("unable to drop schema: %v", err)
}
}
}
func RunPostgresListIndexesTest(t *testing.T, ctx context.Context, pool *pgxpool.Pool) {
schemaName := "testschema_" + strings.ReplaceAll(uuid.New().String(), "-", "")
tableName := "table1_" + strings.ReplaceAll(uuid.New().String(), "-", "")
cleanup := setupPostgresIndex(t, ctx, pool, schemaName, tableName)
defer cleanup(t)
// Primary key index
wantIndexPK := map[string]any{
"schema_name": schemaName,
"table_name": tableName,
"index_name": tableName + "_pkey",
"index_type": "btree",
"is_unique": true,
"is_primary": true,
"is_used": false,
"index_definition": fmt.Sprintf("CREATE UNIQUE INDEX %s_pkey ON %s.%s USING btree (id)", tableName, schemaName, tableName),
// Size and scan counts can vary, so omitting them from strict checks or using ranges might be better in real tests.
}
// Email unique index
wantIndexEmail := map[string]any{
"schema_name": schemaName,
"table_name": tableName,
"index_name": tableName + "_email_idx",
"index_type": "btree",
"is_unique": true,
"is_primary": false,
"is_used": false,
"index_definition": fmt.Sprintf("CREATE UNIQUE INDEX %s_email_idx ON %s.%s USING btree (email)", tableName, schemaName, tableName),
}
// Name non-unique index
wantIndexName := map[string]any{
"schema_name": schemaName,
"table_name": tableName,
"index_name": tableName + "_name_idx",
"index_type": "btree",
"is_unique": false,
"is_primary": false,
"is_used": false,
"index_definition": fmt.Sprintf("CREATE INDEX %s_name_idx ON %s.%s USING btree (name)", tableName, schemaName, tableName),
}
allWantIndexes := []map[string]any{wantIndexEmail, wantIndexName, wantIndexPK}
invokeTcs := []struct {
name string
requestBody io.Reader
wantStatusCode int
want []map[string]any
}{
// List all indexes is skipped because the output might include indexes for other database tables
// defined outside of this test, which could make the test flaky.
{
name: "list_indexes for a specific schema and table",
requestBody: bytes.NewBufferString(fmt.Sprintf(`{"schema_name": "%s", "table_name": "%s"}`, schemaName, tableName)),
wantStatusCode: http.StatusOK,
want: allWantIndexes,
},
{
name: "list_indexes for a specific schema",
requestBody: bytes.NewBufferString(fmt.Sprintf(`{"schema_name": "%s"}`, schemaName)),
wantStatusCode: http.StatusOK,
want: allWantIndexes,
},
{
name: "list_indexes with non-existent schema",
requestBody: bytes.NewBufferString(`{"schema_name": "non_existent_schema"}`),
wantStatusCode: http.StatusOK,
want: nil,
},
{
name: "list_indexes with non-existent table in existing schema",
requestBody: bytes.NewBufferString(fmt.Sprintf(`{"schema_name": "%s", "table_name": "non_existent_table"}`, schemaName)),
wantStatusCode: http.StatusOK,
want: nil,
},
{
name: "list_indexes filter by index name",
requestBody: bytes.NewBufferString(fmt.Sprintf(`{"schema_name": "%s", "table_name": "%s", "index_name": "%s"}`, schemaName, tableName, tableName+"_email_idx")),
wantStatusCode: http.StatusOK,
want: []map[string]any{wantIndexEmail},
},
{
name: "list_indexes filter by non-existent index name",
requestBody: bytes.NewBufferString(fmt.Sprintf(`{"schema_name": "%s", "table_name": "%s", "index_name": "non_existent_idx"}`, schemaName, tableName)),
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_indexes/invoke"
resp, respBody := RunRequest(t, http.MethodPost, api, tc.requestBody, nil)
if resp.StatusCode != tc.wantStatusCode {
t.Fatalf("wrong status code: got %d, want %d, body: %s", resp.StatusCode, tc.wantStatusCode, string(respBody))
}
if tc.wantStatusCode != http.StatusOK {
return
}
var bodyWrapper struct {
Result json.RawMessage `json:"result"`
}
if err := json.Unmarshal(respBody, &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, resultString: %s", err, resultString)
}
// Normalize got by removing fields that are hard to predict (like size)
for _, index := range got {
delete(index, "index_size_bytes")
delete(index, "index_scans")
delete(index, "tuples_read")
delete(index, "tuples_fetched")
}
if diff := cmp.Diff(tc.want, got); diff != "" {
t.Errorf("Unexpected result (-want +got):\n%s", diff)
}
})
}
}
func setupListSequencesTest(t *testing.T, ctx context.Context, pool *pgxpool.Pool) (string, func(t *testing.T)) {
sequenceName := "list_sequences_seq1_" + strings.ReplaceAll(uuid.New().String(), "-", "")
createSequence1Stmt := fmt.Sprintf("CREATE SEQUENCE %s INCREMENT 1 START 1;", sequenceName)
_, err := pool.Exec(ctx, createSequence1Stmt)
if err != nil {
t.Fatalf("unable to create sequence %s: %s", sequenceName, err)
}
return sequenceName, func(t *testing.T) {
_, err := pool.Exec(ctx, fmt.Sprintf("DROP SEQUENCE IF EXISTS %s;", sequenceName))
if err != nil {
t.Errorf("unable to drop sequences: %v", err)
}
}
}
func RunPostgresListSequencesTest(t *testing.T, ctx context.Context, pool *pgxpool.Pool) {
sequenceName, teardown := setupListSequencesTest(t, ctx, pool)
defer teardown(t)
wantSequence := map[string]any{
"sequencename": sequenceName,
"schemaname": "public",
"sequenceowner": "postgres",
"data_type": "bigint",
"start_value": float64(1),
"min_value": float64(1),
"max_value": float64(9223372036854775807),
"increment_by": float64(1),
"last_value": nil,
}
invokeTcs := []struct {
name string
api string
requestBody io.Reader
wantStatusCode int
want []map[string]any
}{
{
name: "invoke list_sequences",
requestBody: bytes.NewBufferString(fmt.Sprintf(`{"sequencename": "%s"}`, sequenceName)),
wantStatusCode: http.StatusOK,
want: []map[string]any{wantSequence},
},
{
name: "invoke list_sequences with non-existent sequence",
requestBody: bytes.NewBufferString(`{"sequencename": "non_existent_sequence"}`),
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_sequences/invoke"
resp, respBody := RunRequest(t, http.MethodPost, api, tc.requestBody, nil)
if resp.StatusCode != tc.wantStatusCode {
t.Fatalf("wrong status code: got %d, want %d, body: %s", resp.StatusCode, tc.wantStatusCode, string(respBody))
}
if tc.wantStatusCode != http.StatusOK {
return
}
var bodyWrapper struct {
Result json.RawMessage `json:"result"`
}
if err := json.Unmarshal(respBody, &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 {