migrate command for terraform stacks (#36482)

Author: dsa0x

internal/collections/set.go

@@ -76,6 +76,15 @@ func (s Set[T]) Add(vs ...T) {
 	}
 }
 
+// AddAll inserts all the members of vs into the set.
+//
+// The behavior is the same as calling Add for each member of vs.
+func (s Set[T]) AddAll(vs Set[T]) {
+	for v := range vs.All() {
+		s.Add(v)
+	}
+}
+
 // Remove removes the given member from the set, or does nothing if no
 // equivalent value was present.
 func (s Set[T]) Remove(v T) {

internal/configs/source_bundle_parser.go

@@ -80,6 +80,11 @@ func (p *SourceBundleParser) IsConfigDir(source sourceaddrs.FinalSource) bool {
 	return (len(primaryPaths) + len(overridePaths)) > 0
 }
 
+// Bundle returns the source bundle that this parser is reading from.
+func (p *SourceBundleParser) Bundle() *sourcebundle.Bundle {
+	return p.sources
+}
+
 func (p *SourceBundleParser) dirSources(source sourceaddrs.FinalSource) (primary, override []sourceaddrs.FinalSource, diags hcl.Diagnostics) {
 	localDir, err := p.sources.LocalPathForSource(source)
 	if err != nil {

internal/rpcapi/dynrpcserver/stacks.go

@@ -8,11 +8,13 @@ import (
 	"sync"
 
 	"github.com/hashicorp/go-slug/sourcebundle"
+
 	"github.com/hashicorp/terraform/internal/depsfile"
 	"github.com/hashicorp/terraform/internal/providercache"
 	"github.com/hashicorp/terraform/internal/stacks/stackconfig"
 	"github.com/hashicorp/terraform/internal/stacks/stackplan"
 	"github.com/hashicorp/terraform/internal/stacks/stackstate"
+	"github.com/hashicorp/terraform/internal/states"
 )
 
 // handle represents an identifier shared between client and server to identify
@@ -138,6 +140,19 @@ func (t *handleTable) CloseStackPlan(hnd handle[*stackplan.Plan]) error {
 	return closeHandle(t, hnd)
 }
 
+func (t *handleTable) NewTerraformState(state *states.State) handle[*states.State] {
+	return newHandle(t, state)
+}
+
+func (t *handleTable) TerraformState(hnd handle[*states.State]) *states.State {
+	ret, _ := readHandle(t, hnd) // non-existent or invalid returns nil
+	return ret
+}
+
+func (t *handleTable) CloseTerraformState(hnd handle[*states.State]) error {
+	return closeHandle(t, hnd)
+}
+
 func (t *handleTable) NewDependencyLocks(locks *depsfile.Locks) handle[*depsfile.Locks] {
 	// NOTE: We intentionally don't track a dependency on a source bundle
 	// here for two reasons:

internal/rpcapi/handles.go

@@ -84,7 +84,7 @@ func serverHandshake(s *grpc.Server, opts *serviceOpts) func(context.Context, *s
 		// doing real work. In future the details of what we register here
 		// might vary based on the negotiated capabilities.
 		dependenciesStub.ActivateRPCServer(newDependenciesServer(handles, services))
-		stacksStub.ActivateRPCServer(newStacksServer(stopper, handles, opts))
+		stacksStub.ActivateRPCServer(newStacksServer(stopper, handles, services, opts))
 		packagesStub.ActivateRPCServer(newPackagesServer(services))
 
 		// If the client requested any extra capabililties that we're going

internal/rpcapi/plugin.go

@@ -4,13 +4,15 @@
 package rpcapi
 
 import (
+	"bytes"
 	"context"
 	"fmt"
 	"io"
 	"time"
 
 	"github.com/hashicorp/go-slug/sourceaddrs"
 	"github.com/hashicorp/go-slug/sourcebundle"
+	"github.com/hashicorp/terraform-svchost/disco"
 	"go.opentelemetry.io/otel/attribute"
 	otelCodes "go.opentelemetry.io/otel/codes"
 	"go.opentelemetry.io/otel/trace"
@@ -22,20 +24,25 @@ import (
 	"github.com/hashicorp/terraform/internal/plans"
 	"github.com/hashicorp/terraform/internal/providercache"
 	"github.com/hashicorp/terraform/internal/providers"
+	"github.com/hashicorp/terraform/internal/rpcapi/terraform1"
 	"github.com/hashicorp/terraform/internal/rpcapi/terraform1/stacks"
 	"github.com/hashicorp/terraform/internal/stacks/stackaddrs"
 	"github.com/hashicorp/terraform/internal/stacks/stackconfig"
+	"github.com/hashicorp/terraform/internal/stacks/stackmigrate"
 	"github.com/hashicorp/terraform/internal/stacks/stackplan"
 	"github.com/hashicorp/terraform/internal/stacks/stackruntime"
 	"github.com/hashicorp/terraform/internal/stacks/stackruntime/hooks"
 	"github.com/hashicorp/terraform/internal/stacks/stackstate"
+	"github.com/hashicorp/terraform/internal/states"
+	"github.com/hashicorp/terraform/internal/states/statefile"
 	"github.com/hashicorp/terraform/internal/tfdiags"
 )
 
 type stacksServer struct {
 	stacks.UnimplementedStacksServer
 
 	stopper            *stopper
+	services           *disco.Disco
 	handles            *handleTable
 	experimentsAllowed bool
 
@@ -53,11 +60,16 @@ type stacksServer struct {
 	planTimestampOverride *time.Time
 }
 
-var _ stacks.StacksServer = (*stacksServer)(nil)
+var (
+	_ stacks.StacksServer = (*stacksServer)(nil)
 
-func newStacksServer(stopper *stopper, handles *handleTable, opts *serviceOpts) *stacksServer {
+	WorkspaceNameEnvVar = "TF_WORKSPACE"
+)
+
+func newStacksServer(stopper *stopper, handles *handleTable, services *disco.Disco, opts *serviceOpts) *stacksServer {
 	return &stacksServer{
 		stopper:            stopper,
+		services:           services,
 		handles:            handles,
 		experimentsAllowed: opts.experimentsAllowed,
 	}
@@ -860,6 +872,140 @@ func (s *stacksServer) InspectExpressionResult(ctx context.Context, req *stacks.
 	return insp.InspectExpressionResult(ctx, req)
 }
 
+func (s *stacksServer) OpenTerraformState(ctx context.Context, request *stacks.OpenTerraformState_Request) (*stacks.OpenTerraformState_Response, error) {
+	switch data := request.State.(type) {
+	case *stacks.OpenTerraformState_Request_ConfigPath:
+		// Load the state from the backend.
+		// This function should return an empty state even if the diags
+		// has errors. This makes it easier for the caller, as they should
+		// close the state handle regardless of the diags.
+		loader := stackmigrate.Loader{Discovery: s.services}
+		state, diags := loader.LoadState(data.ConfigPath)
+
+		hnd := s.handles.NewTerraformState(state)
+		return &stacks.OpenTerraformState_Response{
+			StateHandle: hnd.ForProtobuf(),
+			Diagnostics: diagnosticsToProto(diags),
+		}, nil
+
+	case *stacks.OpenTerraformState_Request_Raw:
+		// load the state from the raw data
+		file, err := statefile.Read(bytes.NewReader(data.Raw))
+		if err != nil {
+			return nil, status.Errorf(codes.InvalidArgument, "invalid raw state data: %s", err)
+		}
+
+		hnd := s.handles.NewTerraformState(file.State)
+		return &stacks.OpenTerraformState_Response{
+			StateHandle: hnd.ForProtobuf(),
+		}, nil
+
+	default:
+		return nil, status.Error(codes.InvalidArgument, "invalid state source")
+	}
+}
+
+func (s *stacksServer) CloseTerraformState(ctx context.Context, request *stacks.CloseTerraformState_Request) (*stacks.CloseTerraformState_Response, error) {
+	hnd := handle[*states.State](request.StateHandle)
+	err := s.handles.CloseTerraformState(hnd)
+	if err != nil {
+		return nil, status.Error(codes.InvalidArgument, err.Error())
+	}
+	return new(stacks.CloseTerraformState_Response), nil
+}
+
+func (s *stacksServer) MigrateTerraformState(request *stacks.MigrateTerraformState_Request, server stacks.Stacks_MigrateTerraformStateServer) error {
+
+	previousStateHandle := handle[*states.State](request.StateHandle)
+	previousState := s.handles.TerraformState(previousStateHandle)
+	if previousState == nil {
+		return status.Error(codes.InvalidArgument, "the given state handle is invalid")
+	}
+
+	configHandle := handle[*stackconfig.Config](request.ConfigHandle)
+	config := s.handles.StackConfig(configHandle)
+	if config == nil {
+		return status.Error(codes.InvalidArgument, "the given config handle is invalid")
+	}
+
+	dependencyLocksHandle := handle[*depsfile.Locks](request.DependencyLocksHandle)
+	dependencyLocks := s.handles.DependencyLocks(dependencyLocksHandle)
+	if dependencyLocks == nil {
+		return status.Error(codes.InvalidArgument, "the given dependency locks handle is invalid")
+	}
+
+	var providerFactories map[addrs.Provider]providers.Factory
+	if s.providerCacheOverride != nil {
+		// This is only used in tests to side load providers without needing a
+		// real provider cache.
+		providerFactories = s.providerCacheOverride
+	} else {
+		providerCacheHandle := handle[*providercache.Dir](request.ProviderCacheHandle)
+		providerCache := s.handles.ProviderPluginCache(providerCacheHandle)
+		if providerCache == nil {
+			return status.Error(codes.InvalidArgument, "the given provider cache handle is invalid")
+		}
+
+		var err error
+		providerFactories, err = providerFactoriesForLocks(dependencyLocks, providerCache)
+		if err != nil {
+			return status.Errorf(codes.InvalidArgument, "provider dependencies are inconsistent: %s", err)
+		}
+	}
+
+	migrate := &stackmigrate.Migration{
+		Providers:     providerFactories,
+		PreviousState: previousState,
+		Config:        config,
+	}
+
+	emit := func(change stackstate.AppliedChange) {
+		proto, err := change.AppliedChangeProto()
+		if err != nil {
+			server.Send(&stacks.MigrateTerraformState_Event{
+				Result: &stacks.MigrateTerraformState_Event_Diagnostic{
+					Diagnostic: &terraform1.Diagnostic{
+						Severity: terraform1.Diagnostic_ERROR,
+						Summary:  "Failed to serialize change",
+						Detail:   fmt.Sprintf("Failed to serialize state change for recording in the migration plan: %s", err),
+					},
+				},
+			})
+			return
+		}
+
+		server.Send(&stacks.MigrateTerraformState_Event{
+			Result: &stacks.MigrateTerraformState_Event_AppliedChange{
+				AppliedChange: proto,
+			},
+		})
+	}
+
+	emitDiag := func(diagnostic tfdiags.Diagnostic) {
+		server.Send(&stacks.MigrateTerraformState_Event{
+			Result: &stacks.MigrateTerraformState_Event_Diagnostic{
+				Diagnostic: diagnosticToProto(diagnostic),
+			},
+		})
+	}
+
+	mapping := request.GetMapping()
+	if mapping == nil {
+		return status.Error(codes.InvalidArgument, "missing migration mapping")
+	}
+	switch mapping := mapping.(type) {
+	case *stacks.MigrateTerraformState_Request_Simple:
+		migrate.Migrate(
+			mapping.Simple.ResourceAddressMap,
+			mapping.Simple.ModuleAddressMap,
+			emit, emitDiag)
+	default:
+		return status.Error(codes.InvalidArgument, "unsupported migration mapping")
+	}
+
+	return nil
+}
+
 func stackPlanHooks(evts *syncPlanStackChangesServer, mainStackSource sourceaddrs.FinalSource) *stackruntime.Hooks {
 	return stackChangeHooks(
 		func(scp *stacks.StackChangeProgress) error {