coder/cli/state.go

package cli

import (
	"fmt"
	"io"
	"os"
	"strconv"

	"github.com/coder/coder/v2/cli/cliui"
	"github.com/coder/coder/v2/codersdk"
	"github.com/coder/serpent"
)

func (r *RootCmd) state() *serpent.Command {
	cmd := &serpent.Command{
		Use:   "state",
		Short: "Manually manage Terraform state to fix broken workspaces",
		Handler: func(inv *serpent.Invocation) error {
			return inv.Command.HelpHandler(inv)
		},
		Children: []*serpent.Command{
			r.statePull(),
			r.statePush(),
		},
	}
	return cmd
}

func (r *RootCmd) statePull() *serpent.Command {
	var buildNumber int64
	cmd := &serpent.Command{
		Use:   "pull <workspace> [file]",
		Short: "Pull a Terraform state file from a workspace.",
		Middleware: serpent.Chain(
			serpent.RequireRangeArgs(1, 2),
		),
		Handler: func(inv *serpent.Invocation) error {
			client, err := r.InitClient(inv)
			if err != nil {
				return err
			}
			var build codersdk.WorkspaceBuild
			if buildNumber == 0 {
				workspace, err := client.ResolveWorkspace(inv.Context(), inv.Args[0])
				if err != nil {
					return err
				}
				build = workspace.LatestBuild
			} else {
				owner, workspace, err := codersdk.SplitWorkspaceIdentifier(inv.Args[0])
				if err != nil {
					return err
				}
				build, err = client.WorkspaceBuildByUsernameAndWorkspaceNameAndBuildNumber(inv.Context(), owner, workspace, strconv.FormatInt(buildNumber, 10))
				if err != nil {
					return err
				}
			}

			state, err := client.WorkspaceBuildState(inv.Context(), build.ID)
			if err != nil {
				return err
			}

			if len(inv.Args) < 2 {
				_, _ = fmt.Fprintln(inv.Stdout, string(state))
				return nil
			}

			return os.WriteFile(inv.Args[1], state, 0o600)
		},
	}
	cmd.Options = serpent.OptionSet{
		buildNumberOption(&buildNumber),
	}
	return cmd
}

func buildNumberOption(n *int64) serpent.Option {
	return serpent.Option{
		Flag:          "build",
		FlagShorthand: "b",
		Description:   "Specify a workspace build to target by name. Defaults to latest.",
		Value:         serpent.Int64Of(n),
	}
}

func (r *RootCmd) statePush() *serpent.Command {
	var buildNumber int64
	var noBuild bool
	cmd := &serpent.Command{
		Use:   "push <workspace> <file>",
		Short: "Push a Terraform state file to a workspace.",
		Middleware: serpent.Chain(
			serpent.RequireNArgs(2),
		),
		Handler: func(inv *serpent.Invocation) error {
			client, err := r.InitClient(inv)
			if err != nil {
				return err
			}
			workspace, err := client.ResolveWorkspace(inv.Context(), inv.Args[0])
			if err != nil {
				return err
			}
			var build codersdk.WorkspaceBuild
			if buildNumber == 0 {
				build = workspace.LatestBuild
			} else {
				owner, workspace, err := codersdk.SplitWorkspaceIdentifier(inv.Args[0])
				if err != nil {
					return err
				}
				build, err = client.WorkspaceBuildByUsernameAndWorkspaceNameAndBuildNumber(inv.Context(), owner, workspace, strconv.FormatInt((buildNumber), 10))
				if err != nil {
					return err
				}
			}

			var state []byte
			if inv.Args[1] == "-" {
				state, err = io.ReadAll(inv.Stdin)
			} else {
				state, err = os.ReadFile(inv.Args[1])
			}
			if err != nil {
				return err
			}

			if noBuild {
				// Update state directly without triggering a build.
				err = client.UpdateWorkspaceBuildState(inv.Context(), build.ID, state)
				if err != nil {
					return err
				}
				_, _ = fmt.Fprintln(inv.Stdout, "State updated successfully.")
				return nil
			}

			build, err = client.CreateWorkspaceBuild(inv.Context(), workspace.ID, codersdk.CreateWorkspaceBuildRequest{
				TemplateVersionID: build.TemplateVersionID,
				Transition:        build.Transition,
				ProvisionerState:  state,
			})
			if err != nil {
				return err
			}
			return cliui.WorkspaceBuild(inv.Context(), inv.Stderr, client, build.ID)
		},
	}
	cmd.Options = serpent.OptionSet{
		buildNumberOption(&buildNumber),
		{
			Flag:          "no-build",
			FlagShorthand: "n",
			Description:   "Update the state without triggering a workspace build. Useful for state-only migrations.",
			Value:         serpent.BoolOf(&noBuild),
		},
	}
	return cmd
}