feat(source/alloydb)!: restructure prebuilt toolsets

.ci/integration.cloudbuild.yaml

@@ -134,7 +134,7 @@ steps:
       - "ALLOYDB_POSTGRES_DATABASE=$_DATABASE_NAME"
       - "ALLOYDB_POSTGRES_REGION=$_REGION"
       - "SERVICE_ACCOUNT_EMAIL=$SERVICE_ACCOUNT_EMAIL"
-    secretEnv: ["ALLOYDB_POSTGRES_USER", "ALLOYDB_POSTGRES_PASS", "CLIENT_ID", "API_KEY"]
+    secretEnv: ["ALLOYDB_POSTGRES_USER", "ALLOYDB_POSTGRES_PASSWORD", "CLIENT_ID", "API_KEY"]
     volumes:
       - name: "go"
         path: "/gopath"
@@ -1058,7 +1058,7 @@ availableSecrets:
     - versionName: projects/$PROJECT_ID/secrets/alloydb_pg_user/versions/latest
       env: ALLOYDB_POSTGRES_USER
     - versionName: projects/$PROJECT_ID/secrets/alloydb_pg_pass/versions/latest
-      env: ALLOYDB_POSTGRES_PASS
+      env: ALLOYDB_POSTGRES_PASSWORD
     - versionName: projects/$PROJECT_ID/secrets/alloydb_ai_nl_user/versions/latest
       env: ALLOYDB_AI_NL_USER
     - versionName: projects/$PROJECT_ID/secrets/alloydb_ai_nl_pass/versions/latest

cmd/internal/tools_file_test.go

@@ -1621,9 +1621,33 @@ func TestPrebuiltTools(t *testing.T) {
 			name: "alloydb omni prebuilt tools",
 			in:   alloydb_omni_config,
 			wantToolset: server.ToolsetConfigs{
-				"alloydb_omni_database_tools": tools.ToolsetConfig{
-					Name:      "alloydb_omni_database_tools",
-					ToolNames: []string{"execute_sql", "list_tables", "list_active_queries", "list_available_extensions", "list_installed_extensions", "list_autovacuum_configurations", "list_columnar_configurations", "list_columnar_recommended_columns", "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", "list_roles", "list_table_stats", "list_stored_procedure"},
+				"data": tools.ToolsetConfig{
+					Name:      "data",
+					ToolNames: []string{"execute_sql", "list_tables", "list_views", "list_schemas,", "list_triggers", "list_indexes", "list_sequences,", "list_stored_procedure"},
+				},
+				"performance": tools.ToolsetConfig{
+					Name:      "performance",
+					ToolNames: []string{"execute_sql", "get_query_plan", "list_query_stats", "get_column_cardinality", "list_table_stats", "list_database_stats", "list_active_queries"},
+				},
+				"monitor": tools.ToolsetConfig{
+					Name:      "monitor",
+					ToolNames: []string{"database_overview", "list_active_queries", "long_running_transactions", "list_locks", "list_database_stats", "list_pg_settings"},
+				},
+				"optimize": tools.ToolsetConfig{
+					Name:      "optimize",
+					ToolNames: []string{"list_pg_settings", "list_memory_configurations", "list_available_extensions", "list_installed_extensions", "list_autovacuum_configurations", "list_columnar_configurations", "list_columnar_recommended_columns"},
+				},
+				"health": tools.ToolsetConfig{
+					Name:      "health",
+					ToolNames: []string{"list_top_bloated_tables", "list_invalid_indexes", "list_table_stats", "list_tablespaces", "database_overview", "list_autovacuum_configurations"},
+				},
+				"replication": tools.ToolsetConfig{
+					Name:      "replication",
+					ToolNames: []string{"replication_stats", "list_replication_slots", "list_publication_tables", "database_overview"},
+				},
+				"access-control": tools.ToolsetConfig{
+					Name:      "access-control",
+					ToolNames: []string{"list_roles", "list_pg_settings", "database_overview"},
 				},
 			},
 		},
@@ -1671,9 +1695,33 @@ func TestPrebuiltTools(t *testing.T) {
 			name: "alloydb prebuilt tools",
 			in:   alloydb_config,
 			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", "list_pg_settings", "list_database_stats", "list_roles", "list_table_stats", "list_stored_procedure"},
+				"admin": tools.ToolsetConfig{
+					Name:      "admin",
+					ToolNames: []string{"create_cluster", "get_cluster", "list_clusters", "create_instance", "get_instance", "list_instances", "database_overview", "wait_for_operation"},
+				},
+				"access-management": tools.ToolsetConfig{
+					Name:      "access-management",
+					ToolNames: []string{"create_user", "list_users", "get_user", "list_roles", "list_pg_settings", "database_overview"},
+				},
+				"data": tools.ToolsetConfig{
+					Name:      "data",
+					ToolNames: []string{"execute_sql,", "list_tables", "list_views", "list_schemas", "list_triggers,", "list_indexes", "list_sequences", "list_stored_procedure"},
+				},
+				"monitor": tools.ToolsetConfig{
+					Name:      "monitor",
+					ToolNames: []string{"list_active_queries", "list_query_stats", "get_query_plan", "get_query_metrics", "get_system_metrics", "long_running_transactions", "list_locks", "list_database_stats"},
+				},
+				"health": tools.ToolsetConfig{
+					Name:      "health",
+					ToolNames: []string{"list_top_bloated_tables", "list_invalid_indexes", "list_table_stats", "get_column_cardinality", "list_autovacuum_configurations", "list_tablespaces", "database_overview", "get_instance"},
+				},
+				"optimize": tools.ToolsetConfig{
+					Name:      "optimize",
+					ToolNames: []string{"list_available_extensions", "list_installed_extensions", "list_memory_configurations", "list_pg_settings", "database_overview", "get_cluster"},
+				},
+				"replication": tools.ToolsetConfig{
+					Name:      "replication",
+					ToolNames: []string{"replication_stats", "list_replication_slots", "list_publication_tables", "list_instances", "get_instance", "database_overview"},
 				},
 			},
 		},
@@ -1960,6 +2008,15 @@ func TestPrebuiltTools(t *testing.T) {
 			if len(toolsFile.Prompts) != 0 {
 				t.Fatalf("expected empty prompts map for prebuilt config, got: %v", toolsFile.Prompts)
 			}
+
+			t.Run("check toolset sizes", func(t *testing.T) {
+				for tsName, ts := range toolsFile.Toolsets {
+					if len(ts.ToolNames) > 10 {
+						t.Logf("WARNING: Toolset %q in config %q has %d tools, which is larger than the recommended maximum of 10.", tsName, tc.name, len(ts.ToolNames))
+						fmt.Printf("WARNING: Toolset %q in config %q has %d tools, which is larger than the recommended maximum of 10.\n", tsName, tc.name, len(ts.ToolNames))
+					}
+				}
+			})
 		})
 	}
 }

internal/prebuiltconfigs/tools/alloydb-omni.yaml

@@ -21,45 +21,36 @@ sources:
         user: ${ALLOYDB_OMNI_USER}
         password: ${ALLOYDB_OMNI_PASSWORD:}
         queryParams: ${ALLOYDB_OMNI_QUERY_PARAMS:}
-
 tools:
     execute_sql:
         kind: postgres-execute-sql
         source: alloydb-omni-source
         description: Use this tool to execute sql.
-
     list_tables:
         kind: postgres-list-tables
         source: alloydb-omni-source
         description: "Lists detailed schema information (object type, columns, constraints, indexes, triggers, owner, comment) as JSON for user-created tables (ordinary or partitioned). Filters by a comma-separated list of names. If names are omitted, lists all tables in user schemas."
-
     list_active_queries:
         kind: postgres-list-active-queries
         source: alloydb-omni-source
         description: "List the top N (default 50) currently running queries (state='active') from pg_stat_activity, ordered by longest-running first. Returns pid, user, database, application_name, client_addr, state, wait_event_type/wait_event, backend/xact/query start times, computed query_duration, and the SQL text."
-
     list_available_extensions:
         kind: postgres-list-available-extensions
         source: alloydb-omni-source
         description: "Discover all PostgreSQL extensions available for installation on this server, returning name, default_version, and description."
-
     list_installed_extensions:
         kind: postgres-list-installed-extensions
         source: alloydb-omni-source
         description: "List all installed PostgreSQL extensions with their name, version, schema, owner, and description."
-
     long_running_transactions:
         kind: postgres-long-running-transactions
         source: alloydb-omni-source
-
     list_locks:
         kind: postgres-list-locks
         source: alloydb-omni-source
-
     replication_stats:
         kind: postgres-replication-stats
         source: alloydb-omni-source
-
     list_autovacuum_configurations:
         kind: postgres-sql
         source: alloydb-omni-source
@@ -69,7 +60,6 @@ tools:
                 setting
             FROM   pg_settings
             WHERE  category = 'Autovacuum';
-
     list_columnar_configurations:
         kind: postgres-sql
         source: alloydb-omni-source
@@ -79,13 +69,11 @@ tools:
                 setting
             FROM   pg_settings
             WHERE  name like 'google_columnar_engine.%';
-
     list_columnar_recommended_columns:
         kind: postgres-sql
         source: alloydb-omni-source
         description: "Lists columns that AlloyDB Omni recommends adding to the columnar engine to improve query performance."
         statement: select * from g_columnar_recommended_columns;
-
     list_memory_configurations:
         kind: postgres-sql
         source: alloydb-omni-source
@@ -107,7 +95,6 @@ tools:
                 WHERE name IN ('shared_buffers', 'wal_buffers', 'effective_cache_size', 'temp_buffers')
             )
             ORDER BY 1 DESC;
-
     list_top_bloated_tables:
         kind: postgres-sql
         source: alloydb-omni-source
@@ -132,7 +119,6 @@ tools:
               description: "The maximum number of results to return."
               type: integer
               default: 50
-
     list_replication_slots:
         kind: postgres-sql
         source: alloydb-omni-source
@@ -151,7 +137,6 @@ tools:
             catalog_xmin,
             pg_size_pretty(pg_wal_lsn_diff(pg_current_wal_lsn(), restart_lsn)) AS retained_wal
             FROM pg_replication_slots;
-
     list_invalid_indexes:
         kind: postgres-sql
         source: alloydb-omni-source
@@ -169,7 +154,6 @@ tools:
             JOIN pg_class ON pg_class.oid = pg_index.indexrelid
             JOIN pg_namespace ON pg_namespace.oid = pg_class.relnamespace
             WHERE indisvalid = FALSE;
-
     get_query_plan:
         kind: postgres-sql
         source: alloydb-omni-source
@@ -181,97 +165,98 @@ tools:
               type: string
               description: "The SQL statement for which you want to generate plan (omit the EXPLAIN keyword)."
               required: true
-
     list_views:
         kind: postgres-list-views
         source: alloydb-omni-source
-
     list_schemas:
         kind: postgres-list-schemas
         source: alloydb-omni-source
-
     list_indexes:
         kind: postgres-list-indexes
         source: alloydb-omni-source
-
     list_sequences:
         kind: postgres-list-sequences
         source: alloydb-omni-source
-
     database_overview:
         kind: postgres-database-overview
         source: alloydb-omni-source
-
     list_triggers:
         kind: postgres-list-triggers
         source: alloydb-omni-source
-
     list_query_stats:
         kind: postgres-list-query-stats
         source: alloydb-omni-source
-
     get_column_cardinality:
         kind: postgres-get-column-cardinality
         source: alloydb-omni-source
-
     list_table_stats:
         kind: postgres-list-table-stats
         source: alloydb-omni-source
-
     list_publication_tables:
         kind: postgres-list-publication-tables
         source: alloydb-omni-source
-
     list_tablespaces:
         kind: postgres-list-tablespaces
         source: alloydb-omni-source
-
     list_pg_settings:
         kind: postgres-list-pg-settings
         source: alloydb-omni-source
-
     list_database_stats:
         kind: postgres-list-database-stats
         source: alloydb-omni-source
-
     list_roles:
         kind: postgres-list-roles
         source: alloydb-omni-source
-
     list_stored_procedure:
         kind: postgres-list-stored-procedure
         source: alloydb-omni-source
-
 toolsets:
-    alloydb_omni_database_tools:
+    data:
         - execute_sql
         - list_tables
+        - list_views
+        - list_schemas,
+        - list_triggers
+        - list_indexes
+        - list_sequences,
+        - list_stored_procedure
+    performance:
+        - execute_sql
+        - get_query_plan
+        - list_query_stats
+        - get_column_cardinality
+        - list_table_stats
+        - list_database_stats
         - list_active_queries
+    monitor:
+        - database_overview
+        - list_active_queries
+        - long_running_transactions
+        - list_locks
+        - list_database_stats
+        - list_pg_settings
+    optimize:
+        - list_pg_settings
+        - list_memory_configurations
         - list_available_extensions
         - list_installed_extensions
         - list_autovacuum_configurations
         - list_columnar_configurations
         - list_columnar_recommended_columns
-        - list_memory_configurations
+    health:
         - list_top_bloated_tables
-        - list_replication_slots
         - list_invalid_indexes
-        - get_query_plan
-        - list_views
-        - list_schemas
+        - list_table_stats
+        - list_tablespaces
         - database_overview
-        - list_triggers
-        - list_indexes
-        - list_sequences
-        - long_running_transactions
-        - list_locks
+        - list_autovacuum_configurations
+    replication:
         - replication_stats
-        - list_query_stats
-        - get_column_cardinality
+        - list_replication_slots
         - list_publication_tables
-        - list_tablespaces
-        - list_pg_settings
-        - list_database_stats
+        - database_overview
+    access-control:
         - list_roles
-        - list_table_stats
-        - list_stored_procedure
+        - list_pg_settings
+        - database_overview
+