From 32367a472fae9653fed7f126428eba0252978bd5 Mon Sep 17 00:00:00 2001 From: Srividya Reddy Date: Wed, 10 Dec 2025 01:25:53 +0530 Subject: [PATCH] feat(tools/postgres): add list_pg_settings, list_database_stats tools for postgres (#2030) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## Description Adds the following tools for Postgres: (1) list_pg_settings: List configuration parameters for the PostgreSQL server. (2) list_database_stats: Lists the key performance and activity statistics for each database in the postgreSQL server. > Should include a concise description of the changes (bug or feature), it's > impact, along with a summary of the solution list_pg_settings: Screenshot 2025-11-25 at 10 19 48 AM Screenshot 2025-11-25 at 10 27 19 AM list_database_stats: Screenshot 2025-11-25 at 10 21 12 AM Screenshot 2025-11-25 at 10 27 47 AM ## 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> Co-authored-by: Averi Kitsch Co-authored-by: Wenxin Du <117315983+duwenxin99@users.noreply.github.com> --- cmd/root.go | 2 + cmd/root_test.go | 6 +- docs/en/reference/prebuilt-tools.md | 9 + docs/en/resources/sources/alloydb-pg.md | 7 + docs/en/resources/sources/cloud-sql-pg.md | 7 + docs/en/resources/sources/postgres.md | 7 + .../postgres/postgres-list-database-stats.md | 95 ++++++ .../postgres/postgres-list-pg-settings.md | 59 ++++ .../tools/alloydb-postgres.yaml | 10 + .../tools/cloud-sql-postgres.yaml | 10 + internal/prebuiltconfigs/tools/postgres.yaml | 10 + .../postgreslistdatabasestats.go | 276 ++++++++++++++++++ .../postgreslistdatabasestats_test.go | 95 ++++++ .../postgreslistpgsettings.go | 204 +++++++++++++ .../postgreslistpgsettings_test.go | 95 ++++++ .../alloydbpg/alloydb_pg_integration_test.go | 2 + .../cloud_sql_pg_integration_test.go | 2 + tests/common.go | 22 +- tests/postgres/postgres_integration_test.go | 2 + tests/tool.go | 235 +++++++++++++++ 20 files changed, 1141 insertions(+), 14 deletions(-) create mode 100644 docs/en/resources/tools/postgres/postgres-list-database-stats.md create mode 100644 docs/en/resources/tools/postgres/postgres-list-pg-settings.md create mode 100644 internal/tools/postgres/postgreslistdatabasestats/postgreslistdatabasestats.go create mode 100644 internal/tools/postgres/postgreslistdatabasestats/postgreslistdatabasestats_test.go create mode 100644 internal/tools/postgres/postgreslistpgsettings/postgreslistpgsettings.go create mode 100644 internal/tools/postgres/postgreslistpgsettings/postgreslistpgsettings_test.go diff --git a/cmd/root.go b/cmd/root.go index 4df639a765..8689a03c69 100644 --- a/cmd/root.go +++ b/cmd/root.go @@ -184,9 +184,11 @@ import ( _ "github.com/googleapis/genai-toolbox/internal/tools/postgres/postgresgetcolumncardinality" _ "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/postgreslistdatabasestats" _ "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/postgreslistlocks" + _ "github.com/googleapis/genai-toolbox/internal/tools/postgres/postgreslistpgsettings" _ "github.com/googleapis/genai-toolbox/internal/tools/postgres/postgreslistpublicationtables" _ "github.com/googleapis/genai-toolbox/internal/tools/postgres/postgreslistquerystats" _ "github.com/googleapis/genai-toolbox/internal/tools/postgres/postgreslistschemas" diff --git a/cmd/root_test.go b/cmd/root_test.go index 025008c54b..40e2b0975a 100644 --- a/cmd/root_test.go +++ b/cmd/root_test.go @@ -1488,7 +1488,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", "list_indexes", "list_sequences", "long_running_transactions", "list_locks", "replication_stats", "list_query_stats", "get_column_cardinality", "list_publication_tables", "list_tablespaces"}, + 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", "long_running_transactions", "list_locks", "replication_stats", "list_query_stats", "get_column_cardinality", "list_publication_tables", "list_tablespaces", "list_pg_settings", "list_database_stats"}, }, }, }, @@ -1518,7 +1518,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", "list_indexes", "list_sequences", "long_running_transactions", "list_locks", "replication_stats", "list_query_stats", "get_column_cardinality", "list_publication_tables", "list_tablespaces"}, + 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", "long_running_transactions", "list_locks", "replication_stats", "list_query_stats", "get_column_cardinality", "list_publication_tables", "list_tablespaces", "list_pg_settings", "list_database_stats"}, }, }, }, @@ -1618,7 +1618,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", "list_indexes", "list_sequences", "long_running_transactions", "list_locks", "replication_stats", "list_query_stats", "get_column_cardinality", "list_publication_tables", "list_tablespaces"}, + 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", "long_running_transactions", "list_locks", "replication_stats", "list_query_stats", "get_column_cardinality", "list_publication_tables", "list_tablespaces", "list_pg_settings", "list_database_stats"}, }, }, }, diff --git a/docs/en/reference/prebuilt-tools.md b/docs/en/reference/prebuilt-tools.md index 053f2423b8..4c57a941c8 100644 --- a/docs/en/reference/prebuilt-tools.md +++ b/docs/en/reference/prebuilt-tools.md @@ -52,6 +52,9 @@ details on how to connect your AI tools (IDEs) to databases via Toolbox and MCP. * `list_sequences`: List sequences in a PostgreSQL database. * `list_publication_tables`: List publication tables in a PostgreSQL database. * `list_tablespaces`: Lists tablespaces in the database. + * `list_pg_settings`: List configuration parameters for the PostgreSQL server. + * `list_database_stats`: Lists the key performance and activity statistics for + each database in the AlloyDB instance. ## AlloyDB Postgres Admin @@ -231,6 +234,9 @@ details on how to connect your AI tools (IDEs) to databases via Toolbox and MCP. * `list_sequences`: List sequences in a PostgreSQL database. * `list_publication_tables`: List publication tables in a PostgreSQL database. * `list_tablespaces`: Lists tablespaces in the database. + * `list_pg_settings`: List configuration parameters for the PostgreSQL server. + * `list_database_stats`: Lists the key performance and activity statistics for + each database in the postgreSQL instance. ## Cloud SQL for PostgreSQL Observability @@ -538,6 +544,9 @@ details on how to connect your AI tools (IDEs) to databases via Toolbox and MCP. * `list_sequences`: List sequences in a PostgreSQL database. * `list_publication_tables`: List publication tables in a PostgreSQL database. * `list_tablespaces`: Lists tablespaces in the database. + * `list_pg_settings`: List configuration parameters for the PostgreSQL server. + * `list_database_stats`: Lists the key performance and activity statistics for + each database in the PostgreSQL server. ## Google Cloud Serverless for Apache Spark diff --git a/docs/en/resources/sources/alloydb-pg.md b/docs/en/resources/sources/alloydb-pg.md index 4c59dc7429..5bc272f63b 100644 --- a/docs/en/resources/sources/alloydb-pg.md +++ b/docs/en/resources/sources/alloydb-pg.md @@ -83,6 +83,13 @@ cluster][alloydb-free-trial]. - [`postgres-list-tablespaces`](../tools/postgres/postgres-list-tablespaces.md) List tablespaces in an AlloyDB for PostgreSQL database. +- [`postgres-list-pg-settings`](../tools/postgres/postgres-list-pg-settings.md) + List configuration parameters for the PostgreSQL server. + +- [`postgres-list-database-stats`](../tools/postgres/postgres-list-database-stats.md) + Lists the key performance and activity statistics for each database in the AlloyDB + instance. + ### Pre-built Configurations - [AlloyDB using MCP](https://googleapis.github.io/genai-toolbox/how-to/connect-ide/alloydb_pg_mcp/) diff --git a/docs/en/resources/sources/cloud-sql-pg.md b/docs/en/resources/sources/cloud-sql-pg.md index 6e73df96f6..a935a41463 100644 --- a/docs/en/resources/sources/cloud-sql-pg.md +++ b/docs/en/resources/sources/cloud-sql-pg.md @@ -79,6 +79,13 @@ to a database by following these instructions][csql-pg-quickstart]. - [`postgres-list-tablespaces`](../tools/postgres/postgres-list-tablespaces.md) List tablespaces in a PostgreSQL database. +- [`postgres-list-pg-settings`](../tools/postgres/postgres-list-pg-settings.md) + List configuration parameters for the PostgreSQL server. + +- [`postgres-list-database-stats`](../tools/postgres/postgres-list-database-stats.md) + Lists the key performance and activity statistics for each database in the postgreSQL + instance. + ### Pre-built Configurations - [Cloud SQL for Postgres using diff --git a/docs/en/resources/sources/postgres.md b/docs/en/resources/sources/postgres.md index 6f1785e529..0997bcf0e3 100644 --- a/docs/en/resources/sources/postgres.md +++ b/docs/en/resources/sources/postgres.md @@ -74,6 +74,13 @@ reputation for reliability, feature robustness, and performance. - [`postgres-list-tablespaces`](../tools/postgres/postgres-list-tablespaces.md) List tablespaces in a PostgreSQL database. +- [`postgres-list-pg-settings`](../tools/postgres/postgres-list-pg-settings.md) + List configuration parameters for the PostgreSQL server. + +- [`postgres-list-database-stats`](../tools/postgres/postgres-list-database-stats.md) + Lists the key performance and activity statistics for each database in the postgreSQL + server. + ### Pre-built Configurations - [PostgreSQL using MCP](https://googleapis.github.io/genai-toolbox/how-to/connect-ide/postgres_mcp/) diff --git a/docs/en/resources/tools/postgres/postgres-list-database-stats.md b/docs/en/resources/tools/postgres/postgres-list-database-stats.md new file mode 100644 index 0000000000..01537bcfa7 --- /dev/null +++ b/docs/en/resources/tools/postgres/postgres-list-database-stats.md @@ -0,0 +1,95 @@ +--- +title: "postgres-list-database-stats" +type: docs +weight: 1 +description: > + The "postgres-list-database-stats" tool lists lists key performance and activity statistics of PostgreSQL databases. +aliases: +- /resources/tools/postgres-list-database-stats +--- + +## About + +The `postgres-list-database-stats` lists the key performance and activity statistics for each PostgreSQL database in the instance, offering insights into cache efficiency, transaction throughput, row-level activity, temporary file usage, and contention. 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-database-stats` lists detailed information as JSON for each database. The tool +takes the following input parameters: + +- `database_name` (optional): A text to filter results by database name. Default: `""` +- `include_templates` (optional): Boolean, set to `true` to include template databases in the results. Default: `false` +- `database_owner` (optional): A text to filter results by database owner. Default: `""` +- `default_tablespace` (optional): A text to filter results by the default tablespace name. Default: `""` +- `order_by` (optional): Specifies the sorting order. Valid values are `'size'` (descending) or `'commit'` (descending). Default: `database_name` ascending. +- `limit` (optional): The maximum number of databases to return. Default: `10` + +## Example + +```yaml +tools: + list_database_stats: + kind: postgres-list-database-stats + source: postgres-source + description: | + Lists the key performance and activity statistics for each PostgreSQL + database in the instance, offering insights into cache efficiency, + transaction throughput row-level activity, temporary file usage, and + contention. It returns: the database name, whether the database is + connectable, database owner, default tablespace name, the percentage of + data blocks found in the buffer cache rather than being read from disk + (a higher value indicates better cache performance), the total number of + disk blocks read from disk, the total number of times disk blocks were + found already in the cache; the total number of committed transactions, + the total number of rolled back transactions, the percentage of rolled + back transactions compared to the total number of completed + transactions, the total number of rows returned by queries, the total + number of live rows fetched by scans, the total number of rows inserted, + the total number of rows updated, the total number of rows deleted, the + number of temporary files created by queries, the total size of + temporary files used by queries in bytes, the number of query + cancellations due to conflicts with recovery, the number of deadlocks + detected, the current number of active backend connections, the + timestamp when the database statistics were last reset, and the total + database size in bytes. +``` + +The response is a json array with the following elements: + +```json +{ + "database_name": "Name of the database", + "is_connectable": "Boolean indicating Whether the database allows connections", + "database_owner": "Username of the database owner", + "default_tablespace": "Name of the default tablespace for the database", + "cache_hit_ratio_percent": "The percentage of data blocks found in the buffer cache rather than being read from disk", + "blocks_read_from_disk": "The total number of disk blocks read for this database", + "blocks_hit_in_cache": "The total number of times disk blocks were found already in the cache.", + "xact_commit": "The total number of committed transactions", + "xact_rollback": "The total number of rolled back transactions", + "rollback_ratio_percent": "The percentage of rolled back transactions compared to the total number of completed transactions", + "rows_returned_by_queries": "The total number of rows returned by queries", + "rows_fetched_by_scans": "The total number of live rows fetched by scans", + "tup_inserted": "The total number of rows inserted", + "tup_updated": "The total number of rows updated", + "tup_deleted": "The total number of rows deleted", + "temp_files": "The number of temporary files created by queries", + "temp_size_bytes": "The total size of temporary files used by queries in bytes", + "conflicts": "Number of query cancellations due to conflicts", + "deadlocks": "Number of deadlocks detected", + "active_connections": "The current number of active backend connections", + "statistics_last_reset": "The timestamp when the database statistics were last reset", + "database_size_bytes": "The total disk size of the database in bytes" +} +``` + +## Reference + +| **field** | **type** | **required** | **description** | +|-------------|:--------:|:------------:|------------------------------------------------------| +| kind | string | true | Must be "postgres-list-database-stats". | +| 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. | diff --git a/docs/en/resources/tools/postgres/postgres-list-pg-settings.md b/docs/en/resources/tools/postgres/postgres-list-pg-settings.md new file mode 100644 index 0000000000..23d5e28e92 --- /dev/null +++ b/docs/en/resources/tools/postgres/postgres-list-pg-settings.md @@ -0,0 +1,59 @@ +--- +title: "postgres-list-pg-settings" +type: docs +weight: 1 +description: > + The "postgres-list-pg-settings" tool lists PostgreSQL run-time configuration settings. +aliases: +- /resources/tools/postgres-list-pg-settings +--- + +## About + +The `postgres-list-pg-settings` tool lists the configuration parameters for the postgres server, their current values, and related information. 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-pg-settings` lists detailed information as JSON for each setting. The tool +takes the following input parameters: + +- `setting_name` (optional): A text to filter results by setting name. Default: `""` +- `limit` (optional): The maximum number of rows to return. Default: `50`. + +## Example + +```yaml +tools: + list_indexes: + kind: postgres-list-pg-settings + source: postgres-source + description: | + Lists configuration parameters for the postgres server ordered lexicographically, + with a default limit of 50 rows. It returns the parameter name, its current setting, + unit of measurement, a short description, the source of the current setting (e.g., + default, configuration file, session), and whether a restart is required when the + parameter value is changed." +``` + +The response is a json array with the following elements: + +```json +{ + "name": "Setting name", + "current_value": "Current value of the setting", + "unit": "Unit of the setting", + "short_desc": "Short description of the setting", + "source": "Source of the current value (e.g., default, configuration file, session)", + "requires_restart": "Indicates if a server restart is required to apply a change ('Yes', 'No', or 'No (Reload sufficient)')" +} +``` + +## Reference + +| **field** | **type** | **required** | **description** | +|-------------|:--------:|:------------:|------------------------------------------------------| +| kind | string | true | Must be "postgres-list-pg-settings". | +| 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. | diff --git a/internal/prebuiltconfigs/tools/alloydb-postgres.yaml b/internal/prebuiltconfigs/tools/alloydb-postgres.yaml index f9cf892f1e..a85f1b1ea5 100644 --- a/internal/prebuiltconfigs/tools/alloydb-postgres.yaml +++ b/internal/prebuiltconfigs/tools/alloydb-postgres.yaml @@ -208,6 +208,14 @@ tools: kind: postgres-list-tablespaces source: alloydb-pg-source + list_pg_settings: + kind: postgres-list-pg-settings + source: alloydb-pg-source + + list_database_stats: + kind: postgres-list-database-stats + source: alloydb-pg-source + toolsets: alloydb_postgres_database_tools: - execute_sql @@ -234,3 +242,5 @@ toolsets: - get_column_cardinality - list_publication_tables - list_tablespaces + - list_pg_settings + - list_database_stats diff --git a/internal/prebuiltconfigs/tools/cloud-sql-postgres.yaml b/internal/prebuiltconfigs/tools/cloud-sql-postgres.yaml index fd062f214e..83aabbc196 100644 --- a/internal/prebuiltconfigs/tools/cloud-sql-postgres.yaml +++ b/internal/prebuiltconfigs/tools/cloud-sql-postgres.yaml @@ -210,6 +210,14 @@ tools: kind: postgres-list-tablespaces source: cloudsql-pg-source + list_pg_settings: + kind: postgres-list-pg-settings + source: cloudsql-pg-source + + list_database_stats: + kind: postgres-list-database-stats + source: cloudsql-pg-source + toolsets: cloud_sql_postgres_database_tools: - execute_sql @@ -236,3 +244,5 @@ toolsets: - get_column_cardinality - list_publication_tables - list_tablespaces + - list_pg_settings + - list_database_stats diff --git a/internal/prebuiltconfigs/tools/postgres.yaml b/internal/prebuiltconfigs/tools/postgres.yaml index b09c32f52c..6ada1b8647 100644 --- a/internal/prebuiltconfigs/tools/postgres.yaml +++ b/internal/prebuiltconfigs/tools/postgres.yaml @@ -209,6 +209,14 @@ tools: kind: postgres-list-tablespaces source: postgresql-source + list_pg_settings: + kind: postgres-list-pg-settings + source: postgresql-source + + list_database_stats: + kind: postgres-list-database-stats + source: postgresql-source + toolsets: postgres_database_tools: - execute_sql @@ -235,3 +243,5 @@ toolsets: - get_column_cardinality - list_publication_tables - list_tablespaces + - list_pg_settings + - list_database_stats diff --git a/internal/tools/postgres/postgreslistdatabasestats/postgreslistdatabasestats.go b/internal/tools/postgres/postgreslistdatabasestats/postgreslistdatabasestats.go new file mode 100644 index 0000000000..31edc08f11 --- /dev/null +++ b/internal/tools/postgres/postgreslistdatabasestats/postgreslistdatabasestats.go @@ -0,0 +1,276 @@ +// 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 postgreslistdatabasestats + +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/googleapis/genai-toolbox/internal/util/parameters" + "github.com/jackc/pgx/v5/pgxpool" +) + +const kind string = "postgres-list-database-stats" + +// SQL query to list database statistics +const listDatabaseStats = ` + WITH database_stats AS ( + SELECT + s.datname AS database_name, + -- Database Metadata + d.datallowconn AS is_connectable, + pg_get_userbyid(d.datdba) AS database_owner, + ts.spcname AS default_tablespace, + + -- Cache Performance + CASE + WHEN (s.blks_hit + s.blks_read) = 0 THEN 0 + ELSE round((s.blks_hit * 100.0) / (s.blks_hit + s.blks_read), 2) + END AS cache_hit_ratio_percent, + s.blks_read AS blocks_read_from_disk, + s.blks_hit AS blocks_hit_in_cache, + + -- Transaction Throughput + s.xact_commit, + s.xact_rollback, + round(s.xact_rollback * 100.0 / (s.xact_commit + s.xact_rollback + 1), 2) AS rollback_ratio_percent, + + -- Tuple Activity + s.tup_returned AS rows_returned_by_queries, + s.tup_fetched AS rows_fetched_by_scans, + s.tup_inserted, + s.tup_updated, + s.tup_deleted, + + -- Temporary File Usage + s.temp_files, + s.temp_bytes AS temp_size_bytes, + + -- Conflicts & Deadlocks + s.conflicts, + s.deadlocks, + + -- General Info + s.numbackends AS active_connections, + s.stats_reset AS statistics_last_reset, + pg_database_size(s.datid) AS database_size_bytes + FROM + pg_stat_database s + JOIN + pg_database d ON d.oid = s.datid + JOIN + pg_tablespace ts ON ts.oid = d.dattablespace + WHERE + -- Exclude cloudsql internal databases + s.datname NOT IN ('cloudsqladmin') + -- Exclude template databases if not requested + AND ( $2::boolean IS TRUE OR d.datistemplate IS FALSE ) + ) + SELECT * + FROM database_stats + WHERE + ($1::text IS NULL OR database_name LIKE '%' || $1::text || '%') + AND ($3::text IS NULL OR database_owner LIKE '%' || $3::text || '%') + AND ($4::text IS NULL OR default_tablespace LIKE '%' || $4::text || '%') + ORDER BY + CASE WHEN $5::text = 'size' THEN database_size_bytes END DESC, + CASE WHEN $5::text = 'commit' THEN xact_commit END DESC, + database_name + LIMIT COALESCE($6::int, 10); +` + +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 := parameters.Parameters{ + parameters.NewStringParameterWithDefault("database_name", "", "Optional: A specific database name pattern to search for."), + parameters.NewBooleanParameterWithDefault("include_templates", false, "Optional: Whether to include template databases in the results."), + parameters.NewStringParameterWithDefault("database_owner", "", "Optional: A specific database owner name pattern to search for."), + parameters.NewStringParameterWithDefault("default_tablespace", "", "Optional: A specific default tablespace name pattern to search for."), + parameters.NewStringParameterWithDefault("order_by", "", "Optional: The field to order the results by. Valid values are 'size' and 'commit'."), + parameters.NewIntParameterWithDefault("limit", 10, "Optional: The maximum number of rows to return."), + } + description := cfg.Description + if description == "" { + description = + "Lists the key performance and activity statistics for each PostgreSQL database" + + "in the instance, offering insights into cache efficiency, transaction throughput" + + "row-level activity, temporary file " + + "usage, and contention. " + + "It returns: the database name, whether the database is connectable, " + + "database owner, default tablespace name, the percentage of data blocks " + + "found in the buffer cache rather than being read from disk (a higher " + + "value indicates better cache performance), the total number of disk " + + "blocks read from disk, the total number of times disk blocks were found " + + "already in the cache; the total number of committed transactions, the " + + "total number of rolled back transactions, the percentage of rolled back " + + "transactions compared to the total number of completed transactions, the " + + "total number of rows returned by queries, the total number of live rows " + + "fetched by scans, the total number of rows inserted, the total number " + + "of rows updated, the total number of rows deleted, the number of " + + "temporary files created by queries, the total size of all temporary " + + "files created by queries in bytes, the number of query cancellations due " + + "to conflicts with recovery, the number of deadlocks detected, the current " + + "number of active connections to the database, the timestamp of the " + + "last statistics reset, and total database size in bytes." + } + mcpManifest := tools.GetMcpManifest(cfg.Name, description, cfg.AuthRequired, allParameters, nil) + + // finish tool setup + return Tool{ + Config: cfg, + 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 { + Config + allParams parameters.Parameters `yaml:"allParams"` + pool *pgxpool.Pool + manifest tools.Manifest + mcpManifest tools.McpManifest +} + +func (t Tool) Invoke(ctx context.Context, resourceMgr tools.SourceProvider, params parameters.ParamValues, accessToken tools.AccessToken) (any, error) { + paramsMap := params.AsMap() + + newParams, err := parameters.GetParams(t.allParams, paramsMap) + if err != nil { + return nil, fmt.Errorf("unable to extract standard params %w", err) + } + sliceParams := newParams.AsSlice() + + results, err := t.pool.Query(ctx, listDatabaseStats, 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) + } + + // this will catch actual query execution errors + if err := results.Err(); err != nil { + return nil, fmt.Errorf("unable to execute query: %w", err) + } + + return out, nil +} + +func (t Tool) ParseParams(data map[string]any, claims map[string]map[string]any) (parameters.ParamValues, error) { + return parameters.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(resourceMgr tools.SourceProvider) bool { + return false +} + +func (t Tool) ToConfig() tools.ToolConfig { + return t.Config +} + +func (t Tool) GetAuthTokenHeaderName() string { + return "Authorization" +} diff --git a/internal/tools/postgres/postgreslistdatabasestats/postgreslistdatabasestats_test.go b/internal/tools/postgres/postgreslistdatabasestats/postgreslistdatabasestats_test.go new file mode 100644 index 0000000000..760370f630 --- /dev/null +++ b/internal/tools/postgres/postgreslistdatabasestats/postgreslistdatabasestats_test.go @@ -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 postgreslistdatabasestats_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/postgreslistdatabasestats" +) + +func TestParseFromYamlPostgresListDatabaseStats(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-database-stats + source: my-postgres-instance + description: some description + authRequired: + - my-google-auth-service + - other-auth-service + `, + want: server.ToolConfigs{ + "example_tool": postgreslistdatabasestats.Config{ + Name: "example_tool", + Kind: "postgres-list-database-stats", + 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-database-stats + source: my-postgres-instance + description: some description + `, + want: server.ToolConfigs{ + "example_tool": postgreslistdatabasestats.Config{ + Name: "example_tool", + Kind: "postgres-list-database-stats", + 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) + } + }) + } + +} diff --git a/internal/tools/postgres/postgreslistpgsettings/postgreslistpgsettings.go b/internal/tools/postgres/postgreslistpgsettings/postgreslistpgsettings.go new file mode 100644 index 0000000000..ae5b9ff2dd --- /dev/null +++ b/internal/tools/postgres/postgreslistpgsettings/postgreslistpgsettings.go @@ -0,0 +1,204 @@ +// 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 postgreslistpgsettings + +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/googleapis/genai-toolbox/internal/util/parameters" + "github.com/jackc/pgx/v5/pgxpool" +) + +const kind string = "postgres-list-pg-settings" + +const listPgSettingsStatement = ` + SELECT + name, + setting AS current_value, + unit, + short_desc, + source, + CASE context + WHEN 'postmaster' THEN 'Yes' + WHEN 'sighup' THEN 'No (Reload sufficient)' + ELSE 'No' + END + AS requires_restart + FROM pg_settings + WHERE ($1::text IS NULL OR name LIKE '%' || $1::text || '%') + ORDER BY name + LIMIT COALESCE($2::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 := parameters.Parameters{ + parameters.NewStringParameterWithDefault("setting_name", "", "Optional: A specific configuration parameter name pattern to search for."), + parameters.NewIntParameterWithDefault("limit", 50, "Optional: The maximum number of rows to return."), + } + description := cfg.Description + if description == "" { + description = "Lists configuration parameters for the postgres server ordered lexicographically, with a default limit of 50 rows. It returns the parameter name, its current setting, unit of measurement, a short description, the source of the current setting (e.g., default, configuration file, session), and whether a restart is required when the parameter value is changed." + } + mcpManifest := tools.GetMcpManifest(cfg.Name, description, cfg.AuthRequired, allParameters, nil) + + // finish tool setup + return Tool{ + Config: cfg, + 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 { + Config + allParams parameters.Parameters `yaml:"allParams"` + pool *pgxpool.Pool + manifest tools.Manifest + mcpManifest tools.McpManifest +} + +func (t Tool) Invoke(ctx context.Context, resourceMgr tools.SourceProvider, params parameters.ParamValues, accessToken tools.AccessToken) (any, error) { + paramsMap := params.AsMap() + + newParams, err := parameters.GetParams(t.allParams, paramsMap) + if err != nil { + return nil, fmt.Errorf("unable to extract standard params %w", err) + } + sliceParams := newParams.AsSlice() + + results, err := t.pool.Query(ctx, listPgSettingsStatement, 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) + } + + // this will catch actual query execution errors + if err := results.Err(); err != nil { + return nil, fmt.Errorf("unable to execute query: %w", err) + } + + return out, nil +} + +func (t Tool) ParseParams(data map[string]any, claims map[string]map[string]any) (parameters.ParamValues, error) { + return parameters.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(resourceMgr tools.SourceProvider) bool { + return false +} + +func (t Tool) ToConfig() tools.ToolConfig { + return t.Config +} + +func (t Tool) GetAuthTokenHeaderName() string { + return "Authorization" +} diff --git a/internal/tools/postgres/postgreslistpgsettings/postgreslistpgsettings_test.go b/internal/tools/postgres/postgreslistpgsettings/postgreslistpgsettings_test.go new file mode 100644 index 0000000000..a2aa9fe78b --- /dev/null +++ b/internal/tools/postgres/postgreslistpgsettings/postgreslistpgsettings_test.go @@ -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 postgreslistpgsettings_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/postgreslistpgsettings" +) + +func TestParseFromYamlPostgreslistPgSettings(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-pg-settings + source: my-postgres-instance + description: some description + authRequired: + - my-google-auth-service + - other-auth-service + `, + want: server.ToolConfigs{ + "example_tool": postgreslistpgsettings.Config{ + Name: "example_tool", + Kind: "postgres-list-pg-settings", + 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-pg-settings + source: my-postgres-instance + description: some description + `, + want: server.ToolConfigs{ + "example_tool": postgreslistpgsettings.Config{ + Name: "example_tool", + Kind: "postgres-list-pg-settings", + 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) + } + }) + } + +} diff --git a/tests/alloydbpg/alloydb_pg_integration_test.go b/tests/alloydbpg/alloydb_pg_integration_test.go index 6cbbbd6632..e392a3074e 100644 --- a/tests/alloydbpg/alloydb_pg_integration_test.go +++ b/tests/alloydbpg/alloydb_pg_integration_test.go @@ -197,6 +197,8 @@ func TestAlloyDBPgToolEndpoints(t *testing.T) { tests.RunPostgresGetColumnCardinalityTest(t, ctx, pool) tests.RunPostgresListPublicationTablesTest(t, ctx, pool) tests.RunPostgresListTableSpacesTest(t) + tests.RunPostgresListPgSettingsTest(t, ctx, pool) + tests.RunPostgresListDatabaseStatsTest(t, ctx, pool) } // Test connection with different IP type diff --git a/tests/cloudsqlpg/cloud_sql_pg_integration_test.go b/tests/cloudsqlpg/cloud_sql_pg_integration_test.go index 3b602024ae..98436a986e 100644 --- a/tests/cloudsqlpg/cloud_sql_pg_integration_test.go +++ b/tests/cloudsqlpg/cloud_sql_pg_integration_test.go @@ -181,6 +181,8 @@ func TestCloudSQLPgSimpleToolEndpoints(t *testing.T) { tests.RunPostgresGetColumnCardinalityTest(t, ctx, pool) tests.RunPostgresListPublicationTablesTest(t, ctx, pool) tests.RunPostgresListTableSpacesTest(t) + tests.RunPostgresListPgSettingsTest(t, ctx, pool) + tests.RunPostgresListDatabaseStatsTest(t, ctx, pool) } // Test connection with different IP type diff --git a/tests/common.go b/tests/common.go index c9f9c95e14..442681235f 100644 --- a/tests/common.go +++ b/tests/common.go @@ -209,6 +209,8 @@ func AddPostgresPrebuiltConfig(t *testing.T, config map[string]any) map[string]a PostgresGetColumnCardinalityToolKind = "postgres-get-column-cardinality" PostgresListPublicationTablesToolKind = "postgres-list-publication-tables" PostgresListTablespacesToolKind = "postgres-list-tablespaces" + PostgresListPGSettingsToolKind = "postgres-list-pg-settings" + PostgresListDatabaseStatsToolKind = "postgres-list-database-stats" ) tools, ok := config["tools"].(map[string]any) @@ -225,34 +227,28 @@ func AddPostgresPrebuiltConfig(t *testing.T, config map[string]any) map[string]a "source": "my-instance", "description": "Lists active queries in the database.", } - tools["list_installed_extensions"] = map[string]any{ "kind": PostgresListInstalledExtensionsToolKind, "source": "my-instance", "description": "Lists installed extensions in the database.", } - tools["list_available_extensions"] = map[string]any{ "kind": PostgresListAvailableExtensionsToolKind, "source": "my-instance", "description": "Lists available extensions in the database.", } - tools["list_views"] = map[string]any{ "kind": PostgresListViewsToolKind, "source": "my-instance", } - tools["list_schemas"] = map[string]any{ "kind": PostgresListSchemasToolKind, "source": "my-instance", } - tools["database_overview"] = map[string]any{ "kind": PostgresDatabaseOverviewToolKind, "source": "my-instance", } - tools["list_triggers"] = map[string]any{ "kind": PostgresListTriggersToolKind, "source": "my-instance", @@ -261,27 +257,22 @@ func AddPostgresPrebuiltConfig(t *testing.T, config map[string]any) map[string]a "kind": PostgresListIndexesToolKind, "source": "my-instance", } - tools["list_sequences"] = map[string]any{ "kind": PostgresListSequencesToolKind, "source": "my-instance", } - tools["list_publication_tables"] = map[string]any{ "kind": PostgresListPublicationTablesToolKind, "source": "my-instance", } - tools["long_running_transactions"] = map[string]any{ "kind": PostgresLongRunningTransactionsToolKind, "source": "my-instance", } - tools["list_locks"] = map[string]any{ "kind": PostgresListLocksToolKind, "source": "my-instance", } - tools["replication_stats"] = map[string]any{ "kind": PostgresReplicationStatsToolKind, "source": "my-instance", @@ -298,6 +289,15 @@ func AddPostgresPrebuiltConfig(t *testing.T, config map[string]any) map[string]a "kind": PostgresListTablespacesToolKind, "source": "my-instance", } + tools["list_pg_settings"] = map[string]any{ + "kind": PostgresListPGSettingsToolKind, + "source": "my-instance", + } + tools["list_database_stats"] = map[string]any{ + "kind": PostgresListDatabaseStatsToolKind, + "source": "my-instance", + } + config["tools"] = tools return config } diff --git a/tests/postgres/postgres_integration_test.go b/tests/postgres/postgres_integration_test.go index 79a97322f1..b870bdc1aa 100644 --- a/tests/postgres/postgres_integration_test.go +++ b/tests/postgres/postgres_integration_test.go @@ -160,4 +160,6 @@ func TestPostgres(t *testing.T) { tests.RunPostgresGetColumnCardinalityTest(t, ctx, pool) tests.RunPostgresListPublicationTablesTest(t, ctx, pool) tests.RunPostgresListTableSpacesTest(t) + tests.RunPostgresListPgSettingsTest(t, ctx, pool) + tests.RunPostgresListDatabaseStatsTest(t, ctx, pool) } diff --git a/tests/tool.go b/tests/tool.go index bd6713becb..a5bc40ed85 100644 --- a/tests/tool.go +++ b/tests/tool.go @@ -2272,6 +2272,241 @@ func RunPostgresListTableSpacesTest(t *testing.T) { } } +func RunPostgresListPgSettingsTest(t *testing.T, ctx context.Context, pool *pgxpool.Pool) { + targetSetting := "maintenance_work_mem" + var name, setting, unit, shortDesc, source, contextVal string + + // We query the raw pg_settings to get the data needed to reconstruct the logic + // defined in your listPgSettingQuery. + err := pool.QueryRow(ctx, ` + SELECT name, setting, unit, short_desc, source, context + FROM pg_settings + WHERE name = $1 + `, targetSetting).Scan(&name, &setting, &unit, &shortDesc, &source, &contextVal) + + if err != nil { + t.Fatalf("Setup failed: could not fetch postgres setting '%s': %v", targetSetting, err) + } + + // Replicate the SQL CASE logic for 'requires_restart' field + requiresRestart := "No" + switch contextVal { + case "postmaster": + requiresRestart = "Yes" + case "sighup": + requiresRestart = "No (Reload sufficient)" + } + + expectedObject := map[string]interface{}{ + "name": name, + "current_value": setting, + "unit": unit, + "short_desc": shortDesc, + "source": source, + "requires_restart": requiresRestart, + } + expectedJSON, _ := json.Marshal([]interface{}{expectedObject}) + + invokeTcs := []struct { + name string + requestBody io.Reader + wantStatusCode int + want string + }{ + { + name: "invoke list_pg_settings with specific setting", + requestBody: bytes.NewBuffer([]byte(fmt.Sprintf(`{"setting_name": "%s"}`, targetSetting))), + wantStatusCode: http.StatusOK, + want: string(expectedJSON), + }, + { + name: "invoke list_pg_settings with non-existent setting", + requestBody: bytes.NewBuffer([]byte(`{"setting_name": "non_existent_config_xyz"}`)), + wantStatusCode: http.StatusOK, + want: `null`, + }, + } + + for _, tc := range invokeTcs { + t.Run(tc.name, func(t *testing.T) { + const api = "http://127.0.0.1:5000/api/tool/list_pg_settings/invoke" + resp, body := 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(body)) + } + if tc.wantStatusCode != http.StatusOK { + return + } + + var bodyWrapper struct { + Result json.RawMessage `json:"result"` + } + if err := json.Unmarshal(body, &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, want any + if err := json.Unmarshal([]byte(resultString), &got); err != nil { + t.Fatalf("failed to unmarshal nested result string: %v", err) + } + if err := json.Unmarshal([]byte(tc.want), &want); err != nil { + t.Fatalf("failed to unmarshal want string: %v", err) + } + + if diff := cmp.Diff(want, got); diff != "" { + t.Errorf("Unexpected result (-want +got):\n%s", diff) + } + }) + } +} + +// RunPostgresDatabaseStatsTest tests the database_stats tool by comparing API results +// against a direct query to the database. +func RunPostgresListDatabaseStatsTest(t *testing.T, ctx context.Context, pool *pgxpool.Pool) { + dbName1 := "test_db_stats_1" + dbOwner1 := "test_user1" + dbName2 := "test_db_stats_2" + dbOwner2 := "test_user2" + + cleanup1 := setUpDatabase(t, ctx, pool, dbName1, dbOwner1) + defer cleanup1() + cleanup2 := setUpDatabase(t, ctx, pool, dbName2, dbOwner2) + defer cleanup2() + + requiredKeys := map[string]bool{ + "database_name": true, + "database_owner": true, + "default_tablespace": true, + "is_connectable": true, + } + + db1Want := map[string]interface{}{ + "database_name": dbName1, + "database_owner": dbOwner1, + "default_tablespace": "pg_default", + "is_connectable": true, + } + + db2Want := map[string]interface{}{ + "database_name": dbName2, + "database_owner": dbOwner2, + "default_tablespace": "pg_default", + "is_connectable": true, + } + + invokeTcs := []struct { + name string + requestBody io.Reader + wantStatusCode int + want []map[string]interface{} + }{ + { + name: "invoke database_stats filtering by specific database name", + requestBody: bytes.NewBuffer([]byte(fmt.Sprintf(`{"database_name": "%s"}`, dbName1))), + wantStatusCode: http.StatusOK, + want: []map[string]interface{}{db1Want}, + }, + { + name: "invoke database_stats filtering by specific owner", + requestBody: bytes.NewBuffer([]byte(fmt.Sprintf(`{"database_owner": "%s"}`, dbOwner2))), + wantStatusCode: http.StatusOK, + want: []map[string]interface{}{db2Want}, + }, + { + name: "filter by tablespace", + requestBody: bytes.NewBuffer([]byte(`{"default_tablespace": "pg_default"}`)), + wantStatusCode: http.StatusOK, + want: []map[string]interface{}{db1Want, db2Want}, + }, + { + name: "sort by size (desc)", + requestBody: bytes.NewBuffer([]byte(`{"sort_by": "size"}`)), + wantStatusCode: http.StatusOK, + want: []map[string]interface{}{db1Want, db2Want}, + }, + } + + for _, tc := range invokeTcs { + t.Run(tc.name, func(t *testing.T) { + const api = "http://127.0.0.1:5000/api/tool/list_database_stats/invoke" + resp, body := 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(body)) + } + + var bodyWrapper struct { + Result json.RawMessage `json:"result"` + } + if err := json.Unmarshal(body, &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]interface{} + if err := json.Unmarshal([]byte(resultString), &got); err != nil { + t.Fatalf("failed to unmarshal nested result string: %v", err) + } + + // Configuration for comparison + opts := []cmp.Option{ + // Ensure consistent order based on name for comparison + cmpopts.SortSlices(func(a, b map[string]interface{}) bool { + return a["database_name"].(string) < b["database_name"].(string) + }), + + // Ignore Volatile Keys which change in every run and only compare the keys in 'requiredKeys' + cmpopts.IgnoreMapEntries(func(key string, _ interface{}) bool { + return !requiredKeys[key] + }), + + // Ignore Irrelevant Databases + cmpopts.IgnoreSliceElements(func(v map[string]interface{}) bool { + name, ok := v["database_name"].(string) + if !ok { + return true + } + return name != dbName1 && name != dbName2 + }), + } + + if diff := cmp.Diff(tc.want, got, opts...); diff != "" { + t.Errorf("Unexpected result (-want +got):\n%s", diff) + } + }) + } +} + +func setUpDatabase(t *testing.T, ctx context.Context, pool *pgxpool.Pool, dbName, dbOwner string) func() { + _, err := pool.Exec(ctx, fmt.Sprintf("CREATE ROLE %s LOGIN PASSWORD 'password';", dbOwner)) + if err != nil { + _, _ = pool.Exec(ctx, fmt.Sprintf("DROP ROLE %s;", dbOwner)) + t.Fatalf("failed to create %s: %v", dbOwner, err) + } + _, err = pool.Exec(ctx, fmt.Sprintf("GRANT %s TO current_user;", dbOwner)) + if err != nil { + t.Fatalf("failed to grant %s to current_user: %v", dbOwner, err) + } + _, err = pool.Exec(ctx, fmt.Sprintf("CREATE DATABASE %s OWNER %s;", dbName, dbOwner)) + if err != nil { + t.Fatalf("failed to create %s: %v", dbName, err) + } + return func() { + _, _ = pool.Exec(ctx, fmt.Sprintf("DROP DATABASE IF EXISTS %s;", dbName)) + _, _ = pool.Exec(ctx, fmt.Sprintf("DROP ROLE IF EXISTS %s;", dbOwner)) + } +} + // RunMySQLListTablesTest run tests against the mysql-list-tables tool func RunMySQLListTablesTest(t *testing.T, databaseName, tableNameParam, tableNameAuth, expectedOwner string) { var ownerWant any