Support workspace for state migrator

tfmigrate/state_import_action_test.go

@@ -44,7 +44,7 @@ resource "aws_iam_user" "baz" {
 		NewStateImportAction("aws_iam_user.baz", "baz"),
 	}
 
-	m := NewStateMigrator(tf.Dir(), actions, &MigratorOption{}, false)
+	m := NewStateMigrator(tf.Dir(), "default", actions, &MigratorOption{}, false)
 	err = m.Plan(ctx)
 	if err != nil {
 		t.Fatalf("failed to run migrator plan: %s", err)

tfmigrate/state_migrator.go

@@ -27,6 +27,8 @@ type StateMigratorConfig struct {
 	// Force option controls behaviour in case of unexpected diff in plan.
 	// When set forces applying even if plan shows diff.
 	Force bool `hcl:"force,optional"`
+	// Workspace is the state workspace which the migration works with.
+	Workspace string `hcl:"workspace,optional"`
 }
 
 // StateMigratorConfig implements a MigratorConfig.
@@ -54,7 +56,12 @@ func (c *StateMigratorConfig) NewMigrator(o *MigratorOption) (Migrator, error) {
 		actions = append(actions, action)
 	}
 
-	return NewStateMigrator(dir, actions, o, c.Force), nil
+	//use default workspace if not specified by user
+	if len(c.Workspace) == 0 {
+		c.Workspace = "default"
+	}
+
+	return NewStateMigrator(dir, c.Workspace, actions, o, c.Force), nil
 }
 
 // StateMigrator implements the Migrator interface.
@@ -68,23 +75,26 @@ type StateMigrator struct {
 	o *MigratorOption
 	// force operation in case of unexpected diff
 	force bool
+	// workspace is the state workspace which the migration works with.
+	workspace string
 }
 
 var _ Migrator = (*StateMigrator)(nil)
 
 // NewStateMigrator returns a new StateMigrator instance.
-func NewStateMigrator(dir string, actions []StateAction, o *MigratorOption, force bool) *StateMigrator {
+func NewStateMigrator(dir string, workspace string, actions []StateAction, o *MigratorOption, force bool) *StateMigrator {
 	e := tfexec.NewExecutor(dir, os.Environ())
 	tf := tfexec.NewTerraformCLI(e)
 	if o != nil && len(o.ExecPath) > 0 {
 		tf.SetExecPath(o.ExecPath)
 	}
 
 	return &StateMigrator{
-		tf:      tf,
-		actions: actions,
-		o:       o,
-		force:   force,
+		tf:        tf,
+		actions:   actions,
+		o:         o,
+		force:     force,
+		workspace: workspace,
 	}
 }
 
@@ -94,7 +104,7 @@ func NewStateMigrator(dir string, actions []StateAction, o *MigratorOption, forc
 // the Migrator interface between a single and multi state migrator.
 func (m *StateMigrator) plan(ctx context.Context) (*tfexec.State, error) {
 	// setup work dir.
-	currentState, switchBackToRemotekFunc, err := setupWorkDir(ctx, m.tf, "default")
+	currentState, switchBackToRemotekFunc, err := setupWorkDir(ctx, m.tf, m.workspace)
 	if err != nil {
 		return nil, err
 	}

tfmigrate/state_migrator_test.go

@@ -50,6 +50,23 @@ func TestStateMigratorConfigNewMigrator(t *testing.T) {
 			},
 			ok: true,
 		},
+		{
+			desc: "valid in non-default workspace",
+			config: &StateMigratorConfig{
+				Dir: "dir1",
+				Actions: []string{
+					"mv aws_security_group.foo aws_security_group.foo2",
+					"mv aws_security_group.bar aws_security_group.bar2",
+					"rm aws_security_group.baz",
+					"import aws_security_group.qux qux",
+				},
+				Workspace: "workspace1",
+			},
+			o: &MigratorOption{
+				ExecPath: "direnv exec . terraform",
+			},
+			ok: true,
+		},
 		{
 			desc: "invalid action",
 			config: &StateMigratorConfig{
@@ -106,81 +123,99 @@ func TestStateMigratorConfigNewMigrator(t *testing.T) {
 func TestAccStateMigratorApply(t *testing.T) {
 	tfexec.SkipUnlessAcceptanceTestEnabled(t)
 
-	backend := tfexec.GetTestAccBackendS3Config(t.Name())
+	cases := []struct {
+		desc      string
+		workspace string
+	}{
+		{
+			desc:      "default workspace",
+			workspace: "default",
+		},
+		{
+			desc:      "non-default workspace",
+			workspace: "workspace1",
+		},
+	}
 
-	source := `
+	for _, tc := range cases {
+		t.Run(tc.desc, func(t *testing.T) {
+			backend := tfexec.GetTestAccBackendS3Config(t.Name())
+
+			source := `
 resource "aws_security_group" "foo" {}
 resource "aws_security_group" "bar" {}
 resource "aws_security_group" "baz" {}
 resource "aws_iam_user" "qux" {
 	name = "qux"
 }
 `
-	tf := tfexec.SetupTestAccWithApply(t, "default", backend+source)
-	ctx := context.Background()
+			tf := tfexec.SetupTestAccWithApply(t, tc.workspace, backend+source)
+			ctx := context.Background()
 
-	updatedSource := `
+			updatedSource := `
 resource "aws_security_group" "foo2" {}
 resource "aws_security_group" "baz" {}
 resource "aws_iam_user" "qux" {
 	name = "qux"
 }
 `
 
-	tfexec.UpdateTestAccSource(t, tf, backend+updatedSource)
+			tfexec.UpdateTestAccSource(t, tf, backend+updatedSource)
 
-	_, err := tf.StateRm(ctx, nil, []string{"aws_iam_user.qux"})
-	if err != nil {
-		t.Fatalf("failed to run terraform state rm: %s", err)
-	}
+			_, err := tf.StateRm(ctx, nil, []string{"aws_iam_user.qux"})
+			if err != nil {
+				t.Fatalf("failed to run terraform state rm: %s", err)
+			}
 
-	changed, err := tf.PlanHasChange(ctx, nil)
-	if err != nil {
-		t.Fatalf("failed to run PlanHasChange: %s", err)
-	}
-	if !changed {
-		t.Fatalf("expect to have changes")
-	}
+			changed, err := tf.PlanHasChange(ctx, nil)
+			if err != nil {
+				t.Fatalf("failed to run PlanHasChange: %s", err)
+			}
+			if !changed {
+				t.Fatalf("expect to have changes")
+			}
 
-	actions := []StateAction{
-		NewStateMvAction("aws_security_group.foo", "aws_security_group.foo2"),
-		NewStateRmAction([]string{"aws_security_group.bar"}),
-		NewStateImportAction("aws_iam_user.qux", "qux"),
-	}
+			actions := []StateAction{
+				NewStateMvAction("aws_security_group.foo", "aws_security_group.foo2"),
+				NewStateRmAction([]string{"aws_security_group.bar"}),
+				NewStateImportAction("aws_iam_user.qux", "qux"),
+			}
 
-	m := NewStateMigrator(tf.Dir(), actions, &MigratorOption{}, false)
-	err = m.Plan(ctx)
-	if err != nil {
-		t.Fatalf("failed to run migrator plan: %s", err)
-	}
+			m := NewStateMigrator(tf.Dir(), tc.workspace, actions, &MigratorOption{}, false)
+			err = m.Plan(ctx)
+			if err != nil {
+				t.Fatalf("failed to run migrator plan: %s", err)
+			}
 
-	err = m.Apply(ctx)
-	if err != nil {
-		t.Fatalf("failed to run migrator apply: %s", err)
-	}
+			err = m.Apply(ctx)
+			if err != nil {
+				t.Fatalf("failed to run migrator apply: %s", err)
+			}
 
-	got, err := tf.StateList(ctx, nil, nil)
-	if err != nil {
-		t.Fatalf("failed to run terraform state list: %s", err)
-	}
+			got, err := tf.StateList(ctx, nil, nil)
+			if err != nil {
+				t.Fatalf("failed to run terraform state list: %s", err)
+			}
 
-	want := []string{
-		"aws_security_group.foo2",
-		"aws_security_group.baz",
-		"aws_iam_user.qux",
-	}
-	sort.Strings(got)
-	sort.Strings(want)
-	if !reflect.DeepEqual(got, want) {
-		t.Errorf("got state: %v, want state: %v", got, want)
-	}
+			want := []string{
+				"aws_security_group.foo2",
+				"aws_security_group.baz",
+				"aws_iam_user.qux",
+			}
+			sort.Strings(got)
+			sort.Strings(want)
+			if !reflect.DeepEqual(got, want) {
+				t.Errorf("got state: %v, want state: %v", got, want)
+			}
 
-	changed, err = tf.PlanHasChange(ctx, nil)
-	if err != nil {
-		t.Fatalf("failed to run PlanHasChange: %s", err)
-	}
-	if changed {
-		t.Fatalf("expect not to have changes")
+			changed, err = tf.PlanHasChange(ctx, nil)
+			if err != nil {
+				t.Fatalf("failed to run PlanHasChange: %s", err)
+			}
+			if changed {
+				t.Fatalf("expect not to have changes")
+			}
+		})
 	}
 }
 
@@ -219,7 +254,7 @@ resource "aws_security_group" "baz" {}
 	o := &MigratorOption{}
 	o.PlanOut = "foo.tfplan"
 
-	m := NewStateMigrator(tf.Dir(), actions, o, true)
+	m := NewStateMigrator(tf.Dir(), "default", actions, o, true)
 	err = m.Plan(ctx)
 	if err != nil {
 		t.Fatalf("failed to run migrator plan: %s", err)

tfmigrate/state_mv_action_test.go

@@ -41,7 +41,7 @@ resource "aws_security_group" "baz" {}
 		NewStateMvAction("aws_security_group.bar", "aws_security_group.bar2"),
 	}
 
-	m := NewStateMigrator(tf.Dir(), actions, &MigratorOption{}, false)
+	m := NewStateMigrator(tf.Dir(), "default", actions, &MigratorOption{}, false)
 	err = m.Plan(ctx)
 	if err != nil {
 		t.Fatalf("failed to run migrator plan: %s", err)

tfmigrate/state_rm_action_test.go

@@ -40,7 +40,7 @@ resource "aws_security_group" "baz" {}
 		NewStateRmAction([]string{"aws_security_group.qux"}),
 	}
 
-	m := NewStateMigrator(tf.Dir(), actions, &MigratorOption{}, false)
+	m := NewStateMigrator(tf.Dir(), "default", actions, &MigratorOption{}, false)
 	err = m.Plan(ctx)
 	if err != nil {
 		t.Fatalf("failed to run migrator plan: %s", err)