feat: Added back TF_WORKSPACE support

through `WorkSpaceOption`. usage: ```hcl var options = []tfexec.ApplyOption{ tfexec.WorkSpace("some_workspace"), } ```
hashicorp · May 11, 2022 · 4384294 · 4384294
1 parent d286988
commit 4384294
Show file tree

Hide file tree

Showing 7 changed files with 100 additions and 8 deletions.
diff --git a/tfexec/refresh.go b/tfexec/refresh.go
@@ -17,6 +17,9 @@ type refreshConfig struct {
 	targets      []string
 	vars         []string
 	varFiles     []string
+
+	// WorkSpace: when leveraging TF_WORKSPACE env variable
+	workSpace string
 }
 
 var defaultRefreshOptions = refreshConfig{
@@ -69,6 +72,10 @@ func (opt *VarFileOption) configureRefresh(conf *refreshConfig) {
 	conf.varFiles = append(conf.varFiles, opt.path)
 }
 
+func (opt *WorkSpaceOption) configureRefresh(conf *refreshConfig) {
+	conf.workSpace = opt.workSpace
+}
+
 // Refresh represents the terraform refresh subcommand.
 func (tf *Terraform) Refresh(ctx context.Context, opts ...RefreshCmdOption) error {
 	cmd, err := tf.refreshCmd(ctx, opts...)
@@ -133,5 +140,10 @@ func (tf *Terraform) refreshCmd(ctx context.Context, opts ...RefreshCmdOption) (
 		mergeEnv[reattachEnvVar] = reattachStr
 	}
 
+	// Check if TF_WORKSPACE is specified.
+	if c.workSpace != "" {
+		mergeEnv[workspaceEnvVar] = c.workSpace
+	}
+
 	return tf.buildTerraformCmd(ctx, mergeEnv, args...), nil
 }
diff --git a/tfexec/refresh_test.go b/tfexec/refresh_test.go
@@ -34,7 +34,20 @@ func TestRefreshCmd(t *testing.T) {
 	})
 
 	t.Run("override all defaults", func(t *testing.T) {
-		refreshCmd, err := tf.refreshCmd(context.Background(), Backup("testbackup"), LockTimeout("200s"), State("teststate"), StateOut("teststateout"), VarFile("testvarfile"), Lock(false), Target("target1"), Target("target2"), Var("var1=foo"), Var("var2=bar"), Dir("refreshdir"))
+		refreshCmd, err := tf.refreshCmd(context.Background(),
+			Backup("testbackup"),
+			LockTimeout("200s"),
+			State("teststate"),
+			StateOut("teststateout"),
+			VarFile("testvarfile"),
+			Lock(false),
+			Target("target1"),
+			Target("target2"),
+			Var("var1=foo"),
+			Var("var2=bar"),
+			Dir("refreshdir"),
+			WorkSpace("some_workspace"))
+
 		if err != nil {
 			t.Fatal(err)
 		}
@@ -54,6 +67,6 @@ func TestRefreshCmd(t *testing.T) {
 			"-var", "var1=foo",
 			"-var", "var2=bar",
 			"refreshdir",
-		}, nil, refreshCmd)
+		}, map[string]string{"TF_WORKSPACE": "some_workspace"}, refreshCmd)
 	})
 }
diff --git a/tfexec/show.go b/tfexec/show.go
@@ -11,6 +11,9 @@ import (
 
 type showConfig struct {
 	reattachInfo ReattachInfo
+
+	// WorkSpace: when leveraging TF_WORKSPACE env variable
+	workSpace string
 }
 
 var defaultShowOptions = showConfig{}
@@ -23,6 +26,10 @@ func (opt *ReattachOption) configureShow(conf *showConfig) {
 	conf.reattachInfo = opt.info
 }
 
+func (opt *WorkSpaceOption) configureShow(conf *showConfig) {
+	conf.workSpace = opt.workSpace
+}
+
 // Show reads the default state path and outputs the state.
 // To read a state or plan file, ShowState or ShowPlan must be used instead.
 func (tf *Terraform) Show(ctx context.Context, opts ...ShowOption) (*tfjson.State, error) {
@@ -46,6 +53,11 @@ func (tf *Terraform) Show(ctx context.Context, opts ...ShowOption) (*tfjson.Stat
 		mergeEnv[reattachEnvVar] = reattachStr
 	}
 
+	// Check if TF_WORKSPACE is specified.
+	if c.workSpace != "" {
+		mergeEnv[workspaceEnvVar] = c.workSpace
+	}
+
 	showCmd := tf.showCmd(ctx, true, mergeEnv)
 
 	var ret tfjson.State
@@ -89,6 +101,11 @@ func (tf *Terraform) ShowStateFile(ctx context.Context, statePath string, opts .
 		mergeEnv[reattachEnvVar] = reattachStr
 	}
 
+	// Check if TF_WORKSPACE is specified.
+	if c.workSpace != "" {
+		mergeEnv[workspaceEnvVar] = c.workSpace
+	}
+
 	showCmd := tf.showCmd(ctx, true, mergeEnv, statePath)
 
 	var ret tfjson.State
@@ -132,6 +149,11 @@ func (tf *Terraform) ShowPlanFile(ctx context.Context, planPath string, opts ...
 		mergeEnv[reattachEnvVar] = reattachStr
 	}
 
+	// Check if TF_WORKSPACE is specified.
+	if c.workSpace != "" {
+		mergeEnv[workspaceEnvVar] = c.workSpace
+	}
+
 	showCmd := tf.showCmd(ctx, true, mergeEnv, planPath)
 
 	var ret tfjson.Plan
@@ -171,6 +193,11 @@ func (tf *Terraform) ShowPlanFileRaw(ctx context.Context, planPath string, opts
 		mergeEnv[reattachEnvVar] = reattachStr
 	}
 
+	// Check if TF_WORKSPACE is specified.
+	if c.workSpace != "" {
+		mergeEnv[workspaceEnvVar] = c.workSpace
+	}
+
 	showCmd := tf.showCmd(ctx, false, mergeEnv, planPath)
 
 	var outBuf strings.Builder

diff --git a/tfexec/state_mv.go b/tfexec/state_mv.go
@@ -14,6 +14,9 @@ type stateMvConfig struct {
 	lockTimeout string
 	state       string
 	stateOut    string
+
+	// WorkSpace: when leveraging TF_WORKSPACE env variable
+	workSpace string
 }
 
 var defaultStateMvOptions = stateMvConfig{
@@ -54,6 +57,10 @@ func (opt *StateOutOption) configureStateMv(conf *stateMvConfig) {
 	conf.stateOut = opt.path
 }
 
+func (opt *WorkSpaceOption) configureStateMv(conf *stateMvConfig) {
+	conf.workSpace = opt.workSpace
+}
+
 // StateMv represents the terraform state mv subcommand.
 func (tf *Terraform) StateMv(ctx context.Context, source string, destination string, opts ...StateMvCmdOption) error {
 	cmd, err := tf.stateMvCmd(ctx, source, destination, opts...)
@@ -101,5 +108,12 @@ func (tf *Terraform) stateMvCmd(ctx context.Context, source string, destination
 	args = append(args, source)
 	args = append(args, destination)
 
-	return tf.buildTerraformCmd(ctx, nil, args...), nil
+	mergeEnv := map[string]string{}
+
+	// Check if TF_WORKSPACE is specified.
+	if c.workSpace != "" {
+		mergeEnv[workspaceEnvVar] = c.workSpace
+	}
+
+	return tf.buildTerraformCmd(ctx, mergeEnv, args...), nil
 }
diff --git a/tfexec/state_mv_test.go b/tfexec/state_mv_test.go
@@ -36,7 +36,13 @@ func TestStateMvCmd(t *testing.T) {
 	})
 
 	t.Run("override all defaults", func(t *testing.T) {
-		stateMvCmd, err := tf.stateMvCmd(context.Background(), "testsrc", "testdest", Backup("testbackup"), BackupOut("testbackupout"), LockTimeout("200s"), State("teststate"), StateOut("teststateout"), Lock(false))
+		stateMvCmd, err := tf.stateMvCmd(context.Background(), "testsrc", "testdest", Backup("testbackup"),
+			BackupOut("testbackupout"),
+			LockTimeout("200s"),
+			State("teststate"),
+			StateOut("teststateout"),
+			Lock(false),
+			WorkSpace("some_workspace"))
 		if err != nil {
 			t.Fatal(err)
 		}
@@ -53,6 +59,6 @@ func TestStateMvCmd(t *testing.T) {
 			"-lock=false",
 			"testsrc",
 			"testdest",
-		}, nil, stateMvCmd)
+		}, map[string]string{"TF_WORKSPACE": "some_workspace"}, stateMvCmd)
 	})
 }
diff --git a/tfexec/state_rm.go b/tfexec/state_rm.go
@@ -14,6 +14,9 @@ type stateRmConfig struct {
 	lockTimeout string
 	state       string
 	stateOut    string
+
+	// WorkSpace: when leveraging TF_WORKSPACE env variable
+	workSpace string
 }
 
 var defaultStateRmOptions = stateRmConfig{
@@ -54,6 +57,10 @@ func (opt *StateOutOption) configureStateRm(conf *stateRmConfig) {
 	conf.stateOut = opt.path
 }
 
+func (opt *WorkSpaceOption) configureStateRm(conf *stateRmConfig) {
+	conf.workSpace = opt.workSpace
+}
+
 // StateRm represents the terraform state rm subcommand.
 func (tf *Terraform) StateRm(ctx context.Context, address string, opts ...StateRmCmdOption) error {
 	cmd, err := tf.stateRmCmd(ctx, address, opts...)
@@ -100,5 +107,12 @@ func (tf *Terraform) stateRmCmd(ctx context.Context, address string, opts ...Sta
 	// positional arguments
 	args = append(args, address)
 
-	return tf.buildTerraformCmd(ctx, nil, args...), nil
+	mergeEnv := map[string]string{}
+
+	// Check if TF_WORKSPACE is specified.
+	if c.workSpace != "" {
+		mergeEnv[workspaceEnvVar] = c.workSpace
+	}
+
+	return tf.buildTerraformCmd(ctx, mergeEnv, args...), nil
 }
diff --git a/tfexec/state_rm_test.go b/tfexec/state_rm_test.go
@@ -35,7 +35,13 @@ func TestStateRmCmd(t *testing.T) {
 	})
 
 	t.Run("override all defaults", func(t *testing.T) {
-		stateRmCmd, err := tf.stateRmCmd(context.Background(), "testAddress", Backup("testbackup"), BackupOut("testbackupout"), LockTimeout("200s"), State("teststate"), StateOut("teststateout"), Lock(false))
+		stateRmCmd, err := tf.stateRmCmd(context.Background(), "testAddress", Backup("testbackup"),
+			BackupOut("testbackupout"),
+			LockTimeout("200s"),
+			State("teststate"),
+			StateOut("teststateout"),
+			Lock(false),
+			WorkSpace("some_workspace"))
 		if err != nil {
 			t.Fatal(err)
 		}
@@ -51,6 +57,6 @@ func TestStateRmCmd(t *testing.T) {
 			"-state-out=teststateout",
 			"-lock=false",
 			"testAddress",
-		}, nil, stateRmCmd)
+		}, map[string]string{"TF_WORKSPACE": "some_workspace"}, stateRmCmd)
 	})
 }