feat: Added back TF_WORKSPACE support

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

Hide file tree

Showing 11 changed files with 80 additions and 11 deletions.
diff --git a/tfexec/apply.go b/tfexec/apply.go
@@ -25,6 +25,9 @@ type applyConfig struct {
 	// Vars: each var must be supplied as a single string, e.g. 'foo=bar'
 	vars     []string
 	varFiles []string
+
+	// WorkSpace: when leveraging TF_WORKSPACE env variable
+	workSpace string
 }
 
 var defaultApplyOptions = applyConfig{
@@ -90,6 +93,10 @@ func (opt *ReattachOption) configureApply(conf *applyConfig) {
 	conf.reattachInfo = opt.info
 }
 
+func (opt *WorkSpaceOption) configureApply(conf *applyConfig) {
+	conf.workSpace = opt.workSpace
+}
+
 // Apply represents the terraform apply subcommand.
 func (tf *Terraform) Apply(ctx context.Context, opts ...ApplyOption) error {
 	cmd, err := tf.applyCmd(ctx, opts...)
@@ -165,5 +172,10 @@ func (tf *Terraform) applyCmd(ctx context.Context, opts ...ApplyOption) (*exec.C
 		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/apply_test.go b/tfexec/apply_test.go
@@ -36,6 +36,7 @@ func TestApplyCmd(t *testing.T) {
 			Var("var1=foo"),
 			Var("var2=bar"),
 			DirOrPlan("testfile"),
+			WorkSpace("some_workspace"),
 		)
 		if err != nil {
 			t.Fatal(err)
@@ -62,6 +63,7 @@ func TestApplyCmd(t *testing.T) {
 			"-var", "var1=foo",
 			"-var", "var2=bar",
 			"testfile",
-		}, nil, applyCmd)
+		}, map[string]string{"TF_WORKSPACE": "some_workspace"}, applyCmd)
 	})
+
 }
diff --git a/tfexec/cmd.go b/tfexec/cmd.go
@@ -39,7 +39,6 @@ var prohibitedEnvVars = []string{
 	logEnvVar,
 	reattachEnvVar,
 	appendUserAgentEnvVar,
-	workspaceEnvVar,
 	disablePluginTLSEnvVar,
 	skipProviderVerifyEnvVar,
 }
@@ -156,9 +155,6 @@ func (tf *Terraform) buildEnv(mergeEnv map[string]string) []string {
 	// constant automation override env vars
 	env[automationEnvVar] = "1"
 
-	// force usage of workspace methods for switching
-	env[workspaceEnvVar] = ""
-
 	if tf.disablePluginTLS {
 		env[disablePluginTLSEnvVar] = "1"
 	}

diff --git a/tfexec/cmd_test.go b/tfexec/cmd_test.go
@@ -41,7 +41,6 @@ func defaultEnv() []string {
 		"TF_IN_AUTOMATION=1",
 		"TF_LOG_PATH=",
 		"TF_LOG=",
-		"TF_WORKSPACE=",
 	}
 }
 

diff --git a/tfexec/destroy.go b/tfexec/destroy.go
@@ -24,6 +24,9 @@ type destroyConfig struct {
 	// Vars: each var must be supplied as a single string, e.g. 'foo=bar'
 	vars     []string
 	varFiles []string
+
+	// WorkSpace: when leveraging TF_WORKSPACE env variable
+	workSpace string
 }
 
 var defaultDestroyOptions = destroyConfig{
@@ -86,6 +89,10 @@ func (opt *ReattachOption) configureDestroy(conf *destroyConfig) {
 	conf.reattachInfo = opt.info
 }
 
+func (opt *WorkSpaceOption) configureDestroy(conf *destroyConfig) {
+	conf.workSpace = opt.workSpace
+}
+
 // Destroy represents the terraform destroy subcommand.
 func (tf *Terraform) Destroy(ctx context.Context, opts ...DestroyOption) error {
 	cmd, err := tf.destroyCmd(ctx, opts...)
@@ -152,5 +159,10 @@ func (tf *Terraform) destroyCmd(ctx context.Context, opts ...DestroyOption) (*ex
 		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/destroy_test.go b/tfexec/destroy_test.go
@@ -37,7 +37,20 @@ func TestDestroyCmd(t *testing.T) {
 	})
 
 	t.Run("override all defaults", func(t *testing.T) {
-		destroyCmd, err := tf.destroyCmd(context.Background(), Backup("testbackup"), LockTimeout("200s"), State("teststate"), StateOut("teststateout"), VarFile("testvarfile"), Lock(false), Parallelism(99), Refresh(false), Target("target1"), Target("target2"), Var("var1=foo"), Var("var2=bar"), Dir("destroydir"))
+		destroyCmd, err := tf.destroyCmd(context.Background(),
+			Backup("testbackup"),
+			LockTimeout("200s"),
+			State("teststate"),
+			StateOut("teststateout"),
+			VarFile("testvarfile"),
+			Lock(false),
+			Parallelism(99),
+			Refresh(false),
+			Target("target1"),
+			Target("target2"),
+			Var("var1=foo"),
+			Var("var2=bar"), Dir("destroydir"),
+			WorkSpace("some_workspace"))
 		if err != nil {
 			t.Fatal(err)
 		}
@@ -60,6 +73,6 @@ func TestDestroyCmd(t *testing.T) {
 			"-var", "var1=foo",
 			"-var", "var2=bar",
 			"destroydir",
-		}, nil, destroyCmd)
+		}, map[string]string{"TF_WORKSPACE": "some_workspace"}, destroyCmd)
 	})
 }
diff --git a/tfexec/import.go b/tfexec/import.go
@@ -19,6 +19,9 @@ type importConfig struct {
 	stateOut           string
 	vars               []string
 	varFiles           []string
+
+	// WorkSpace: when leveraging TF_WORKSPACE env variable
+	workSpace string
 }
 
 var defaultImportOptions = importConfig{
@@ -72,6 +75,10 @@ func (opt *VarFileOption) configureImport(conf *importConfig) {
 	conf.varFiles = append(conf.varFiles, opt.path)
 }
 
+func (opt *WorkSpaceOption) configureImport(conf *importConfig) {
+	conf.workSpace = opt.workSpace
+}
+
 // Import represents the terraform import subcommand.
 func (tf *Terraform) Import(ctx context.Context, address, id string, opts ...ImportOption) error {
 	cmd, err := tf.importCmd(ctx, address, id, opts...)
@@ -137,5 +144,10 @@ func (tf *Terraform) importCmd(ctx context.Context, address, id string, opts ...
 		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/import_test.go b/tfexec/import_test.go
@@ -46,6 +46,7 @@ func TestImportCmd(t *testing.T) {
 			Var("var1=foo"),
 			Var("var2=bar"),
 			AllowMissingConfig(true),
+			WorkSpace("some_workspace"),
 		)
 		if err != nil {
 			t.Fatal(err)
@@ -66,6 +67,6 @@ func TestImportCmd(t *testing.T) {
 			"-var", "var2=bar",
 			"my-addr2",
 			"my-id2",
-		}, nil, importCmd)
+		}, map[string]string{"TF_WORKSPACE": "some_workspace"}, importCmd)
 	})
 }
diff --git a/tfexec/options.go b/tfexec/options.go
@@ -409,3 +409,12 @@ type VerifyPluginsOption struct {
 func VerifyPlugins(verifyPlugins bool) *VerifyPluginsOption {
 	return &VerifyPluginsOption{verifyPlugins}
 }
+
+// WorkSpaceOption is used to leverage the TF_WORKSPACE environment variable
+type WorkSpaceOption struct {
+	workSpace string
+}
+
+func WorkSpace(workSpace string) *WorkSpaceOption {
+	return &WorkSpaceOption{workSpace}
+}
diff --git a/tfexec/plan.go b/tfexec/plan.go
@@ -21,6 +21,9 @@ type planConfig struct {
 	targets      []string
 	vars         []string
 	varFiles     []string
+
+	// WorkSpace: when leveraging TF_WORKSPACE env variable
+	workSpace string
 }
 
 var defaultPlanOptions = planConfig{
@@ -88,6 +91,10 @@ func (opt *DestroyFlagOption) configurePlan(conf *planConfig) {
 	conf.destroy = opt.destroy
 }
 
+func (opt *WorkSpaceOption) configurePlan(conf *planConfig) {
+	conf.workSpace = opt.workSpace
+}
+
 // Plan executes `terraform plan` with the specified options and waits for it
 // to complete.
 //
@@ -176,5 +183,10 @@ func (tf *Terraform) planCmd(ctx context.Context, opts ...PlanOption) (*exec.Cmd
 		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/plan_test.go b/tfexec/plan_test.go
@@ -52,7 +52,8 @@ func TestPlanCmd(t *testing.T) {
 			Var("android=paranoid"),
 			Var("brain_size=planet"),
 			VarFile("trillian"),
-			Dir("earth"))
+			Dir("earth"),
+			WorkSpace("some_workspace"))
 		if err != nil {
 			t.Fatal(err)
 		}
@@ -77,6 +78,6 @@ func TestPlanCmd(t *testing.T) {
 			"-var", "android=paranoid",
 			"-var", "brain_size=planet",
 			"earth",
-		}, nil, planCmd)
+		}, map[string]string{"TF_WORKSPACE": "some_workspace"}, planCmd)
 	})
 }