mirror of
https://github.com/googleapis/genai-toolbox.git
synced 2026-01-11 08:28:11 -05:00
feat(tools/postgres): add list_pg_settings, list_database_stats tools for postgres (#2030)
## 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: <img width="1526" height="803" alt="Screenshot 2025-11-25 at 10 19 48 AM" src="https://github.com/user-attachments/assets/73634b9b-4936-4bf0-a94b-6b31fe3642a1" /> <img width="1064" height="715" alt="Screenshot 2025-11-25 at 10 27 19 AM" src="https://github.com/user-attachments/assets/36c13585-27e4-4294-b451-1c1a963c0d6c" /> list_database_stats: <img width="1511" height="779" alt="Screenshot 2025-11-25 at 10 21 12 AM" src="https://github.com/user-attachments/assets/d283e018-ea81-427d-b1b4-7aaf79b9696b" /> <img width="1017" height="506" alt="Screenshot 2025-11-25 at 10 27 47 AM" src="https://github.com/user-attachments/assets/47b72bd7-7114-4f2a-8a9d-cecc80bf47e9" /> ## 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 <akitsch@google.com> Co-authored-by: Wenxin Du <117315983+duwenxin99@users.noreply.github.com>
This commit is contained in:
Srividya Reddy
@@ -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"
|
||||
|
||||
@@ -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"},
|
||||
},
|
||||
},
|
||||
},
|
||||
|
||||
@@ -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
|
||||
|
||||
|
||||
@@ -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/)
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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/)
|
||||
|
||||
@@ -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. |
|
||||
@@ -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. |
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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"
|
||||
}
|
||||
@@ -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)
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
}
|
||||
@@ -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"
|
||||
}
|
||||
@@ -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)
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
}
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
}
|
||||
|
||||
@@ -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)
|
||||
}
|
||||
|
||||
235
tests/tool.go
235
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
|
||||
|
||||
Reference in New Issue
Block a user