feat: add dry-run and prune flag for apply command

Author: Yuan325

cmd/internal/apply/command.go

@@ -27,14 +27,15 @@ import (
 
 	"github.com/goccy/go-yaml"
 	"github.com/googleapis/genai-toolbox/cmd/internal"
-	"github.com/googleapis/genai-toolbox/internal/log"
 	"github.com/spf13/cobra"
 )
 
 type applyCmd struct {
 	*cobra.Command
 	port    int
 	address string
+	dryRun  bool
+	prune   bool
 }
 
 func NewCommand(opts *internal.ToolboxOptions) *cobra.Command {
@@ -48,6 +49,8 @@ func NewCommand(opts *internal.ToolboxOptions) *cobra.Command {
 	internal.ConfigFileFlags(flags, opts)
 	flags.StringVarP(&cmd.address, "address", "a", "127.0.0.1", "Address of the server that is running on.")
 	flags.IntVarP(&cmd.port, "port", "p", 5000, "Port of the server that is listening on.")
+	flags.BoolVar(&cmd.dryRun, "dry-run", false, "Simulate the apply process and report intended updates.")
+	flags.BoolVar(&cmd.prune, "prune", false, "Remove resources from the server that are not defined in the provided YAML files.")
 	cmd.RunE = func(*cobra.Command, []string) error { return runApply(cmd, opts) }
 	return cmd.Command
 }
@@ -92,6 +95,16 @@ func (p primitives) Remove(kind, name string) {
 	}
 }
 
+func (p primitives) GetRemaining() []struct{ kind, name string } {
+	var remaining []struct{ kind, name string }
+	for kind, names := range p {
+		for name := range names {
+			remaining = append(remaining, struct{ kind, name string }{kind, name})
+		}
+	}
+	return remaining
+}
+
 // adminRequest is a generic helper for admin api requests.
 func adminRequest(ctx context.Context, method, url string, body any) ([]byte, error) {
 	var bodyReader io.Reader
@@ -171,6 +184,17 @@ func applyPrimitive(ctx context.Context, address string, port int, kind, name st
 	return nil
 }
 
+func deletePrimitive(ctx context.Context, address string, port int, kind, name string) error {
+	url := fmt.Sprintf("http://%s:%d/admin/%s/%s", address, port, kind, name)
+
+	_, err := adminRequest(ctx, http.MethodDelete, url, nil)
+	if err != nil {
+		return err
+	}
+
+	return nil
+}
+
 func runApply(cmd *applyCmd, opts *internal.ToolboxOptions) error {
 	ctx, cancel := context.WithCancel(cmd.Context())
 	defer cancel()
@@ -196,9 +220,11 @@ func runApply(cmd *applyCmd, opts *internal.ToolboxOptions) error {
 		return err
 	}
 
-	opts.Logger.InfoContext(ctx, "starting apply sequence", "count", len(filePaths))
+	if cmd.dryRun {
+		opts.Logger.InfoContext(ctx, "simulating dry run")
+	}
 	for _, filePath := range filePaths {
-		if err := processFile(ctx, opts.Logger, filePath, p, cmd.address, cmd.port); err != nil {
+		if err := processFile(ctx, opts, filePath, p, cmd); err != nil {
 			opts.Logger.ErrorContext(ctx, err.Error())
 			return err
 		}
@@ -207,14 +233,15 @@ func runApply(cmd *applyCmd, opts *internal.ToolboxOptions) error {
 	return nil
 }
 
-func processFile(ctx context.Context, logger log.Logger, path string, p primitives, address string, port int) error {
+func processFile(ctx context.Context, opts *internal.ToolboxOptions, path string, p primitives, cmd *applyCmd) error {
 	f, err := os.Open(path)
 	if err != nil {
 		return fmt.Errorf("unable to open file at %q: %w", path, err)
 	}
 	defer f.Close() // Safe closure
 
 	decoder := yaml.NewDecoder(f)
+	opts.Logger.InfoContext(ctx, "starting apply sequence", "count", len(path))
 	// loop through documents with the `---` separator
 	for {
 		var doc map[string]any
@@ -231,15 +258,15 @@ func processFile(ctx context.Context, logger log.Logger, path string, p primitiv
 		kind, kOk := doc["kind"].(string)
 		name, nOk := doc["name"].(string)
 		if !kOk || !nOk || kind == "" || name == "" {
-			logger.WarnContext(ctx, fmt.Sprintf("invalid primitive schema: missing metadata in %s: kind and name are required", path))
+			opts.Logger.WarnContext(ctx, fmt.Sprintf("invalid primitive schema: missing metadata in %s: kind and name are required", path))
 			continue
 		}
 
 		delete(doc, "kind")
 
 		if p.Exists(kind, name) {
 			p.Remove(kind, name)
-			remoteBody, err := getPrimitiveByName(ctx, address, port, kind, name)
+			remoteBody, err := getPrimitiveByName(ctx, cmd.address, cmd.port, kind, name)
 			if err != nil {
 				return err
 			}
@@ -253,19 +280,41 @@ func processFile(ctx context.Context, logger log.Logger, path string, p primitiv
 			}
 
 			if bytes.Equal(localJSON, remoteJSON) {
-				logger.DebugContext(ctx, "skipping: no changes detected", "kind", kind, "name", name)
+				opts.Logger.DebugContext(ctx, "skipping: no changes detected", "kind", kind, "name", name)
 				continue
 			}
-			logger.DebugContext(ctx, "change detected, updating resource", "kind", kind, "name", name)
+			opts.Logger.DebugContext(ctx, "change detected, updating resource", "kind", kind, "name", name)
 		}
 
 		// TODO: check --prune flag: if prune, delete primitives that are left
 		// in the primitive list
-		// TODO: check for --dry-run flag.
 
-		if err := applyPrimitive(ctx, address, port, kind, name, doc); err != nil {
+		if cmd.dryRun {
+			fmt.Fprintf(opts.IOStreams.Out, "[dry-run] would update %s: %s\n", kind, name)
+			continue
+		}
+		if err := applyPrimitive(ctx, cmd.address, cmd.port, kind, name, doc); err != nil {
 			return err
 		}
 	}
+	if cmd.prune {
+		remaining := p.GetRemaining()
+		opts.Logger.InfoContext(ctx, "starting prune sequence", "count", len(remaining))
+		for _, resource := range remaining {
+			if cmd.dryRun {
+				fmt.Fprintf(opts.IOStreams.Out, "[dry-run] would delete %s: %s\n", resource.kind, resource.name)
+				continue
+			}
+
+			// Perform the actual deletion
+			if err := deletePrimitive(ctx, cmd.address, cmd.port, resource.kind, resource.name); err != nil {
+				return fmt.Errorf("failed to prune %s/%s: %w", resource.kind, resource.name, err)
+			}
+
+			// Remove from map so we don't try to delete it again
+			// if processFile is called for multiple files
+			p.Remove(resource.kind, resource.name)
+		}
+	}
 	return nil
 }