mirror of
https://github.com/googleapis/genai-toolbox.git
synced 2026-01-11 08:28:11 -05:00
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:
@@ -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"
|
||||
|
||||
@@ -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"},
|
||||
},
|
||||
},
|
||||
},
|
||||
|
||||
@@ -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
|
||||
|
||||
|
||||
@@ -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/)
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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/)
|
||||
|
||||
67
docs/en/resources/tools/postgres/postgres-list-indexes.md
Normal file
67
docs/en/resources/tools/postgres/postgres-list-indexes.md
Normal 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. |
|
||||
61
docs/en/resources/tools/postgres/postgres-list-sequences.md
Normal file
61
docs/en/resources/tools/postgres/postgres-list-sequences.md
Normal 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. |
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
}
|
||||
@@ -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)
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
}
|
||||
@@ -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
|
||||
}
|
||||
@@ -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)
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
}
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
}
|
||||
|
||||
@@ -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)
|
||||
}
|
||||
|
||||
251
tests/tool.go
251
tests/tool.go
@@ -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 {
|
||||
|
||||
Reference in New Issue
Block a user