Added workspace support

tfexec/cmd.go

@@ -156,8 +156,7 @@ 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] = ""
+	env[workspaceEnvVar] = tf.workspace
 
 	if tf.disablePluginTLS {
 		env[disablePluginTLSEnvVar] = "1"

tfexec/terraform.go

@@ -47,6 +47,7 @@ type Terraform struct {
 	disablePluginTLS   bool
 	skipProviderVerify bool
 	env                map[string]string
+	workspace          string
 
 	stdout  io.Writer
 	stderr  io.Writer
@@ -101,6 +102,38 @@ func (tf *Terraform) SetEnv(env map[string]string) error {
 	return nil
 }
 
+// SetWorkspace sets the workspace for subsequent calls to Terraform commands. It creates the workspace
+// if it does not exists. This does not affect the working directory, use WorkspaceSelect to change the
+// workspace in the working directory.
+func (tf *Terraform) SetWorkspace(ctx context.Context, workspace string, opts ...WorkspaceNewCmdOption) error {
+	lastWorkspace := tf.workspace
+
+	// set it to empty for `terraform workspace` commands to use filesystem workspace
+	tf.workspace = ""
+	ws, current, err := tf.WorkspaceList(ctx)
+
+	if err != nil {
+		tf.workspace = lastWorkspace
+		return err
+	}
+	for _, w := range ws {
+		if w == workspace {
+			tf.workspace = workspace
+			return nil
+		}
+	}
+	if err := tf.WorkspaceNew(ctx, workspace, opts...); err != nil {
+		tf.workspace = lastWorkspace
+		return err
+	}
+	if err := tf.WorkspaceSelect(ctx, current); err != nil {
+		tf.workspace = lastWorkspace
+		return err
+	}
+	tf.workspace = workspace
+	return nil
+}
+
 // SetLogger specifies a logger for tfexec to use.
 func (tf *Terraform) SetLogger(logger printfer) {
 	tf.logger = logger