mirror of
https://github.com/googleapis/genai-toolbox.git
synced 2026-01-10 07:58:12 -05:00
feat(tools/postgres-list-publication-tables): add new postgres-list-publication-tables tool (#1919)
## Description Adds a postgresql custom list_publication_tables tool, that returns the details of publication tables present in database. Test Output: <img width="845" height="239" alt="Screenshot 2025-11-11 at 12 50 59 AM" src="https://github.com/user-attachments/assets/b7606e44-c5f6-4fc7-865e-7efadd112eff" /> <img width="1529" height="648" alt="Screenshot 2025-11-11 at 1 15 18 AM" src="https://github.com/user-attachments/assets/6192b772-f0bc-4fb4-8032-ca487434d77c" /> > 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 Co-authored-by: Averi Kitsch <akitsch@google.com> Co-authored-by: Wenxin Du <117315983+duwenxin99@users.noreply.github.com>
This commit is contained in:
@@ -187,6 +187,7 @@ import (
|
||||
_ "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/postgreslistpublicationtables"
|
||||
_ "github.com/googleapis/genai-toolbox/internal/tools/postgres/postgreslistquerystats"
|
||||
_ "github.com/googleapis/genai-toolbox/internal/tools/postgres/postgreslistschemas"
|
||||
_ "github.com/googleapis/genai-toolbox/internal/tools/postgres/postgreslistsequences"
|
||||
|
||||
@@ -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"},
|
||||
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"},
|
||||
},
|
||||
},
|
||||
},
|
||||
@@ -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"},
|
||||
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"},
|
||||
},
|
||||
},
|
||||
},
|
||||
@@ -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"},
|
||||
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"},
|
||||
},
|
||||
},
|
||||
},
|
||||
|
||||
@@ -50,6 +50,7 @@ details on how to connect your AI tools (IDEs) to databases via Toolbox and MCP.
|
||||
* `list_triggers`: Lists triggers in the database.
|
||||
* `list_indexes`: List available user indexes in a PostgreSQL database.
|
||||
* `list_sequences`: List sequences in a PostgreSQL database.
|
||||
* `list_publication_tables`: List publication tables in a PostgreSQL database.
|
||||
|
||||
## AlloyDB Postgres Admin
|
||||
|
||||
@@ -227,6 +228,7 @@ details on how to connect your AI tools (IDEs) to databases via Toolbox and MCP.
|
||||
* `list_triggers`: Lists triggers in the database.
|
||||
* `list_indexes`: List available user indexes in a PostgreSQL database.
|
||||
* `list_sequences`: List sequences in a PostgreSQL database.
|
||||
* `list_publication_tables`: List publication tables in a PostgreSQL database.
|
||||
|
||||
## Cloud SQL for PostgreSQL Observability
|
||||
|
||||
@@ -532,6 +534,7 @@ details on how to connect your AI tools (IDEs) to databases via Toolbox and MCP.
|
||||
* `list_triggers`: Lists triggers in the database.
|
||||
* `list_indexes`: List available user indexes in a PostgreSQL database.
|
||||
* `list_sequences`: List sequences in a PostgreSQL database.
|
||||
* `list_publication_tables`: List publication tables in a PostgreSQL database.
|
||||
|
||||
## Google Cloud Serverless for Apache Spark
|
||||
|
||||
|
||||
@@ -77,6 +77,9 @@ cluster][alloydb-free-trial].
|
||||
- [`postgres-get-column-cardinality`](../tools/postgres/postgres-get-column-cardinality.md)
|
||||
List cardinality of columns in a table in a PostgreSQL database.
|
||||
|
||||
- [`postgres-list-publication-tables`](../tools/postgres/postgres-list-publication-tables.md)
|
||||
List publication tables in a PostgreSQL database.
|
||||
|
||||
### Pre-built Configurations
|
||||
|
||||
- [AlloyDB using MCP](https://googleapis.github.io/genai-toolbox/how-to/connect-ide/alloydb_pg_mcp/)
|
||||
|
||||
@@ -73,6 +73,9 @@ to a database by following these instructions][csql-pg-quickstart].
|
||||
- [`postgres-get-column-cardinality`](../tools/postgres/postgres-get-column-cardinality.md)
|
||||
List cardinality of columns in a table in a PostgreSQL database.
|
||||
|
||||
- [`postgres-list-publication-tables`](../tools/postgres/postgres-list-publication-tables.md)
|
||||
List publication tables in a PostgreSQL database.
|
||||
|
||||
### Pre-built Configurations
|
||||
|
||||
- [Cloud SQL for Postgres using
|
||||
|
||||
@@ -68,6 +68,9 @@ reputation for reliability, feature robustness, and performance.
|
||||
- [`postgres-get-column-cardinality`](../tools/postgres/postgres-get-column-cardinality.md)
|
||||
List cardinality of columns in a table in a PostgreSQL database.
|
||||
|
||||
- [`postgres-list-publication-tables`](../tools/postgres/postgres-list-publication-tables.md)
|
||||
List publication tables 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,66 @@
|
||||
---
|
||||
title: "postgres-list-publication-tables"
|
||||
type: docs
|
||||
weight: 1
|
||||
description: >
|
||||
The "postgres-list-publication-tables" tool lists publication tables in a Postgres database.
|
||||
aliases:
|
||||
- /resources/tools/postgres-list-publication-tables
|
||||
---
|
||||
|
||||
## About
|
||||
|
||||
The `postgres-list-publication-tables` tool lists all publication tables 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-publication-tables` lists detailed information as JSON for publication tables. A publication table in PostgreSQL is a
|
||||
table that is explicitly included as a source for replication within a publication (a set of changes generated from a table or group
|
||||
of tables) as part of the logical replication feature. The tool takes the following input parameters:
|
||||
|
||||
- `table_names` (optional): Filters by a comma-separated list of table names. Default: `""`
|
||||
- `publication_names` (optional): Filters by a comma-separated list of publication names. Default: `""`
|
||||
- `schema_names` (optional): Filters by a comma-separated list of schema names. Default: `""`
|
||||
- `limit` (optional): The maximum number of rows to return. Default: `50`
|
||||
|
||||
## Example
|
||||
|
||||
```yaml
|
||||
tools:
|
||||
list_indexes:
|
||||
kind: postgres-list-publication-tables
|
||||
source: postgres-source
|
||||
description: |
|
||||
Lists all tables that are explicitly part of a publication in the database.
|
||||
Tables that are part of a publication via 'FOR ALL TABLES' are not included,
|
||||
unless they are also explicitly added to the publication.
|
||||
Returns the publication name, schema name, and table name, along with
|
||||
definition details indicating if it publishes all tables, whether it
|
||||
replicates inserts, updates, deletes, or truncates, and the publication
|
||||
owner.
|
||||
```
|
||||
|
||||
The response is a JSON array with the following elements:
|
||||
```json
|
||||
{
|
||||
"publication_name": "Name of the publication",
|
||||
"schema_name": "Name of the schema the table belongs to",
|
||||
"table_name": "Name of the table",
|
||||
"publishes_all_tables": "boolean indicating if the publication was created with FOR ALL TABLES",
|
||||
"publishes_inserts": "boolean indicating if INSERT operations are replicated",
|
||||
"publishes_updates": "boolean indicating if UPDATE operations are replicated",
|
||||
"publishes_deletes": "boolean indicating if DELETE operations are replicated",
|
||||
"publishes_truncates": "boolean indicating if TRUNCATE operations are replicated",
|
||||
"publication_owner": "Username of the database role that owns the publication"
|
||||
}
|
||||
```
|
||||
|
||||
## Reference
|
||||
|
||||
| **field** | **type** | **required** | **description** |
|
||||
|-------------|:--------:|:------------:|------------------------------------------------------|
|
||||
| kind | string | true | Must be "postgres-list-publication-tables". |
|
||||
| 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. |
|
||||
@@ -200,6 +200,10 @@ tools:
|
||||
kind: postgres-get-column-cardinality
|
||||
source: alloydb-pg-source
|
||||
|
||||
list_publication_tables:
|
||||
kind: postgres-list-publication-tables
|
||||
source: alloydb-pg-source
|
||||
|
||||
toolsets:
|
||||
alloydb_postgres_database_tools:
|
||||
- execute_sql
|
||||
@@ -224,3 +228,4 @@ toolsets:
|
||||
- replication_stats
|
||||
- list_query_stats
|
||||
- get_column_cardinality
|
||||
- list_publication_tables
|
||||
|
||||
@@ -202,6 +202,10 @@ tools:
|
||||
kind: postgres-get-column-cardinality
|
||||
source: cloudsql-pg-source
|
||||
|
||||
list_publication_tables:
|
||||
kind: postgres-list-publication-tables
|
||||
source: cloudsql-pg-source
|
||||
|
||||
toolsets:
|
||||
cloud_sql_postgres_database_tools:
|
||||
- execute_sql
|
||||
@@ -226,3 +230,4 @@ toolsets:
|
||||
- replication_stats
|
||||
- list_query_stats
|
||||
- get_column_cardinality
|
||||
- list_publication_tables
|
||||
|
||||
@@ -201,6 +201,10 @@ tools:
|
||||
kind: postgres-get-column-cardinality
|
||||
source: postgresql-source
|
||||
|
||||
list_publication_tables:
|
||||
kind: postgres-list-publication-tables
|
||||
source: postgresql-source
|
||||
|
||||
toolsets:
|
||||
postgres_database_tools:
|
||||
- execute_sql
|
||||
@@ -225,4 +229,4 @@ toolsets:
|
||||
- replication_stats
|
||||
- list_query_stats
|
||||
- get_column_cardinality
|
||||
|
||||
- list_publication_tables
|
||||
|
||||
@@ -0,0 +1,216 @@
|
||||
// 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 postgreslistpublicationtables
|
||||
|
||||
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-publication-tables"
|
||||
|
||||
const listPublicationTablesStatement = `
|
||||
WITH
|
||||
publication_details AS (
|
||||
SELECT
|
||||
pt.pubname AS publication_name,
|
||||
pt.schemaname AS schema_name,
|
||||
pt.tablename AS table_name,
|
||||
-- Definition details
|
||||
p.puballtables AS publishes_all_tables,
|
||||
p.pubinsert AS publishes_inserts,
|
||||
p.pubupdate AS publishes_updates,
|
||||
p.pubdelete AS publishes_deletes,
|
||||
p.pubtruncate AS publishes_truncates,
|
||||
-- Owner information
|
||||
pg_catalog.pg_get_userbyid(p.pubowner) AS publication_owner
|
||||
FROM pg_catalog.pg_publication_tables pt
|
||||
JOIN pg_catalog.pg_publication p
|
||||
ON pt.pubname = p.pubname
|
||||
)
|
||||
SELECT *
|
||||
FROM publication_details
|
||||
WHERE
|
||||
(NULLIF(TRIM($1::text), '') IS NULL OR table_name = ANY(regexp_split_to_array(TRIM($1::text), '\s*,\s*')))
|
||||
AND (NULLIF(TRIM($2::text), '') IS NULL OR publication_name = ANY(regexp_split_to_array(TRIM($2::text), '\s*,\s*')))
|
||||
AND (NULLIF(TRIM($3::text), '') IS NULL OR schema_name = ANY(regexp_split_to_array(TRIM($3::text), '\s*,\s*')))
|
||||
ORDER BY
|
||||
publication_name, schema_name, table_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 := parameters.Parameters{
|
||||
parameters.NewStringParameterWithDefault("table_names", "", "Optional: Filters by a comma-separated list of table names."),
|
||||
parameters.NewStringParameterWithDefault("publication_names", "", "Optional: Filters by a comma-separated list of publication names."),
|
||||
parameters.NewStringParameterWithDefault("schema_names", "", "Optional: Filters by a comma-separated list of schema names."),
|
||||
parameters.NewIntParameterWithDefault("limit", 50, "Optional: The maximum number of rows to return."),
|
||||
}
|
||||
description := cfg.Description
|
||||
if description == "" {
|
||||
description = "Lists all publication tables in the database. Returns the publication name, schema name, and table name, along with definition details indicating if it publishes all tables, whether it replicates inserts, updates, deletes, or truncates, and the publication owner."
|
||||
}
|
||||
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, listPublicationTablesStatement, 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)
|
||||
}
|
||||
|
||||
if err := results.Err(); err != nil {
|
||||
return err.Error(), 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 postgreslistpublicationtables_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/postgreslistpublicationtables"
|
||||
)
|
||||
|
||||
func TestParseFromYamlPostgresListPublicationTables(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-publication-tables
|
||||
source: my-postgres-instance
|
||||
description: some description
|
||||
authRequired:
|
||||
- my-google-auth-service
|
||||
- other-auth-service
|
||||
`,
|
||||
want: server.ToolConfigs{
|
||||
"example_tool": postgreslistpublicationtables.Config{
|
||||
Name: "example_tool",
|
||||
Kind: "postgres-list-publication-tables",
|
||||
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-publication-tables
|
||||
source: my-postgres-instance
|
||||
description: some description
|
||||
`,
|
||||
want: server.ToolConfigs{
|
||||
"example_tool": postgreslistpublicationtables.Config{
|
||||
Name: "example_tool",
|
||||
Kind: "postgres-list-publication-tables",
|
||||
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)
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
}
|
||||
@@ -195,6 +195,7 @@ func TestAlloyDBPgToolEndpoints(t *testing.T) {
|
||||
tests.RunPostgresLongRunningTransactionsTest(t, ctx, pool)
|
||||
tests.RunPostgresListQueryStatsTest(t, ctx, pool)
|
||||
tests.RunPostgresGetColumnCardinalityTest(t, ctx, pool)
|
||||
tests.RunPostgresListPublicationTablesTest(t, ctx, pool)
|
||||
}
|
||||
|
||||
// Test connection with different IP type
|
||||
|
||||
@@ -179,6 +179,7 @@ func TestCloudSQLPgSimpleToolEndpoints(t *testing.T) {
|
||||
tests.RunPostgresLongRunningTransactionsTest(t, ctx, pool)
|
||||
tests.RunPostgresListQueryStatsTest(t, ctx, pool)
|
||||
tests.RunPostgresGetColumnCardinalityTest(t, ctx, pool)
|
||||
tests.RunPostgresListPublicationTablesTest(t, ctx, pool)
|
||||
}
|
||||
|
||||
// Test connection with different IP type
|
||||
|
||||
@@ -207,6 +207,7 @@ func AddPostgresPrebuiltConfig(t *testing.T, config map[string]any) map[string]a
|
||||
PostgresReplicationStatsToolKind = "postgres-replication-stats"
|
||||
PostgresListQueryStatsToolKind = "postgres-list-query-stats"
|
||||
PostgresGetColumnCardinalityToolKind = "postgres-get-column-cardinality"
|
||||
PostgresListPublicationTablesToolKind = "postgres-list-publication-tables"
|
||||
)
|
||||
|
||||
tools, ok := config["tools"].(map[string]any)
|
||||
@@ -265,14 +266,21 @@ func AddPostgresPrebuiltConfig(t *testing.T, config map[string]any) map[string]a
|
||||
"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",
|
||||
|
||||
@@ -158,4 +158,5 @@ func TestPostgres(t *testing.T) {
|
||||
tests.RunPostgresReplicationStatsTest(t, ctx, pool)
|
||||
tests.RunPostgresListQueryStatsTest(t, ctx, pool)
|
||||
tests.RunPostgresGetColumnCardinalityTest(t, ctx, pool)
|
||||
tests.RunPostgresListPublicationTablesTest(t, ctx, pool)
|
||||
}
|
||||
|
||||
148
tests/tool.go
148
tests/tool.go
@@ -1660,6 +1660,154 @@ func RunPostgresListTriggersTest(t *testing.T, ctx context.Context, pool *pgxpoo
|
||||
}
|
||||
}
|
||||
|
||||
func setupPostgresPublicationTable(t *testing.T, ctx context.Context, pool *pgxpool.Pool, tableName string, pubName string) func(t *testing.T) {
|
||||
t.Helper()
|
||||
createTableStmt := fmt.Sprintf("CREATE TABLE %s (id SERIAL PRIMARY KEY, name TEXT);", tableName)
|
||||
if _, err := pool.Exec(ctx, createTableStmt); err != nil {
|
||||
t.Fatalf("unable to create table %s: %v", tableName, err)
|
||||
}
|
||||
|
||||
createPubStmt := fmt.Sprintf("CREATE PUBLICATION %s FOR TABLE %s;", pubName, tableName)
|
||||
if _, err := pool.Exec(ctx, createPubStmt); err != nil {
|
||||
if _, dropErr := pool.Exec(ctx, fmt.Sprintf("DROP TABLE IF EXISTS %s;", tableName)); dropErr != nil {
|
||||
t.Errorf("unable to drop table after failing to create publication: %v", dropErr)
|
||||
}
|
||||
t.Fatalf("unable to create publication %s: %v", pubName, err)
|
||||
}
|
||||
|
||||
return func(t *testing.T) {
|
||||
t.Helper()
|
||||
if _, err := pool.Exec(ctx, fmt.Sprintf("DROP PUBLICATION IF EXISTS %s;", pubName)); err != nil {
|
||||
t.Errorf("unable to drop publication %s: %v", pubName, err)
|
||||
}
|
||||
if _, err := pool.Exec(ctx, fmt.Sprintf("DROP TABLE IF EXISTS %s;", tableName)); err != nil {
|
||||
t.Errorf("unable to drop table %s: %v", tableName, err)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func RunPostgresListPublicationTablesTest(t *testing.T, ctx context.Context, pool *pgxpool.Pool) {
|
||||
table1Name := "pub_table_1"
|
||||
pub1Name := "pub_1"
|
||||
|
||||
table2Name := "pub_table_2"
|
||||
pub2Name := "pub_2"
|
||||
|
||||
cleanup := setupPostgresPublicationTable(t, ctx, pool, table1Name, pub1Name)
|
||||
defer cleanup(t)
|
||||
cleanup2 := setupPostgresPublicationTable(t, ctx, pool, table2Name, pub2Name)
|
||||
defer cleanup2(t)
|
||||
|
||||
// Fetch the current user to match the publication_owner
|
||||
var currentUser string
|
||||
err := pool.QueryRow(ctx, "SELECT current_user;").Scan(¤tUser)
|
||||
if err != nil {
|
||||
t.Fatalf("unable to fetch current user: %v", err)
|
||||
}
|
||||
|
||||
wantTable1 := map[string]any{
|
||||
"publication_name": pub1Name,
|
||||
"schema_name": "public",
|
||||
"table_name": table1Name,
|
||||
"publishes_all_tables": false,
|
||||
"publishes_inserts": true,
|
||||
"publishes_updates": true,
|
||||
"publishes_deletes": true,
|
||||
"publishes_truncates": true,
|
||||
"publication_owner": currentUser,
|
||||
}
|
||||
|
||||
wantTable2 := map[string]any{
|
||||
"publication_name": pub2Name,
|
||||
"schema_name": "public",
|
||||
"table_name": table2Name,
|
||||
"publishes_all_tables": false,
|
||||
"publishes_inserts": true,
|
||||
"publishes_updates": true,
|
||||
"publishes_deletes": true,
|
||||
"publishes_truncates": true,
|
||||
"publication_owner": currentUser,
|
||||
}
|
||||
|
||||
invokeTcs := []struct {
|
||||
name string
|
||||
requestBody io.Reader
|
||||
wantStatusCode int
|
||||
want []map[string]any
|
||||
}{
|
||||
{
|
||||
name: "list all publication tables",
|
||||
requestBody: bytes.NewBufferString(`{}`),
|
||||
wantStatusCode: http.StatusOK,
|
||||
want: []map[string]any{wantTable1, wantTable2},
|
||||
},
|
||||
{
|
||||
name: "list all tables for the created publication",
|
||||
requestBody: bytes.NewBufferString(fmt.Sprintf(`{"publication_names": "%s"}`, pub1Name)),
|
||||
wantStatusCode: http.StatusOK,
|
||||
want: []map[string]any{wantTable1},
|
||||
},
|
||||
{
|
||||
name: "filter by table_name",
|
||||
requestBody: bytes.NewBufferString(fmt.Sprintf(`{"table_names": "%s, %s"}`, table1Name, table2Name)),
|
||||
wantStatusCode: http.StatusOK,
|
||||
want: []map[string]any{wantTable1, wantTable2},
|
||||
},
|
||||
{
|
||||
name: "filter by schema_name and table_name",
|
||||
requestBody: bytes.NewBufferString(fmt.Sprintf(`{"schema_names": "public", "table_name": "%s , %s"}`, table1Name, table2Name)),
|
||||
wantStatusCode: http.StatusOK,
|
||||
want: []map[string]any{wantTable1, wantTable2},
|
||||
},
|
||||
{
|
||||
name: "invoke list_publication_tables with non-existent table",
|
||||
requestBody: bytes.NewBufferString(`{"table_names": "non_existent_table"}`),
|
||||
wantStatusCode: http.StatusOK,
|
||||
want: nil,
|
||||
},
|
||||
{
|
||||
name: "invoke list_publication_tables with non-existent publication",
|
||||
requestBody: bytes.NewBufferString(`{"publication_names": "non_existent_pub"}`),
|
||||
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_publication_tables/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