mirror of
https://github.com/googleapis/genai-toolbox.git
synced 2026-01-11 16:38:15 -05:00
feat(tools/postgres): add list_triggers, database_overview tools for postgres (#1912)
## Description Adds the following tools for Postgres: (1) list_triggers: Lists triggers in the database. . (2) database_overview: Fetches the current state of the PostgreSQL server. list_triggers: <img width="1712" height="703" alt="Screenshot 2025-11-09 at 8 16 53 PM" src="https://github.com/user-attachments/assets/1974e303-b559-4efc-b129-444ba97c7715" /> <img width="874" height="513" alt="Screenshot 2025-11-09 at 8 19 43 PM" src="https://github.com/user-attachments/assets/59ddcd15-224b-4e9a-906d-ec2645835873" /> database_overview: <img width="1521" height="683" alt="Screenshot 2025-11-09 at 8 53 03 PM" src="https://github.com/user-attachments/assets/4ae86e74-aa78-410c-a9cc-f33ae3268fb6" /> <img width="850" height="241" alt="Screenshot 2025-11-09 at 8 49 53 PM" src="https://github.com/user-attachments/assets/abae2c7a-5f3e-4433-86de-3606e3298ec5" /> > Should include a concise description of the changes (bug or feature), it's > impact, along with a summary of the solution ## PR Checklist > Thank you for opening a Pull Request! Before submitting your PR, there are a > few things you can do to make sure it goes smoothly: - [x] Make sure you reviewed [CONTRIBUTING.md](https://github.com/googleapis/genai-toolbox/blob/main/CONTRIBUTING.md) - [x] Make sure to open an issue as a [bug/issue](https://github.com/googleapis/genai-toolbox/issues/new/choose) before writing your code! That way we can discuss the change, evaluate designs, and agree on the general idea - [x] Ensure the tests and linter pass - [x] Code coverage does not decrease (if any source code was changed) - [x] Appropriate docs were updated (if necessary) - [x] Make sure to add `!` if this involve a breaking change 🛠️ Fixes #1738
This commit is contained in:
@@ -176,12 +176,14 @@ import (
|
||||
_ "github.com/googleapis/genai-toolbox/internal/tools/oceanbase/oceanbasesql"
|
||||
_ "github.com/googleapis/genai-toolbox/internal/tools/oracle/oracleexecutesql"
|
||||
_ "github.com/googleapis/genai-toolbox/internal/tools/oracle/oraclesql"
|
||||
_ "github.com/googleapis/genai-toolbox/internal/tools/postgres/postgresdatabaseoverview"
|
||||
_ "github.com/googleapis/genai-toolbox/internal/tools/postgres/postgresexecutesql"
|
||||
_ "github.com/googleapis/genai-toolbox/internal/tools/postgres/postgreslistactivequeries"
|
||||
_ "github.com/googleapis/genai-toolbox/internal/tools/postgres/postgreslistavailableextensions"
|
||||
_ "github.com/googleapis/genai-toolbox/internal/tools/postgres/postgreslistinstalledextensions"
|
||||
_ "github.com/googleapis/genai-toolbox/internal/tools/postgres/postgreslistschemas"
|
||||
_ "github.com/googleapis/genai-toolbox/internal/tools/postgres/postgreslisttables"
|
||||
_ "github.com/googleapis/genai-toolbox/internal/tools/postgres/postgreslisttriggers"
|
||||
_ "github.com/googleapis/genai-toolbox/internal/tools/postgres/postgreslistviews"
|
||||
_ "github.com/googleapis/genai-toolbox/internal/tools/postgres/postgressql"
|
||||
_ "github.com/googleapis/genai-toolbox/internal/tools/redis"
|
||||
|
||||
@@ -1477,7 +1477,7 @@ func TestPrebuiltTools(t *testing.T) {
|
||||
wantToolset: server.ToolsetConfigs{
|
||||
"alloydb_postgres_database_tools": tools.ToolsetConfig{
|
||||
Name: "alloydb_postgres_database_tools",
|
||||
ToolNames: []string{"execute_sql", "list_tables", "list_active_queries", "list_available_extensions", "list_installed_extensions", "list_autovacuum_configurations", "list_memory_configurations", "list_top_bloated_tables", "list_replication_slots", "list_invalid_indexes", "get_query_plan", "list_views", "list_schemas"},
|
||||
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"},
|
||||
},
|
||||
},
|
||||
},
|
||||
@@ -1507,7 +1507,7 @@ func TestPrebuiltTools(t *testing.T) {
|
||||
wantToolset: server.ToolsetConfigs{
|
||||
"cloud_sql_postgres_database_tools": tools.ToolsetConfig{
|
||||
Name: "cloud_sql_postgres_database_tools",
|
||||
ToolNames: []string{"execute_sql", "list_tables", "list_active_queries", "list_available_extensions", "list_installed_extensions", "list_autovacuum_configurations", "list_memory_configurations", "list_top_bloated_tables", "list_replication_slots", "list_invalid_indexes", "get_query_plan", "list_views", "list_schemas"},
|
||||
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"},
|
||||
},
|
||||
},
|
||||
},
|
||||
@@ -1607,7 +1607,7 @@ func TestPrebuiltTools(t *testing.T) {
|
||||
wantToolset: server.ToolsetConfigs{
|
||||
"postgres_database_tools": tools.ToolsetConfig{
|
||||
Name: "postgres_database_tools",
|
||||
ToolNames: []string{"execute_sql", "list_tables", "list_active_queries", "list_available_extensions", "list_installed_extensions", "list_autovacuum_configurations", "list_memory_configurations", "list_top_bloated_tables", "list_replication_slots", "list_invalid_indexes", "get_query_plan", "list_views", "list_schemas"},
|
||||
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"},
|
||||
},
|
||||
},
|
||||
},
|
||||
|
||||
@@ -46,6 +46,8 @@ details on how to connect your AI tools (IDEs) to databases via Toolbox and MCP.
|
||||
* `list_views`: Lists views in the database from pg_views with a default
|
||||
limit of 50 rows. Returns schemaname, viewname and the ownername.
|
||||
* `list_schemas`: Lists schemas in the database.
|
||||
* `database_overview`: Fetches the current state of the PostgreSQL server.
|
||||
* `list_triggers`: Lists triggers in the database.
|
||||
|
||||
## AlloyDB Postgres Admin
|
||||
|
||||
@@ -216,6 +218,8 @@ details on how to connect your AI tools (IDEs) to databases via Toolbox and MCP.
|
||||
* `list_views`: Lists views in the database from pg_views with a default
|
||||
limit of 50 rows. Returns schemaname, viewname and the ownername.
|
||||
* `list_schemas`: Lists schemas in the database.
|
||||
* `database_overview`: Fetches the current state of the PostgreSQL server.
|
||||
* `list_triggers`: Lists triggers in the database.
|
||||
|
||||
## Cloud SQL for PostgreSQL Observability
|
||||
|
||||
@@ -513,6 +517,8 @@ details on how to connect your AI tools (IDEs) to databases via Toolbox and MCP.
|
||||
* `list_views`: Lists views in the database from pg_views with a default
|
||||
limit of 50 rows. Returns schemaname, viewname and the ownername.
|
||||
* `list_schemas`: Lists schemas in the database.
|
||||
* `database_overview`: Fetches the current state of the PostgreSQL server.
|
||||
* `list_triggers`: Lists triggers in the database.
|
||||
|
||||
## Google Cloud Serverless for Apache Spark
|
||||
|
||||
|
||||
@@ -51,6 +51,12 @@ cluster][alloydb-free-trial].
|
||||
- [`postgres-list-schemas`](../tools/postgres/postgres-list-schemas.md)
|
||||
List schemas in an AlloyDB for PostgreSQL database.
|
||||
|
||||
- [`postgres-database-overview`](../tools/postgres/postgres-database-overview.md)
|
||||
Fetches the current state of the PostgreSQL server.
|
||||
|
||||
- [`postgres-list-triggers`](../tools/postgres/postgres-list-triggers.md)
|
||||
List triggers in an AlloyDB for PostgreSQL database.
|
||||
|
||||
### Pre-built Configurations
|
||||
|
||||
- [AlloyDB using MCP](https://googleapis.github.io/genai-toolbox/how-to/connect-ide/alloydb_pg_mcp/)
|
||||
|
||||
@@ -47,6 +47,12 @@ to a database by following these instructions][csql-pg-quickstart].
|
||||
- [`postgres-list-schemas`](../tools/postgres/postgres-list-schemas.md)
|
||||
List schemas in a PostgreSQL database.
|
||||
|
||||
- [`postgres-database-overview`](../tools/postgres/postgres-database-overview.md)
|
||||
Fetches the current state of the PostgreSQL server.
|
||||
|
||||
- [`postgres-list-triggers`](../tools/postgres/postgres-list-triggers.md)
|
||||
List triggers in a PostgreSQL database.
|
||||
|
||||
### Pre-built Configurations
|
||||
|
||||
- [Cloud SQL for Postgres using
|
||||
|
||||
@@ -41,6 +41,12 @@ reputation for reliability, feature robustness, and performance.
|
||||
- [`postgres-list-schemas`](../tools/postgres/postgres-list-views.md)
|
||||
List schemas in a PostgreSQL database.
|
||||
|
||||
- [`postgres-database-overview`](../tools/postgres/postgres-database-overview.md)
|
||||
Fetches the current state of the PostgreSQL server.
|
||||
|
||||
- [`postgres-list-triggers`](../tools/postgres/postgres-list-triggers.md)
|
||||
List triggers in a PostgreSQL database.
|
||||
|
||||
### Pre-built Configurations
|
||||
|
||||
- [PostgreSQL using MCP](https://googleapis.github.io/genai-toolbox/how-to/connect-ide/postgres_mcp/)
|
||||
|
||||
@@ -0,0 +1,51 @@
|
||||
---
|
||||
title: "postgres-database-overview"
|
||||
type: docs
|
||||
weight: 1
|
||||
description: >
|
||||
The "postgres-database-overview" fetches the current state of the PostgreSQL server.
|
||||
aliases:
|
||||
- /resources/tools/postgres-database-overview
|
||||
---
|
||||
|
||||
## About
|
||||
|
||||
The `postgres-database-overview` fetches the current state of the PostgreSQL server. 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-database-overview` fetches the current state of the PostgreSQL server This tool does not take any input parameters.
|
||||
|
||||
## Example
|
||||
|
||||
```yaml
|
||||
tools:
|
||||
database_overview:
|
||||
kind: postgres-database-overview
|
||||
source: cloudsql-pg-source
|
||||
description: |
|
||||
fetches the current state of the PostgreSQL server. It returns the postgres version, whether it's a replica, uptime duration, maximum connection limit, number of current connections, number of active connections and the percentage of connections in use.
|
||||
```
|
||||
|
||||
The response is a JSON object with the following elements:
|
||||
```json
|
||||
{
|
||||
"pg_version": "PostgreSQL server version string",
|
||||
"is_replica": "boolean indicating if the instance is in recovery mode",
|
||||
"uptime": "interval string representing the total server uptime",
|
||||
"max_connections": "integer maximum number of allowed connections",
|
||||
"current_connections": "integer number of current connections",
|
||||
"active_connections": "integer number of currently active connections",
|
||||
"pct_connections_used": "float percentage of max_connections currently in use"
|
||||
}
|
||||
```
|
||||
|
||||
## Reference
|
||||
|
||||
| **field** | **type** | **required** | **description** |
|
||||
|-------------|:--------:|:-------------:|------------------------------------------------------|
|
||||
| kind | string | true | Must be "postgres-database-overview". |
|
||||
| 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. |
|
||||
60
docs/en/resources/tools/postgres/postgres-list-triggers.md
Normal file
60
docs/en/resources/tools/postgres/postgres-list-triggers.md
Normal file
@@ -0,0 +1,60 @@
|
||||
---
|
||||
title: "postgres-list-triggers"
|
||||
type: docs
|
||||
weight: 1
|
||||
description: >
|
||||
The "postgres-list-triggers" tool lists triggers in a Postgres database.
|
||||
aliases:
|
||||
- /resources/tools/postgres-list-triggers
|
||||
---
|
||||
|
||||
## About
|
||||
|
||||
The `postgres-list-triggers` tool lists available non-internal triggers in the database. It's compatible with any of the following sources:
|
||||
|
||||
- [alloydb-postgres](../../sources/alloydb-pg.md)
|
||||
- [cloud-sql-postgres](../../sources/cloud-sql-pg.md)
|
||||
- [postgres](../../sources/postgres.md)
|
||||
|
||||
`postgres-list-triggers` lists detailed information as JSON for triggers. The tool takes the following input parameters:
|
||||
|
||||
- `trigger_name` (optional): A text to filter results by trigger name. The input is used within a LIKE clause. Default: `""`
|
||||
- `schema_name` (optional): A text to filter results by schema name. The input is used within a LIKE clause. Default: `""`
|
||||
- `table_name` (optional): A text to filter results by table name. The input is used within a LIKE clause. Default: `""`
|
||||
- `limit` (optional): The maximum number of triggers to return. Default: `50`
|
||||
|
||||
|
||||
## Example
|
||||
|
||||
```yaml
|
||||
```yaml
|
||||
tools:
|
||||
list_triggers:
|
||||
kind: postgres-list-triggers
|
||||
source: postgres-source
|
||||
description: |
|
||||
Lists all non-internal triggers in a database. Returns trigger name, schema name, table name, wether its enabled or disabled, timing (e.g BEFORE/AFTER of the event), the events that cause the trigger to fire such as INSERT, UPDATE, or DELETE, whether the trigger activates per ROW or per STATEMENT, the handler function executed by the trigger and full definition.
|
||||
```
|
||||
|
||||
The response is a json array with the following elements:
|
||||
```json
|
||||
{
|
||||
"trigger_name": "trigger name",
|
||||
"schema_name": "schema name",
|
||||
"table_name": "table name",
|
||||
"status": "Whether the trigger is currently active (ENABLED, DISABLED, REPLICA, ALWAYS).",
|
||||
"timing": "When it runs relative to the event (BEFORE, AFTER, INSTEAD OF).",
|
||||
"events": "The specific operations that fire it (INSERT, UPDATE, DELETE, TRUNCATE)",
|
||||
"activation_level": "Granularity of execution (ROW vs STATEMENT).",
|
||||
"function_name": "The function it executes",
|
||||
"definition": "Full SQL definition of the trigger"
|
||||
}
|
||||
```
|
||||
|
||||
## Reference
|
||||
|
||||
| **field** | **type** | **required** | **description** |
|
||||
|-------------|:--------:|:-------------:|------------------------------------------------------|
|
||||
| kind | string | true | Must be "postgres-list-triggers". |
|
||||
| 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. |
|
||||
@@ -164,6 +164,14 @@ tools:
|
||||
kind: postgres-list-schemas
|
||||
source: alloydb-pg-source
|
||||
|
||||
database_overview:
|
||||
kind: postgres-database-overview
|
||||
source: alloydb-pg-source
|
||||
|
||||
list_triggers:
|
||||
kind: postgres-list-triggers
|
||||
source: alloydb-pg-source
|
||||
|
||||
toolsets:
|
||||
alloydb_postgres_database_tools:
|
||||
- execute_sql
|
||||
@@ -179,3 +187,5 @@ toolsets:
|
||||
- get_query_plan
|
||||
- list_views
|
||||
- list_schemas
|
||||
- database_overview
|
||||
- list_triggers
|
||||
|
||||
@@ -163,6 +163,14 @@ tools:
|
||||
kind: postgres-list-schemas
|
||||
source: cloudsql-pg-source
|
||||
|
||||
database_overview:
|
||||
kind: postgres-database-overview
|
||||
source: cloudsql-pg-source
|
||||
|
||||
list_triggers:
|
||||
kind: postgres-list-triggers
|
||||
source: cloudsql-pg-source
|
||||
|
||||
toolsets:
|
||||
cloud_sql_postgres_database_tools:
|
||||
- execute_sql
|
||||
@@ -178,3 +186,5 @@ toolsets:
|
||||
- get_query_plan
|
||||
- list_views
|
||||
- list_schemas
|
||||
- database_overview
|
||||
- list_triggers
|
||||
|
||||
@@ -162,6 +162,14 @@ tools:
|
||||
kind: postgres-list-schemas
|
||||
source: postgresql-source
|
||||
|
||||
database_overview:
|
||||
kind: postgres-database-overview
|
||||
source: postgresql-source
|
||||
|
||||
list_triggers:
|
||||
kind: postgres-list-triggers
|
||||
source: postgresql-source
|
||||
|
||||
toolsets:
|
||||
postgres_database_tools:
|
||||
- execute_sql
|
||||
@@ -177,3 +185,5 @@ toolsets:
|
||||
- get_query_plan
|
||||
- list_views
|
||||
- list_schemas
|
||||
- database_overview
|
||||
- list_triggers
|
||||
|
||||
@@ -0,0 +1,180 @@
|
||||
// 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 postgresdatabaseoverview
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
|
||||
yaml "github.com/goccy/go-yaml"
|
||||
"github.com/googleapis/genai-toolbox/internal/sources"
|
||||
"github.com/googleapis/genai-toolbox/internal/sources/alloydbpg"
|
||||
"github.com/googleapis/genai-toolbox/internal/sources/cloudsqlpg"
|
||||
"github.com/googleapis/genai-toolbox/internal/sources/postgres"
|
||||
"github.com/googleapis/genai-toolbox/internal/tools"
|
||||
"github.com/jackc/pgx/v5/pgxpool"
|
||||
)
|
||||
|
||||
const kind string = "postgres-database-overview"
|
||||
|
||||
const databaseOverviewStatement = `
|
||||
SELECT
|
||||
current_setting('server_version') AS pg_version,
|
||||
pg_is_in_recovery() AS is_replica,
|
||||
(now() - pg_postmaster_start_time())::TEXT AS uptime,
|
||||
current_setting('max_connections')::int AS max_connections,
|
||||
(SELECT count(*) FROM pg_stat_activity) AS current_connections,
|
||||
(SELECT count(*) FROM pg_stat_activity WHERE state = 'active') AS active_connections,
|
||||
round(
|
||||
(100.0 * (SELECT count(*) FROM pg_stat_activity) / current_setting('max_connections')::int),
|
||||
2
|
||||
) AS pct_connections_used;
|
||||
`
|
||||
|
||||
func init() {
|
||||
if !tools.Register(kind, newConfig) {
|
||||
panic(fmt.Sprintf("tool kind %q already registered", kind))
|
||||
}
|
||||
}
|
||||
|
||||
func newConfig(ctx context.Context, name string, decoder *yaml.Decoder) (tools.ToolConfig, error) {
|
||||
actual := Config{Name: name}
|
||||
if err := decoder.DecodeContext(ctx, &actual); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return actual, nil
|
||||
}
|
||||
|
||||
type compatibleSource interface {
|
||||
PostgresPool() *pgxpool.Pool
|
||||
}
|
||||
|
||||
// validate compatible sources are still compatible
|
||||
var _ compatibleSource = &alloydbpg.Source{}
|
||||
var _ compatibleSource = &cloudsqlpg.Source{}
|
||||
var _ compatibleSource = &postgres.Source{}
|
||||
|
||||
var compatibleSources = [...]string{alloydbpg.SourceKind, cloudsqlpg.SourceKind, postgres.SourceKind}
|
||||
|
||||
type Config struct {
|
||||
Name string `yaml:"name" validate:"required"`
|
||||
Kind string `yaml:"kind" validate:"required"`
|
||||
Source string `yaml:"source" validate:"required"`
|
||||
Description string `yaml:"description"`
|
||||
AuthRequired []string `yaml:"authRequired"`
|
||||
}
|
||||
|
||||
// validate interface
|
||||
var _ tools.ToolConfig = Config{}
|
||||
|
||||
func (cfg Config) ToolConfigKind() string {
|
||||
return kind
|
||||
}
|
||||
|
||||
func (cfg Config) Initialize(srcs map[string]sources.Source) (tools.Tool, error) {
|
||||
// verify source exists
|
||||
rawS, ok := srcs[cfg.Source]
|
||||
if !ok {
|
||||
return nil, fmt.Errorf("no source named %q configured", cfg.Source)
|
||||
}
|
||||
|
||||
// verify the source is compatible
|
||||
s, ok := rawS.(compatibleSource)
|
||||
if !ok {
|
||||
return nil, fmt.Errorf("invalid source for %q tool: source kind must be one of %q", kind, compatibleSources)
|
||||
}
|
||||
|
||||
allParameters := tools.Parameters{}
|
||||
description := cfg.Description
|
||||
if description == "" {
|
||||
description = "Fetches the current state of the PostgreSQL server, returning the version, whether it's a replica, uptime duration, maximum connection limit, number of current connections, number of active connections, and the percentage of connections in use."
|
||||
}
|
||||
mcpManifest := tools.GetMcpManifest(cfg.Name, description, cfg.AuthRequired, allParameters)
|
||||
|
||||
// finish tool setup
|
||||
return Tool{
|
||||
name: cfg.Name,
|
||||
kind: kind,
|
||||
authRequired: cfg.AuthRequired,
|
||||
allParams: allParameters,
|
||||
pool: s.PostgresPool(),
|
||||
manifest: tools.Manifest{
|
||||
Description: cfg.Description,
|
||||
Parameters: allParameters.Manifest(),
|
||||
AuthRequired: cfg.AuthRequired,
|
||||
},
|
||||
mcpManifest: mcpManifest,
|
||||
}, nil
|
||||
}
|
||||
|
||||
// validate interface
|
||||
var _ tools.Tool = Tool{}
|
||||
|
||||
type Tool struct {
|
||||
name string `yaml:"name"`
|
||||
kind string `yaml:"kind"`
|
||||
authRequired []string `yaml:"authRequired"`
|
||||
allParams tools.Parameters `yaml:"allParams"`
|
||||
pool *pgxpool.Pool
|
||||
manifest tools.Manifest
|
||||
mcpManifest tools.McpManifest
|
||||
}
|
||||
|
||||
func (t Tool) Invoke(ctx context.Context, params tools.ParamValues, accessToken tools.AccessToken) (any, error) {
|
||||
sliceParams := params.AsSlice()
|
||||
|
||||
results, err := t.pool.Query(ctx, databaseOverviewStatement, sliceParams...)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("unable to execute query: %w", err)
|
||||
}
|
||||
defer results.Close()
|
||||
|
||||
fields := results.FieldDescriptions()
|
||||
var out []map[string]any
|
||||
|
||||
for results.Next() {
|
||||
values, err := results.Values()
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("unable to parse row: %w", err)
|
||||
}
|
||||
rowMap := make(map[string]any)
|
||||
for i, field := range fields {
|
||||
rowMap[string(field.Name)] = values[i]
|
||||
}
|
||||
out = append(out, rowMap)
|
||||
}
|
||||
|
||||
return out, nil
|
||||
}
|
||||
|
||||
func (t Tool) ParseParams(data map[string]any, claims map[string]map[string]any) (tools.ParamValues, error) {
|
||||
return tools.ParseParams(t.allParams, data, claims)
|
||||
}
|
||||
|
||||
func (t Tool) Manifest() tools.Manifest {
|
||||
return t.manifest
|
||||
}
|
||||
|
||||
func (t Tool) McpManifest() tools.McpManifest {
|
||||
return t.mcpManifest
|
||||
}
|
||||
|
||||
func (t Tool) Authorized(verifiedAuthServices []string) bool {
|
||||
return tools.IsAuthorized(t.authRequired, verifiedAuthServices)
|
||||
}
|
||||
|
||||
func (t Tool) RequiresClientAuthorization() bool {
|
||||
return false
|
||||
}
|
||||
@@ -0,0 +1,95 @@
|
||||
// Copyright 2025 Google LLC
|
||||
//
|
||||
// Licensed under the Apache License, Version 2.0 (the "License");
|
||||
// you may not use this file except in compliance with the License.
|
||||
// You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software
|
||||
// distributed under the License is distributed on an "AS IS" BASIS,
|
||||
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
// See the License for the specific language governing permissions and
|
||||
// limitations under the License.
|
||||
|
||||
package postgresdatabaseoverview_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/postgresdatabaseoverview"
|
||||
)
|
||||
|
||||
func TestParseFromYamlPostgresDatabaseOverview(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-database-overview
|
||||
source: my-postgres-instance
|
||||
description: some description
|
||||
authRequired:
|
||||
- my-google-auth-service
|
||||
- other-auth-service
|
||||
`,
|
||||
want: server.ToolConfigs{
|
||||
"example_tool": postgresdatabaseoverview.Config{
|
||||
Name: "example_tool",
|
||||
Kind: "postgres-database-overview",
|
||||
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-database-overview
|
||||
source: my-postgres-instance
|
||||
description: some description
|
||||
`,
|
||||
want: server.ToolConfigs{
|
||||
"example_tool": postgresdatabaseoverview.Config{
|
||||
Name: "example_tool",
|
||||
Kind: "postgres-database-overview",
|
||||
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,217 @@
|
||||
// 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 postgreslisttriggers
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
|
||||
yaml "github.com/goccy/go-yaml"
|
||||
"github.com/googleapis/genai-toolbox/internal/sources"
|
||||
"github.com/googleapis/genai-toolbox/internal/sources/alloydbpg"
|
||||
"github.com/googleapis/genai-toolbox/internal/sources/cloudsqlpg"
|
||||
"github.com/googleapis/genai-toolbox/internal/sources/postgres"
|
||||
"github.com/googleapis/genai-toolbox/internal/tools"
|
||||
"github.com/jackc/pgx/v5/pgxpool"
|
||||
)
|
||||
|
||||
const kind string = "postgres-list-triggers"
|
||||
|
||||
const listTriggersStatement = `
|
||||
WITH
|
||||
trigger_list AS (
|
||||
SELECT
|
||||
t.tgname AS trigger_name,
|
||||
n.nspname AS schema_name,
|
||||
c.relname AS table_name,
|
||||
CASE t.tgenabled
|
||||
WHEN 'O' THEN 'ENABLED'
|
||||
WHEN 'D' THEN 'DISABLED'
|
||||
WHEN 'R' THEN 'REPLICA'
|
||||
WHEN 'A' THEN 'ALWAYS'
|
||||
END AS status,
|
||||
CASE
|
||||
WHEN (t.tgtype::int & 2) = 2 THEN 'BEFORE'
|
||||
WHEN (t.tgtype::int & 64) = 64 THEN 'INSTEAD OF'
|
||||
ELSE 'AFTER'
|
||||
END AS timing,
|
||||
concat_ws(
|
||||
', ',
|
||||
CASE WHEN (t.tgtype::int & 4) = 4 THEN 'INSERT' END,
|
||||
CASE WHEN (t.tgtype::int & 16) = 16 THEN 'UPDATE' END,
|
||||
CASE WHEN (t.tgtype::int & 8) = 8 THEN 'DELETE' END,
|
||||
CASE WHEN (t.tgtype::int & 32) = 32 THEN 'TRUNCATE' END) AS events,
|
||||
CASE WHEN (t.tgtype::int & 1) = 1 THEN 'ROW' ELSE 'STATEMENT' END AS activation_level,
|
||||
p.proname AS function_name,
|
||||
pg_get_triggerdef(t.oid) AS definition
|
||||
FROM pg_trigger t
|
||||
JOIN pg_class c
|
||||
ON t.tgrelid = c.oid
|
||||
JOIN pg_namespace n
|
||||
ON c.relnamespace = n.oid
|
||||
LEFT JOIN pg_proc p
|
||||
ON t.tgfoid = p.oid
|
||||
WHERE NOT t.tgisinternal
|
||||
)
|
||||
SELECT *
|
||||
FROM trigger_list
|
||||
WHERE
|
||||
($1::text IS NULL OR trigger_name LIKE '%' || $1::text || '%')
|
||||
AND ($2::text IS NULL OR schema_name LIKE '%' || $2::text || '%')
|
||||
AND ($3::text IS NULL OR table_name LIKE '%' || $3::text || '%')
|
||||
ORDER BY schema_name, table_name, trigger_name
|
||||
LIMIT COALESCE($4::int, 50);
|
||||
`
|
||||
|
||||
func init() {
|
||||
if !tools.Register(kind, newConfig) {
|
||||
panic(fmt.Sprintf("tool kind %q already registered", kind))
|
||||
}
|
||||
}
|
||||
|
||||
func newConfig(ctx context.Context, name string, decoder *yaml.Decoder) (tools.ToolConfig, error) {
|
||||
actual := Config{Name: name}
|
||||
if err := decoder.DecodeContext(ctx, &actual); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return actual, nil
|
||||
}
|
||||
|
||||
type compatibleSource interface {
|
||||
PostgresPool() *pgxpool.Pool
|
||||
}
|
||||
|
||||
// validate compatible sources are still compatible
|
||||
var _ compatibleSource = &alloydbpg.Source{}
|
||||
var _ compatibleSource = &cloudsqlpg.Source{}
|
||||
var _ compatibleSource = &postgres.Source{}
|
||||
|
||||
var compatibleSources = [...]string{alloydbpg.SourceKind, cloudsqlpg.SourceKind, postgres.SourceKind}
|
||||
|
||||
type Config struct {
|
||||
Name string `yaml:"name" validate:"required"`
|
||||
Kind string `yaml:"kind" validate:"required"`
|
||||
Source string `yaml:"source" validate:"required"`
|
||||
Description string `yaml:"description"`
|
||||
AuthRequired []string `yaml:"authRequired"`
|
||||
}
|
||||
|
||||
// validate interface
|
||||
var _ tools.ToolConfig = Config{}
|
||||
|
||||
func (cfg Config) ToolConfigKind() string {
|
||||
return kind
|
||||
}
|
||||
|
||||
func (cfg Config) Initialize(srcs map[string]sources.Source) (tools.Tool, error) {
|
||||
// verify source exists
|
||||
rawS, ok := srcs[cfg.Source]
|
||||
if !ok {
|
||||
return nil, fmt.Errorf("no source named %q configured", cfg.Source)
|
||||
}
|
||||
|
||||
// verify the source is compatible
|
||||
s, ok := rawS.(compatibleSource)
|
||||
if !ok {
|
||||
return nil, fmt.Errorf("invalid source for %q tool: source kind must be one of %q", kind, compatibleSources)
|
||||
}
|
||||
|
||||
allParameters := tools.Parameters{
|
||||
tools.NewStringParameterWithDefault("trigger_name", "", "Optional: A specific trigger name pattern to search for."),
|
||||
tools.NewStringParameterWithDefault("schema_name", "", "Optional: A specific schema name pattern to search for."),
|
||||
tools.NewStringParameterWithDefault("table_name", "", "Optional: A specific table name pattern to search for."),
|
||||
tools.NewIntParameterWithDefault("limit", 50, "Optional: The maximum number of rows to return."),
|
||||
}
|
||||
description := cfg.Description
|
||||
if description == "" {
|
||||
description = "Lists all non-internal triggers in a database. Returns trigger name, schema name, table name, whether its enabled or disabled, timing (e.g BEFORE/AFTER of the event), the events that cause the trigger to fire such as INSERT, UPDATE, or DELETE, whether the trigger activates per ROW or per STATEMENT, the handler function executed by the trigger and full definition."
|
||||
}
|
||||
mcpManifest := tools.GetMcpManifest(cfg.Name, description, cfg.AuthRequired, allParameters)
|
||||
|
||||
// finish tool setup
|
||||
return Tool{
|
||||
name: cfg.Name,
|
||||
kind: kind,
|
||||
authRequired: cfg.AuthRequired,
|
||||
allParams: allParameters,
|
||||
pool: s.PostgresPool(),
|
||||
manifest: tools.Manifest{
|
||||
Description: cfg.Description,
|
||||
Parameters: allParameters.Manifest(),
|
||||
AuthRequired: cfg.AuthRequired,
|
||||
},
|
||||
mcpManifest: mcpManifest,
|
||||
}, nil
|
||||
}
|
||||
|
||||
// validate interface
|
||||
var _ tools.Tool = Tool{}
|
||||
|
||||
type Tool struct {
|
||||
name string `yaml:"name"`
|
||||
kind string `yaml:"kind"`
|
||||
authRequired []string `yaml:"authRequired"`
|
||||
allParams tools.Parameters `yaml:"allParams"`
|
||||
pool *pgxpool.Pool
|
||||
manifest tools.Manifest
|
||||
mcpManifest tools.McpManifest
|
||||
}
|
||||
|
||||
func (t Tool) Invoke(ctx context.Context, params tools.ParamValues, accessToken tools.AccessToken) (any, error) {
|
||||
sliceParams := params.AsSlice()
|
||||
|
||||
results, err := t.pool.Query(ctx, listTriggersStatement, sliceParams...)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("unable to execute query: %w", err)
|
||||
}
|
||||
defer results.Close()
|
||||
|
||||
fields := results.FieldDescriptions()
|
||||
var out []map[string]any
|
||||
|
||||
for results.Next() {
|
||||
values, err := results.Values()
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("unable to parse row: %w", err)
|
||||
}
|
||||
rowMap := make(map[string]any)
|
||||
for i, field := range fields {
|
||||
rowMap[string(field.Name)] = values[i]
|
||||
}
|
||||
out = append(out, rowMap)
|
||||
}
|
||||
|
||||
return out, nil
|
||||
}
|
||||
|
||||
func (t Tool) ParseParams(data map[string]any, claims map[string]map[string]any) (tools.ParamValues, error) {
|
||||
return tools.ParseParams(t.allParams, data, claims)
|
||||
}
|
||||
|
||||
func (t Tool) Manifest() tools.Manifest {
|
||||
return t.manifest
|
||||
}
|
||||
|
||||
func (t Tool) McpManifest() tools.McpManifest {
|
||||
return t.mcpManifest
|
||||
}
|
||||
|
||||
func (t Tool) Authorized(verifiedAuthServices []string) bool {
|
||||
return tools.IsAuthorized(t.authRequired, verifiedAuthServices)
|
||||
}
|
||||
|
||||
func (t Tool) RequiresClientAuthorization() bool {
|
||||
return false
|
||||
}
|
||||
@@ -0,0 +1,95 @@
|
||||
// Copyright 2025 Google LLC
|
||||
//
|
||||
// Licensed under the Apache License, Version 2.0 (the "License");
|
||||
// you may not use this file except in compliance with the License.
|
||||
// You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software
|
||||
// distributed under the License is distributed on an "AS IS" BASIS,
|
||||
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
// See the License for the specific language governing permissions and
|
||||
// limitations under the License.
|
||||
|
||||
package postgreslisttriggers_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/postgreslisttriggers"
|
||||
)
|
||||
|
||||
func TestParseFromYamlPostgreslistTriggers(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-triggers
|
||||
source: my-postgres-instance
|
||||
description: some description
|
||||
authRequired:
|
||||
- my-google-auth-service
|
||||
- other-auth-service
|
||||
`,
|
||||
want: server.ToolConfigs{
|
||||
"example_tool": postgreslisttriggers.Config{
|
||||
Name: "example_tool",
|
||||
Kind: "postgres-list-triggers",
|
||||
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-triggers
|
||||
source: my-postgres-instance
|
||||
description: some description
|
||||
`,
|
||||
want: server.ToolConfigs{
|
||||
"example_tool": postgreslisttriggers.Config{
|
||||
Name: "example_tool",
|
||||
Kind: "postgres-list-triggers",
|
||||
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)
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
}
|
||||
@@ -186,6 +186,8 @@ func TestAlloyDBPgToolEndpoints(t *testing.T) {
|
||||
tests.RunPostgresListActiveQueriesTest(t, ctx, pool)
|
||||
tests.RunPostgresListAvailableExtensionsTest(t)
|
||||
tests.RunPostgresListInstalledExtensionsTest(t)
|
||||
tests.RunPostgresDatabaseOverviewTest(t, ctx, pool)
|
||||
tests.RunPostgresListTriggersTest(t, ctx, pool)
|
||||
}
|
||||
|
||||
// Test connection with different IP type
|
||||
|
||||
@@ -170,6 +170,8 @@ func TestCloudSQLPgSimpleToolEndpoints(t *testing.T) {
|
||||
tests.RunPostgresListActiveQueriesTest(t, ctx, pool)
|
||||
tests.RunPostgresListAvailableExtensionsTest(t)
|
||||
tests.RunPostgresListInstalledExtensionsTest(t)
|
||||
tests.RunPostgresDatabaseOverviewTest(t, ctx, pool)
|
||||
tests.RunPostgresListTriggersTest(t, ctx, pool)
|
||||
}
|
||||
|
||||
// Test connection with different IP type
|
||||
|
||||
@@ -198,6 +198,8 @@ func AddPostgresPrebuiltConfig(t *testing.T, config map[string]any) map[string]a
|
||||
PostgresListInstalledExtensionsToolKind = "postgres-list-installed-extensions"
|
||||
PostgresListAvailableExtensionsToolKind = "postgres-list-available-extensions"
|
||||
PostgresListViewsToolKind = "postgres-list-views"
|
||||
PostgresDatabaseOverviewToolKind = "postgres-database-overview"
|
||||
PostgresListTriggersToolKind = "postgres-list-triggers"
|
||||
)
|
||||
|
||||
tools, ok := config["tools"].(map[string]any)
|
||||
@@ -231,10 +233,21 @@ func AddPostgresPrebuiltConfig(t *testing.T, config map[string]any) map[string]a
|
||||
"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",
|
||||
}
|
||||
config["tools"] = tools
|
||||
return config
|
||||
}
|
||||
|
||||
@@ -149,4 +149,6 @@ func TestPostgres(t *testing.T) {
|
||||
tests.RunPostgresListActiveQueriesTest(t, ctx, pool)
|
||||
tests.RunPostgresListAvailableExtensionsTest(t)
|
||||
tests.RunPostgresListInstalledExtensionsTest(t)
|
||||
tests.RunPostgresDatabaseOverviewTest(t, ctx, pool)
|
||||
tests.RunPostgresListTriggersTest(t, ctx, pool)
|
||||
}
|
||||
|
||||
243
tests/tool.go
243
tests/tool.go
@@ -1398,6 +1398,249 @@ func RunPostgresListSchemasTest(t *testing.T, ctx context.Context, pool *pgxpool
|
||||
}
|
||||
}
|
||||
|
||||
func RunPostgresDatabaseOverviewTest(t *testing.T, ctx context.Context, pool *pgxpool.Pool) {
|
||||
const api = "http://127.0.0.1:5000/api/tool/database_overview/invoke"
|
||||
requestBody := bytes.NewBuffer([]byte(`{}`))
|
||||
|
||||
resp, respBody := RunRequest(t, http.MethodPost, api, requestBody, nil)
|
||||
|
||||
if resp.StatusCode != http.StatusOK {
|
||||
t.Fatalf("wrong status code: got %d, want %d, body: %s", resp.StatusCode, http.StatusOK, string(respBody))
|
||||
}
|
||||
|
||||
var bodyWrapper struct {
|
||||
Result json.RawMessage `json:"result"`
|
||||
}
|
||||
if err := json.Unmarshal(respBody, &bodyWrapper); err != nil {
|
||||
t.Fatalf("error decoding response wrapper: %v, body: %s", err, string(respBody))
|
||||
}
|
||||
|
||||
var resultString string
|
||||
if err := json.Unmarshal(bodyWrapper.Result, &resultString); err != nil {
|
||||
resultString = string(bodyWrapper.Result)
|
||||
}
|
||||
|
||||
var got []map[string]any
|
||||
if err := json.Unmarshal([]byte(resultString), &got); err != nil {
|
||||
t.Fatalf("failed to unmarshal nested result string: %v, result string: %s", err, resultString)
|
||||
}
|
||||
|
||||
if len(got) != 1 {
|
||||
t.Fatalf("Expected exactly one row in the result, got %d", len(got))
|
||||
}
|
||||
|
||||
resultRow := got[0]
|
||||
|
||||
// Define expected keys based on the SELECT statement
|
||||
expectedKeys := []string{
|
||||
"pg_version",
|
||||
"is_replica",
|
||||
"uptime",
|
||||
"max_connections",
|
||||
"current_connections",
|
||||
"active_connections",
|
||||
"pct_connections_used",
|
||||
}
|
||||
|
||||
for _, key := range expectedKeys {
|
||||
if _, ok := resultRow[key]; !ok {
|
||||
t.Errorf("Missing expected key in result: %s", key)
|
||||
}
|
||||
}
|
||||
|
||||
// Check types of the fields. JSON numbers are unmarshalled into float64.
|
||||
if _, ok := resultRow["pg_version"].(string); !ok {
|
||||
t.Errorf("Expected 'pg_version' to be a string, got %T", resultRow["pg_version"])
|
||||
}
|
||||
if _, ok := resultRow["is_replica"].(bool); !ok {
|
||||
t.Errorf("Expected 'is_replica' to be a bool, got %T", resultRow["is_replica"])
|
||||
}
|
||||
if _, ok := resultRow["uptime"].(string); !ok {
|
||||
t.Errorf("Expected 'uptime' to be a string, got %T", resultRow["uptime"])
|
||||
}
|
||||
if _, ok := resultRow["max_connections"].(float64); !ok {
|
||||
t.Errorf("Expected 'max_connections' to be a number (float64), got %T", resultRow["max_connections"])
|
||||
}
|
||||
if _, ok := resultRow["current_connections"].(float64); !ok {
|
||||
t.Errorf("Expected 'current_connections' to be a number (float64), got %T", resultRow["current_connections"])
|
||||
}
|
||||
if _, ok := resultRow["active_connections"].(float64); !ok {
|
||||
t.Errorf("Expected 'active_connections' to be a number (float64), got %T", resultRow["active_connections"])
|
||||
}
|
||||
if _, ok := resultRow["pct_connections_used"].(float64); !ok {
|
||||
t.Errorf("Expected 'pct_connections_used' to be a number (float64), got %T", resultRow["pct_connections_used"])
|
||||
}
|
||||
|
||||
// Basic sanity checks on values
|
||||
if maxConn, ok := resultRow["max_connections"].(float64); ok {
|
||||
if maxConn <= 0 {
|
||||
t.Errorf("Expected 'max_connections' to be positive, got %f", maxConn)
|
||||
}
|
||||
}
|
||||
|
||||
if pctUsed, ok := resultRow["pct_connections_used"].(float64); ok {
|
||||
if pctUsed < 0 || pctUsed > 100 {
|
||||
t.Errorf("Expected 'pct_connections_used' to be between 0 and 100, got %f", pctUsed)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func setupPostgresTrigger(t *testing.T, ctx context.Context, pool *pgxpool.Pool, schemaName, tableName, functionName, triggerName string) func() {
|
||||
t.Helper()
|
||||
|
||||
createSchemaStmt := fmt.Sprintf("CREATE SCHEMA %s", schemaName)
|
||||
if _, err := pool.Exec(ctx, createSchemaStmt); err != nil {
|
||||
t.Fatalf("failed to create schema %s: %v", schemaName, err)
|
||||
}
|
||||
|
||||
createTableStmt := fmt.Sprintf("CREATE TABLE %s.%s (id SERIAL PRIMARY KEY, name TEXT)", schemaName, tableName)
|
||||
if _, err := pool.Exec(ctx, createTableStmt); err != nil {
|
||||
t.Fatalf("failed to create table %s.%s: %v", schemaName, tableName, err)
|
||||
}
|
||||
|
||||
createFunctionStmt := fmt.Sprintf(`
|
||||
CREATE OR REPLACE FUNCTION %s.%s() RETURNS TRIGGER AS $$
|
||||
BEGIN
|
||||
RETURN NEW;
|
||||
END;
|
||||
$$ LANGUAGE plpgsql;
|
||||
`, schemaName, functionName)
|
||||
if _, err := pool.Exec(ctx, createFunctionStmt); err != nil {
|
||||
t.Fatalf("failed to create function %s.%s: %v", schemaName, functionName, err)
|
||||
}
|
||||
|
||||
createTriggerStmt := fmt.Sprintf(`
|
||||
CREATE TRIGGER %s
|
||||
AFTER INSERT ON %s.%s
|
||||
FOR EACH ROW
|
||||
EXECUTE FUNCTION %s.%s();
|
||||
`, triggerName, schemaName, tableName, schemaName, functionName)
|
||||
if _, err := pool.Exec(ctx, createTriggerStmt); err != nil {
|
||||
t.Fatalf("failed to create trigger %s: %v", triggerName, err)
|
||||
}
|
||||
|
||||
return func() {
|
||||
dropSchemaStmt := fmt.Sprintf("DROP SCHEMA %s CASCADE", schemaName)
|
||||
if _, err := pool.Exec(ctx, dropSchemaStmt); err != nil {
|
||||
t.Fatalf("failed to drop schema %s: %v", schemaName, err)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func RunPostgresListTriggersTest(t *testing.T, ctx context.Context, pool *pgxpool.Pool) {
|
||||
uniqueID := strings.ReplaceAll(uuid.New().String(), "-", "")
|
||||
schemaName := "test_schema_" + uniqueID
|
||||
tableName := "test_table_" + uniqueID
|
||||
functionName := "test_func_" + uniqueID
|
||||
triggerName := "test_trigger_" + uniqueID
|
||||
|
||||
cleanup := setupPostgresTrigger(t, ctx, pool, schemaName, tableName, functionName, triggerName)
|
||||
defer cleanup()
|
||||
|
||||
// Definition can vary slightly based on server version/settings, so we fetch it to compare.
|
||||
var expectedDef string
|
||||
getDefQuery := fmt.Sprintf("SELECT pg_get_triggerdef(oid) FROM pg_trigger WHERE tgname = '%s'", triggerName)
|
||||
err := pool.QueryRow(ctx, getDefQuery).Scan(&expectedDef)
|
||||
if err != nil {
|
||||
t.Fatalf("failed to fetch trigger definition: %v", err)
|
||||
}
|
||||
|
||||
wantTrigger := map[string]any{
|
||||
"trigger_name": triggerName,
|
||||
"schema_name": schemaName,
|
||||
"table_name": tableName,
|
||||
"status": "ENABLED",
|
||||
"timing": "AFTER",
|
||||
"events": "INSERT",
|
||||
"activation_level": "ROW",
|
||||
"function_name": functionName,
|
||||
"definition": expectedDef,
|
||||
}
|
||||
|
||||
invokeTcs := []struct {
|
||||
name string
|
||||
requestBody io.Reader
|
||||
wantStatusCode int
|
||||
want []map[string]any
|
||||
}{
|
||||
{
|
||||
name: "list all triggers (expecting the one we created)",
|
||||
requestBody: bytes.NewBuffer([]byte(`{}`)),
|
||||
wantStatusCode: http.StatusOK,
|
||||
want: []map[string]any{wantTrigger},
|
||||
},
|
||||
{
|
||||
name: "filter by trigger_name",
|
||||
requestBody: bytes.NewBuffer([]byte(fmt.Sprintf(`{"trigger_name": "%s"}`, triggerName))),
|
||||
wantStatusCode: http.StatusOK,
|
||||
want: []map[string]any{wantTrigger},
|
||||
},
|
||||
{
|
||||
name: "filter by schema_name",
|
||||
requestBody: bytes.NewBuffer([]byte(fmt.Sprintf(`{"schema_name": "%s"}`, schemaName))),
|
||||
wantStatusCode: http.StatusOK,
|
||||
want: []map[string]any{wantTrigger},
|
||||
},
|
||||
{
|
||||
name: "filter by table_name",
|
||||
requestBody: bytes.NewBuffer([]byte(fmt.Sprintf(`{"table_name": "%s"}`, tableName))),
|
||||
wantStatusCode: http.StatusOK,
|
||||
want: []map[string]any{wantTrigger},
|
||||
},
|
||||
{
|
||||
name: "filter by non-existent trigger_name",
|
||||
requestBody: bytes.NewBuffer([]byte(`{"trigger_name": "non_existent_trigger"}`)),
|
||||
wantStatusCode: http.StatusOK,
|
||||
want: nil,
|
||||
},
|
||||
{
|
||||
name: "filter by non-existent schema_name",
|
||||
requestBody: bytes.NewBuffer([]byte(`{"schema_name": "non_existent_schema"}`)),
|
||||
wantStatusCode: http.StatusOK,
|
||||
want: nil,
|
||||
},
|
||||
{
|
||||
name: "filter by non-existent table_name",
|
||||
requestBody: bytes.NewBuffer([]byte(`{"table_name": "non_existent_table"}`)),
|
||||
wantStatusCode: http.StatusOK,
|
||||
want: nil,
|
||||
},
|
||||
}
|
||||
for _, tc := range invokeTcs {
|
||||
t.Run(tc.name, func(t *testing.T) {
|
||||
const api = "http://127.0.0.1:5000/api/tool/list_triggers/invoke"
|
||||
resp, respBody := RunRequest(t, http.MethodPost, api, tc.requestBody, nil)
|
||||
if resp.StatusCode != tc.wantStatusCode {
|
||||
t.Fatalf("wrong status code: got %d, want %d, body: %s", resp.StatusCode, tc.wantStatusCode, string(respBody))
|
||||
}
|
||||
if tc.wantStatusCode != http.StatusOK {
|
||||
return
|
||||
}
|
||||
|
||||
var bodyWrapper struct {
|
||||
Result json.RawMessage `json:"result"`
|
||||
}
|
||||
if err := json.Unmarshal(respBody, &bodyWrapper); err != nil {
|
||||
t.Fatalf("error decoding response wrapper: %v", err)
|
||||
}
|
||||
|
||||
var resultString string
|
||||
if err := json.Unmarshal(bodyWrapper.Result, &resultString); err != nil {
|
||||
resultString = string(bodyWrapper.Result)
|
||||
}
|
||||
|
||||
var got []map[string]any
|
||||
if err := json.Unmarshal([]byte(resultString), &got); err != nil {
|
||||
t.Fatalf("failed to unmarshal nested result string: %v, content: %s", err, resultString)
|
||||
}
|
||||
|
||||
if diff := cmp.Diff(tc.want, got); diff != "" {
|
||||
t.Errorf("Unexpected result (-want +got):\n%s", diff)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func RunPostgresListActiveQueriesTest(t *testing.T, ctx context.Context, pool *pgxpool.Pool) {
|
||||
type queryListDetails struct {
|
||||
ProcessId any `json:"pid"`
|
||||
|
||||
Reference in New Issue
Block a user