From 5932540be37a585ca6e709b64c33f9b96ccc5839 Mon Sep 17 00:00:00 2001
From: Brittan DeYoung <32572259+brittandeyoung@users.noreply.github.com>
Date: Fri, 12 Aug 2022 09:14:52 -0400
Subject: [PATCH] New Terraform Workspace select flag:

---
 internal/command/workspace_command_test.go | 28 ++++++++++++++++++++++
 internal/command/workspace_select.go       | 14 +++++++++--
 2 files changed, 40 insertions(+), 2 deletions(-)

diff --git a/internal/command/workspace_command_test.go b/internal/command/workspace_command_test.go
index 35dd090d0c..7d2052e111 100644
--- a/internal/command/workspace_command_test.go
+++ b/internal/command/workspace_command_test.go
@@ -435,3 +435,31 @@ func TestWorkspace_deleteWithState(t *testing.T) {
 		t.Fatal("env 'test' still exists!")
 	}
 }
+
+func TestWorkspace_selectWithOrCreate(t *testing.T) {
+	// Create a temporary working directory that is empty
+	td := t.TempDir()
+	os.MkdirAll(td, 0755)
+	defer testChdir(t, td)()
+
+	selectCmd := &WorkspaceSelectCommand{}
+
+	current, _ := selectCmd.Workspace()
+	if current != backend.DefaultStateName {
+		t.Fatal("current workspace should be 'default'")
+	}
+
+	args := []string{"-or-create", "test"}
+	ui := new(cli.MockUi)
+	view, _ := testView(t)
+	selectCmd.Meta = Meta{Ui: ui, View: view}
+	if code := selectCmd.Run(args); code != 0 {
+		t.Fatalf("bad: %d\n\n%s", code, ui.ErrorWriter)
+	}
+
+	current, _ = selectCmd.Workspace()
+	if current != "test" {
+		t.Fatalf("current workspace should be 'test', got %q", current)
+	}
+
+}
diff --git a/internal/command/workspace_select.go b/internal/command/workspace_select.go
index e257b59d7b..bd36be210f 100644
--- a/internal/command/workspace_select.go
+++ b/internal/command/workspace_select.go
@@ -18,7 +18,9 @@ func (c *WorkspaceSelectCommand) Run(args []string) int {
 	args = c.Meta.process(args)
 	envCommandShowWarning(c.Ui, c.LegacyName)
 
+	var orCreate bool
 	cmdFlags := c.Meta.defaultFlagSet("workspace select")
+	cmdFlags.BoolVar(&orCreate, "or-create", false, "create workspace if it doesnt exist")
 	cmdFlags.Usage = func() { c.Ui.Error(c.Help()) }
 	if err := cmdFlags.Parse(args); err != nil {
 		c.Ui.Error(fmt.Sprintf("Error parsing command-line flags: %s\n", err.Error()))
@@ -96,8 +98,16 @@ func (c *WorkspaceSelectCommand) Run(args []string) int {
 	}
 
 	if !found {
-		c.Ui.Error(fmt.Sprintf(envDoesNotExist, name))
-		return 1
+		if orCreate {
+			_, err = b.StateMgr(name)
+			if err != nil {
+				c.Ui.Error(err.Error())
+				return 1
+			}
+		} else {
+			c.Ui.Error(fmt.Sprintf(envDoesNotExist, name))
+			return 1
+		}
 	}
 
 	err = c.SetWorkspace(name)