{cli, filestate}: Add 'state upgrade' command

Adds a new `pulumi state upgrade` command
that will upgrade the state for the current state storage
to the latest version supported by that format.
This is a no-op for httpstate,
but for filestate, this will migrate from unversioned storage
to the project-based versioned format.

The upgrade command is backed by a new Upgrade operation
on the filestate Backend.
The upgrade operation operates by effectively renaming stacks
from legacy store to the new project store.

UX notes:

- We'll warn about stacks we failed to upgrade
  but keep going otherwise.
- We'll print a loud warning before upgrading,
  informing users that the state will not be usable
  from older CLIs.

Extracted from #12134

Co-authored-by: Abhinav Gupta <abhinav@pulumi.com>

changelog/pending/20230316--cli-state--add-upgrade-subcommand-to-upgrade-a-pulumi-self-managed-state-to-use-project-layout.yaml

@@ -0,0 +1,4 @@
+changes:
+- type: feat
+  scope: cli/state
+  description: Add 'upgrade' subcommand to upgrade a Pulumi self-managed state to use project layout.

pkg/backend/filestate/backend.go

@@ -60,6 +60,7 @@ import (
 	"github.com/pulumi/pulumi/sdk/v3/go/common/util/contract"
 	"github.com/pulumi/pulumi/sdk/v3/go/common/util/result"
 	"github.com/pulumi/pulumi/sdk/v3/go/common/workspace"
+	"gopkg.in/yaml.v3"
 )
 
 // PulumiFilestateGzipEnvVar is an env var that must be truthy
@@ -93,6 +94,9 @@ var (
 type Backend interface {
 	backend.Backend
 	local() // at the moment, no local specific info, so just use a marker function.
+
+	// Upgrade to the latest state store version.
+	Upgrade(ctx context.Context) error
 }
 
 type localBackend struct {
@@ -296,6 +300,7 @@ func New(ctx context.Context, d diag.Sink, originalURL string, project *workspac
 	}
 	// Otherwise, warn about any old stack files.
 	// This is possible if a user creates a new stack with a new CLI,
+	// or migrates it to project mode with `pulumi state upgrade`,
 	// but someone else interacts with the same state with an old CLI.
 
 	refs, err := newLegacyReferenceStore(wbucket).ListReferences()
@@ -312,11 +317,75 @@ func New(ctx context.Context, d diag.Sink, originalURL string, project *workspac
 	for _, ref := range refs {
 		fmt.Fprintf(&msg, "  - %s\n", ref.Name())
 	}
+	msg.WriteString("Please run 'pulumi state upgrade' to migrate them to the new format.\n")
 	msg.WriteString("Set PULUMI_SELF_MANAGED_STATE_NO_LEGACY_WARNING=1 to disable this warning.")
 	d.Warningf(diag.Message("", msg.String()))
 	return backend, nil
 }
 
+func (b *localBackend) Upgrade(ctx context.Context) error {
+	// We don't use the existing b.store because
+	// this may already be a projectReferenceStore
+	// with new legacy files introduced to it accidentally.
+	olds, err := newLegacyReferenceStore(b.bucket).ListReferences()
+	if err != nil {
+		return fmt.Errorf("read old references: %w", err)
+	}
+
+	newStore := newProjectReferenceStore(b.bucket, b.currentProject.Load)
+	var upgraded int
+	for _, old := range olds {
+		if err := b.upgradeStack(ctx, newStore, old); err != nil {
+			b.d.Warningf(diag.Message("", "Skipping stack %q: %v"), old, err)
+			continue
+		}
+		upgraded++
+	}
+
+	pulumiYaml, err := yaml.Marshal(&pulumiMeta{Version: 1})
+	contract.AssertNoErrorf(err, "Could not marshal filestate.pulumiState to yaml")
+	if err = b.bucket.WriteAll(ctx, "Pulumi.yaml", pulumiYaml, nil); err != nil {
+		return fmt.Errorf("could not write 'Pulumi.yaml': %w", err)
+	}
+	b.store = newStore
+
+	b.d.Infoerrf(diag.Message("", "Upgraded %d stack(s) to project mode"), upgraded)
+	return nil
+}
+
+// upgradeStack upgrades a single stack to use the provided projectReferenceStore.
+func (b *localBackend) upgradeStack(
+	ctx context.Context,
+	newStore *projectReferenceStore,
+	old *localBackendReference,
+) error {
+	contract.Requiref(old.project == "", "old.project", "must be empty")
+
+	chk, err := b.getCheckpoint(old)
+	if err != nil {
+		return err
+	}
+
+	// Try and find the project name from _any_ resource URN
+	var project tokens.Name
+	if chk.Latest != nil {
+		for _, res := range chk.Latest.Resources {
+			project = tokens.Name(res.URN.Project())
+			break
+		}
+	}
+	if project == "" {
+		return errors.New("no project found")
+	}
+
+	new := newStore.newReference(project, old.Name())
+	if err := b.renameStack(ctx, old, new); err != nil {
+		return fmt.Errorf("rename to %v: %w", new, err)
+	}
+
+	return nil
+}
+
 // massageBlobPath takes the path the user provided and converts it to an appropriate form go-cloud
 // can support.  Importantly, s3/azblob/gs paths should not be be touched. This will only affect
 // file:// paths which have a few oddities around them that we want to ensure work properly.

pkg/backend/filestate/backend_test.go

@@ -788,6 +788,7 @@ func TestNew_legacyFileWarning(t *testing.T) {
 			wantOut: "warning: Found legacy stack files in state store:\n" +
 				"  - a\n" +
 				"  - b\n" +
+				"Please run 'pulumi state upgrade' to migrate them to the new format.\n" +
 				"Set PULUMI_SELF_MANAGED_STATE_NO_LEGACY_WARNING=1 to disable this warning.\n",
 		},
 		{
@@ -832,6 +833,118 @@ func TestNew_legacyFileWarning(t *testing.T) {
 	}
 }
 
+func TestLegacyUpgrade(t *testing.T) {
+	t.Parallel()
+
+	// Make a dummy stack file in the legacy location
+	tmpDir := t.TempDir()
+	err := os.MkdirAll(path.Join(tmpDir, ".pulumi", "stacks"), os.ModePerm)
+	require.NoError(t, err)
+	err = os.WriteFile(path.Join(tmpDir, ".pulumi", "stacks", "a.json"), []byte(`{
+		"latest": {
+			"resources": [
+				{
+					"type": "package:module:resource",
+					"urn": "urn:pulumi:stack::project::package:module:resource::name"
+				}
+			]
+		}
+	}`), os.ModePerm)
+	require.NoError(t, err)
+
+	var output bytes.Buffer
+	sink := diag.DefaultSink(&output, &output, diag.FormatOptions{Color: colors.Never})
+
+	// Login to a temp dir filestate backend
+	ctx := context.Background()
+	b, err := New(ctx, sink, "file://"+filepath.ToSlash(tmpDir), nil)
+	require.NoError(t, err)
+	// Check the backend says it's NOT in project mode
+	lb, ok := b.(*localBackend)
+	assert.True(t, ok)
+	assert.NotNil(t, lb)
+	assert.IsType(t, &legacyReferenceStore{}, lb.store)
+
+	err = lb.Upgrade(ctx)
+	require.NoError(t, err)
+	assert.IsType(t, &projectReferenceStore{}, lb.store)
+
+	assert.Contains(t, output.String(), "Upgraded 1 stack(s) to project mode")
+
+	// Check that a has been moved
+	aStackRef, err := lb.parseStackReference("organization/project/a")
+	require.NoError(t, err)
+	stackFileExists, err := lb.bucket.Exists(ctx, lb.stackPath(aStackRef))
+	require.NoError(t, err)
+	assert.True(t, stackFileExists)
+
+	// Write b.json and upgrade again
+	err = os.WriteFile(path.Join(tmpDir, ".pulumi", "stacks", "b.json"), []byte(`{
+		"latest": {
+			"resources": [
+				{
+					"type": "package:module:resource",
+					"urn": "urn:pulumi:stack::other-project::package:module:resource::name"
+				}
+			]
+		}
+	}`), os.ModePerm)
+	require.NoError(t, err)
+
+	err = lb.Upgrade(ctx)
+	require.NoError(t, err)
+
+	// Check that b has been moved
+	bStackRef, err := lb.parseStackReference("organization/other-project/b")
+	require.NoError(t, err)
+	stackFileExists, err = lb.bucket.Exists(ctx, lb.stackPath(bStackRef))
+	require.NoError(t, err)
+	assert.True(t, stackFileExists)
+}
+
+func TestLegacyUpgrade_partial(t *testing.T) {
+	t.Parallel()
+
+	// Verifies that we can upgrade a subset of stacks.
+
+	stateDir := t.TempDir()
+	bucket, err := fileblob.OpenBucket(stateDir, nil)
+	require.NoError(t, err)
+
+	ctx := context.Background()
+	require.NoError(t,
+		bucket.WriteAll(ctx, ".pulumi/stacks/foo.json", []byte(`{
+		"latest": {
+			"resources": [
+				{
+					"type": "package:module:resource",
+					"urn": "urn:pulumi:stack::project::package:module:resource::name"
+				}
+			]
+		}
+	}`), nil))
+	require.NoError(t,
+		// no resources, can't guess project name
+		bucket.WriteAll(ctx, ".pulumi/stacks/bar.json",
+			[]byte(`{"latest": {"resources": []}}`), nil))
+
+	var buff bytes.Buffer
+	sink := diag.DefaultSink(io.Discard, &buff, diag.FormatOptions{Color: colors.Never})
+	b, err := New(ctx, sink, "file://"+filepath.ToSlash(stateDir), nil)
+	require.NoError(t, err)
+
+	require.NoError(t, b.Upgrade(ctx))
+	assert.Contains(t, buff.String(), `Skipping stack "bar": no project found`)
+
+	exists, err := bucket.Exists(ctx, ".pulumi/stacks/project/foo.json")
+	require.NoError(t, err)
+	assert.True(t, exists, "foo was not migrated")
+
+	ref, err := b.ParseStackReference("organization/project/foo")
+	require.NoError(t, err)
+	assert.Equal(t, tokens.QName("organization/project/foo"), ref.FullyQualifiedName())
+}
+
 func TestNew_unsupportedStoreVersion(t *testing.T) {
 	t.Parallel()
 

pkg/cmd/pulumi/pulumi.go

@@ -652,18 +652,27 @@ func isDevVersion(s semver.Version) bool {
 }
 
 func confirmPrompt(prompt string, name string, opts display.Options) bool {
+	out := opts.Stdout
+	if out == nil {
+		out = os.Stdout
+	}
+	in := opts.Stdin
+	if in == nil {
+		in = os.Stdin
+	}
+
 	if prompt != "" {
-		fmt.Print(
+		fmt.Fprint(out,
 			opts.Color.Colorize(
 				fmt.Sprintf("%s%s%s\n", colors.SpecAttention, prompt, colors.Reset)))
 	}
 
-	fmt.Print(
+	fmt.Fprint(out,
 		opts.Color.Colorize(
 			fmt.Sprintf("%sPlease confirm that this is what you'd like to do by typing `%s%s%s`:%s ",
 				colors.SpecAttention, colors.SpecPrompt, name, colors.SpecAttention, colors.Reset)))
 
-	reader := bufio.NewReader(os.Stdin)
+	reader := bufio.NewReader(in)
 	line, _ := reader.ReadString('\n')
 	return strings.TrimSpace(line) == name
 }

pkg/cmd/pulumi/state.go

@@ -1,4 +1,4 @@
-// Copyright 2016-2022, Pulumi Corporation.
+// Copyright 2016-2023, Pulumi Corporation.
 //
 // Licensed under the Apache License, Version 2.0 (the "License");
 // you may not use this file except in compliance with the License.
@@ -50,6 +50,7 @@ troubleshooting a stack or when performing specific edits that otherwise would r
 	cmd.AddCommand(newStateDeleteCommand())
 	cmd.AddCommand(newStateUnprotectCommand())
 	cmd.AddCommand(newStateRenameCommand())
+	cmd.AddCommand(newStateUpgradeCommand())
 	return cmd
 }
 

pkg/cmd/pulumi/state_upgrade.go

@@ -0,0 +1,105 @@
+// Copyright 2016-2023, Pulumi Corporation.
+//
+// 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 main
+
+import (
+	"context"
+	"fmt"
+	"io"
+	"os"
+
+	"github.com/pulumi/pulumi/pkg/v3/backend"
+	"github.com/pulumi/pulumi/pkg/v3/backend/display"
+	"github.com/pulumi/pulumi/pkg/v3/backend/filestate"
+	"github.com/pulumi/pulumi/sdk/v3/go/common/util/cmdutil"
+	"github.com/pulumi/pulumi/sdk/v3/go/common/util/result"
+	"github.com/pulumi/pulumi/sdk/v3/go/common/workspace"
+
+	"github.com/spf13/cobra"
+)
+
+func newStateUpgradeCommand() *cobra.Command {
+	var sucmd stateUpgradeCmd
+	cmd := &cobra.Command{
+		Use:   "upgrade",
+		Short: "Migrates the current backend to the latest supported version",
+		Long: `Migrates the current backend to the latest supported version
+
+This only has an effect on self-managed backends.
+`,
+		Args: cmdutil.NoArgs,
+		Run: cmdutil.RunResultFunc(func(cmd *cobra.Command, args []string) result.Result {
+			if err := sucmd.Run(commandContext()); err != nil {
+				return result.FromError(err)
+			}
+			return nil
+		}),
+	}
+	return cmd
+}
+
+// stateUpgradeCmd implements the 'pulumi state upgrade' command.
+type stateUpgradeCmd struct {
+	Stdin  io.Reader // defaults to os.Stdin
+	Stdout io.Writer // defaults to os.Stdout
+
+	// Used to mock out the currentBackend function for testing.
+	// Defaults to currentBackend function.
+	currentBackend func(context.Context, *workspace.Project, display.Options) (backend.Backend, error)
+}
+
+func (cmd *stateUpgradeCmd) Run(ctx context.Context) error {
+	if cmd.Stdout == nil {
+		cmd.Stdout = os.Stdout
+	}
+	if cmd.Stdin == nil {
+		cmd.Stdin = os.Stdin
+	}
+
+	if cmd.currentBackend == nil {
+		cmd.currentBackend = currentBackend
+	}
+	currentBackend := cmd.currentBackend // shadow top-level currentBackend
+
+	dopts := display.Options{
+		Color:  cmdutil.GetGlobalColorization(),
+		Stdin:  cmd.Stdin,
+		Stdout: cmd.Stdout,
+	}
+
+	b, err := currentBackend(ctx, nil, dopts)
+	if err != nil {
+		return err
+	}
+
+	lb, ok := b.(filestate.Backend)
+	if !ok {
+		// Only the file state backend supports upgrades,
+		// but we don't want to error out here.
+		// Report the no-op.
+		fmt.Fprintln(cmd.Stdout, "Nothing to do")
+		return nil
+	}
+
+	prompt := "This will upgrade the current backend to the latest supported version.\n" +
+		"Older versions of Pulumi will not be able to read the new format.\n" +
+		"Are you sure you want to proceed?"
+	if !confirmPrompt(prompt, "yes", dopts) {
+		fmt.Fprintln(cmd.Stdout, "Upgrade cancelled")
+		return nil
+	}
+
+	return lb.Upgrade(ctx)
+}