From c3bf5d45128954c555ae2d30e1d6300016f77ee9 Mon Sep 17 00:00:00 2001
From: Chris Arcand <chris@chrisarcand.com>
Date: Wed, 3 Nov 2021 10:59:09 -0500
Subject: [PATCH] cloud: Automatically select renamed current workspace

After migrating to TFC with renamed workspaces, automatically select
what was the previous current workspace on behalf of the user. We don't
need to make the user reselect.
---
 internal/command/meta_backend_migrate.go | 32 +++++++++++++++++++++---
 1 file changed, 29 insertions(+), 3 deletions(-)

diff --git a/internal/command/meta_backend_migrate.go b/internal/command/meta_backend_migrate.go
index e708cb3802..3c022a88a5 100644
--- a/internal/command/meta_backend_migrate.go
+++ b/internal/command/meta_backend_migrate.go
@@ -613,6 +613,12 @@ func (m *Meta) backendMigrateTFC(opts *backendMigrateOpts) error {
 func (m *Meta) backendMigrateState_S_TFC(opts *backendMigrateOpts, sourceWorkspaces []string) error {
 	log.Print("[TRACE] backendMigrateState: migrating all named workspaces")
 
+	currentWorkspace, err := m.Workspace()
+	if err != nil {
+		return err
+	}
+	newCurrentWorkspace := ""
+
 	// This map is used later when doing the migration per source/destination.
 	// If a source has 'default', then we ask what the new name should be.
 	// And further down when we actually run state migration for each
@@ -653,15 +659,35 @@ func (m *Meta) backendMigrateState_S_TFC(opts *backendMigrateOpts, sourceWorkspa
 			return fmt.Errorf(strings.TrimSpace(
 				errMigrateMulti), name, opts.SourceType, opts.DestinationType, err)
 		}
+
+		if currentWorkspace == opts.sourceWorkspace {
+			newCurrentWorkspace = opts.destinationWorkspace
+		}
 	}
 
-	// After migrating multiple workspaces, we want to ensure that a workspace is
-	// set or we prompt the user to set a workspace.
-	err = m.selectWorkspace(opts.Destination)
+	// After migrating multiple workspaces, we need to reselect the current workspace as it may
+	// have been renamed. Query the backend first to be sure it now exists, and if it does,
+	// select it.
+	workspaces, err := opts.Destination.Workspaces()
 	if err != nil {
 		return err
 	}
 
+	for _, name := range workspaces {
+		if name == newCurrentWorkspace {
+			if err = m.SetWorkspace(name); err != nil {
+				return err
+			}
+			return nil
+		}
+	}
+
+	// If we couldn't select the workspace automatically from the backend (maybe it was empty
+	// and wasn't migrated, for instance), ask the user to select one.
+	if err = m.selectWorkspace(opts.Destination); err != nil {
+		return err
+	}
+
 	return nil
 }