diff --git a/cmd/root.go b/cmd/root.go index 4c373bf181..dfac1c250f 100644 --- a/cmd/root.go +++ b/cmd/root.go @@ -315,15 +315,15 @@ func Execute() { type Command struct { *cobra.Command - cfg server.ServerConfig - logger log.Logger - tools_file string - tools_files []string - tools_folder string - prebuiltConfig string - inStream io.Reader - outStream io.Writer - errStream io.Writer + cfg server.ServerConfig + logger log.Logger + tools_file string + tools_files []string + tools_folder string + prebuiltConfigs []string + inStream io.Reader + outStream io.Writer + errStream io.Writer } // NewCommand returns a Command object representing an invocation of the CLI. @@ -376,10 +376,10 @@ func NewCommand(opts ...Option) *Command { flags.StringVar(&cmd.cfg.TelemetryServiceName, "telemetry-service-name", "toolbox", "Sets the value of the service.name resource attribute for telemetry data.") // Fetch prebuilt tools sources to customize the help description prebuiltHelp := fmt.Sprintf( - "Use a prebuilt tool configuration by source type. Allowed: '%s'.", + "Use a prebuilt tool configuration by source type. Allowed: '%s'. Can be specified multiple times.", strings.Join(prebuiltconfigs.GetPrebuiltSources(), "', '"), ) - flags.StringVar(&cmd.prebuiltConfig, "prebuilt", "", prebuiltHelp) + flags.StringSliceVar(&cmd.prebuiltConfigs, "prebuilt", []string{}, prebuiltHelp) flags.BoolVar(&cmd.cfg.Stdio, "stdio", false, "Listens via MCP STDIO instead of acting as a remote HTTP server.") flags.BoolVar(&cmd.cfg.DisableReload, "disable-reload", false, "Disables dynamic reloading of tools file.") flags.BoolVar(&cmd.cfg.UI, "ui", false, "Launches the Toolbox UI web server.") @@ -867,24 +867,32 @@ func run(cmd *Command) error { var allToolsFiles []ToolsFile // Load Prebuilt Configuration - if cmd.prebuiltConfig != "" { - buf, err := prebuiltconfigs.Get(cmd.prebuiltConfig) - if err != nil { - cmd.logger.ErrorContext(ctx, err.Error()) - return err - } - logMsg := fmt.Sprint("Using prebuilt tool configuration for ", cmd.prebuiltConfig) - cmd.logger.InfoContext(ctx, logMsg) - // Append prebuilt.source to Version string for the User Agent - cmd.cfg.Version += "+prebuilt." + cmd.prebuiltConfig - parsed, err := parseToolsFile(ctx, buf) - if err != nil { - errMsg := fmt.Errorf("unable to parse prebuilt tool configuration: %w", err) - cmd.logger.ErrorContext(ctx, errMsg.Error()) - return errMsg + if len(cmd.prebuiltConfigs) > 0 { + slices.Sort(cmd.prebuiltConfigs) + sourcesList := strings.Join(cmd.prebuiltConfigs, ", ") + logMsg := fmt.Sprintf("Using prebuilt tool configurations for: %s", sourcesList) + cmd.logger.InfoContext(ctx, logMsg) + + for _, configName := range cmd.prebuiltConfigs { + buf, err := prebuiltconfigs.Get(configName) + if err != nil { + cmd.logger.ErrorContext(ctx, err.Error()) + return err + } + + // Update version string + cmd.cfg.Version += "+prebuilt." + configName + + // Parse into ToolsFile struct + parsed, err := parseToolsFile(ctx, buf) + if err != nil { + errMsg := fmt.Errorf("unable to parse prebuilt tool configuration for '%s': %w", configName, err) + cmd.logger.ErrorContext(ctx, errMsg.Error()) + return errMsg + } + allToolsFiles = append(allToolsFiles, parsed) } - allToolsFiles = append(allToolsFiles, parsed) } // Determine if Custom Files should be loaded @@ -892,7 +900,7 @@ func run(cmd *Command) error { isCustomConfigured := cmd.tools_file != "" || len(cmd.tools_files) > 0 || cmd.tools_folder != "" // Determine if default 'tools.yaml' should be used (No prebuilt AND No custom flags) - useDefaultToolsFile := cmd.prebuiltConfig == "" && !isCustomConfigured + useDefaultToolsFile := len(cmd.prebuiltConfigs) == 0 && !isCustomConfigured if useDefaultToolsFile { cmd.tools_file = "tools.yaml" diff --git a/cmd/root_test.go b/cmd/root_test.go index e825a36366..17058d18ff 100644 --- a/cmd/root_test.go +++ b/cmd/root_test.go @@ -430,17 +430,27 @@ func TestPrebuiltFlag(t *testing.T) { tcs := []struct { desc string args []string - want string + want []string }{ { desc: "default value", args: []string{}, - want: "", + want: []string{}, }, { - desc: "custom pre built flag", - args: []string{"--tools-file", "alloydb"}, - want: "alloydb", + desc: "single prebuilt flag", + args: []string{"--prebuilt", "alloydb"}, + want: []string{"alloydb"}, + }, + { + desc: "multiple prebuilt flags", + args: []string{"--prebuilt", "alloydb", "--prebuilt", "bigquery"}, + want: []string{"alloydb", "bigquery"}, + }, + { + desc: "comma separated prebuilt flags", + args: []string{"--prebuilt", "alloydb,bigquery"}, + want: []string{"alloydb", "bigquery"}, }, } for _, tc := range tcs { @@ -449,8 +459,8 @@ func TestPrebuiltFlag(t *testing.T) { if err != nil { t.Fatalf("unexpected error invoking command: %s", err) } - if c.tools_file != tc.want { - t.Fatalf("got %v, want %v", c.cfg, tc.want) + if diff := cmp.Diff(c.prebuiltConfigs, tc.want); diff != "" { + t.Fatalf("got %v, want %v, diff %s", c.prebuiltConfigs, tc.want, diff) } }) } @@ -2073,6 +2083,12 @@ authSources: return nil }, }, + { + desc: "sqlite called twice error", + args: []string{"--prebuilt", "sqlite", "--prebuilt", "sqlite"}, + wantErr: true, + errString: "resource conflicts detected", + }, { desc: "tool conflict error", args: []string{"--prebuilt", "sqlite", "--tools-file", toolConflictFile}, diff --git a/docs/en/reference/cli.md b/docs/en/reference/cli.md index a1e6a57985..686dbc0c73 100644 --- a/docs/en/reference/cli.md +++ b/docs/en/reference/cli.md @@ -16,7 +16,7 @@ description: > | | `--log-level` | Specify the minimum level logged. Allowed: 'DEBUG', 'INFO', 'WARN', 'ERROR'. | `info` | | | `--logging-format` | Specify logging format to use. Allowed: 'standard' or 'JSON'. | `standard` | | `-p` | `--port` | Port the server will listen on. | `5000` | -| | `--prebuilt` | Use a prebuilt tool configuration by source type. See [Prebuilt Tools Reference](prebuilt-tools.md) for allowed values. | | +| | `--prebuilt` | Use one or more prebuilt tool configuration by source type. See [Prebuilt Tools Reference](prebuilt-tools.md) for allowed values. | | | | `--stdio` | Listens via MCP STDIO instead of acting as a remote HTTP server. | | | | `--telemetry-gcp` | Enable exporting directly to Google Cloud Monitoring. | | | | `--telemetry-otlp` | Enable exporting using OpenTelemetry Protocol (OTLP) to the specified endpoint (e.g. 'http://127.0.0.1:4318') | | @@ -51,6 +51,11 @@ description: > # Server with prebuilt + custom tools configurations ./toolbox --tools-file tools.yaml --prebuilt alloydb-postgres + +# Server with multiple prebuilt tools configurations +./toolbox --prebuilt alloydb-postgres,alloydb-postgres-admin +# OR +./toolbox --prebuilt alloydb-postgres --prebuilt alloydb-postgres-admin ``` ### Tool Configuration Sources @@ -71,7 +76,7 @@ The CLI supports multiple mutually exclusive ways to specify tool configurations **Prebuilt Configurations:** -- `--prebuilt`: Use predefined configurations for specific database types (e.g., +- `--prebuilt`: Use one or more predefined configurations for specific database types (e.g., 'bigquery', 'postgres', 'spanner'). See [Prebuilt Tools Reference](prebuilt-tools.md) for allowed values. diff --git a/docs/en/reference/prebuilt-tools.md b/docs/en/reference/prebuilt-tools.md index d4f347193b..a539af9fb2 100644 --- a/docs/en/reference/prebuilt-tools.md +++ b/docs/en/reference/prebuilt-tools.md @@ -16,6 +16,9 @@ details on how to connect your AI tools (IDEs) to databases via Toolbox and MCP. {{< notice tip >}} You can now use `--prebuilt` along `--tools-file`, `--tools-files`, or `--tools-folder` to combine prebuilt configs with custom tools. + +You can also combine multiple prebuilt configs. + See [Usage Examples](../reference/cli.md#examples). {{< /notice >}}