Add support for snapshot load

cmd/root.go

@@ -78,7 +78,7 @@ func NewRootCmd(cfg *env.Env, tel *telemetry.Client, logger log.Logger) *cobra.C
 		newUpdateCmd(cfg),
 		newDocsCmd(),
 		newAWSCmd(cfg),
-		newSnapshotCmd(cfg),
+		newSnapshotCmd(cfg, tel, logger),
 		newResetCmd(cfg),
 	)
 

cmd/snapshot.go

@@ -1,30 +1,135 @@
 package cmd
 
 import (
+	"context"
 	"fmt"
 	"os"
 	"time"
 
 	"github.com/localstack/lstk/internal/config"
+	"github.com/localstack/lstk/internal/container"
 	"github.com/localstack/lstk/internal/emulator/aws"
 	"github.com/localstack/lstk/internal/endpoint"
 	"github.com/localstack/lstk/internal/env"
+	"github.com/localstack/lstk/internal/log"
 	"github.com/localstack/lstk/internal/output"
 	"github.com/localstack/lstk/internal/runtime"
 	"github.com/localstack/lstk/internal/snapshot"
+	"github.com/localstack/lstk/internal/telemetry"
 	"github.com/localstack/lstk/internal/ui"
 	"github.com/spf13/cobra"
 )
 
-func newSnapshotCmd(cfg *env.Env) *cobra.Command {
+func newSnapshotCmd(cfg *env.Env, tel *telemetry.Client, logger log.Logger) *cobra.Command {
 	cmd := &cobra.Command{
 		Use:   "snapshot",
 		Short: "Manage emulator snapshots",
 	}
 	cmd.AddCommand(newSnapshotSaveCmd(cfg))
+	cmd.AddCommand(newSnapshotLoadCmd(cfg, tel, logger))
 	return cmd
 }
 
+func buildStarter(cfg *env.Env, rt runtime.Runtime, appConfig *config.Config, logger log.Logger, tel *telemetry.Client) snapshot.Starter {
+	return func(ctx context.Context, sink output.Sink) error {
+		opts := buildStartOptions(cfg, appConfig, logger, tel, false)
+		return container.Start(ctx, rt, sink, opts, false)
+	}
+}
+
+func newSnapshotLoadCmd(cfg *env.Env, tel *telemetry.Client, logger log.Logger) *cobra.Command {
+	cmd := &cobra.Command{
+		Use:   "load REF",
+		Short: "Load a snapshot into the running emulator",
+		Long: `Load a snapshot into the running emulator, starting it first if needed.
+
+REF identifies the snapshot to load:
+
+  lstk snapshot load my-baseline           # loads ./my-baseline or ./my-baseline.zip
+  lstk snapshot load ./checkpoint.zip      # loads from explicit path
+  lstk snapshot load pod:my-baseline       # loads from LocalStack Cloud
+
+Merge strategies control how snapshot state is combined with running state:
+
+  --merge=account-region-merge  (default) snapshot wins on (service, account, region) overlap
+  --merge=overwrite             wipe running state, then load
+  --merge=service-merge         snapshot wins per-resource; non-overlapping resources combined`,
+		Args:    cobra.ExactArgs(1),
+		PreRunE: initConfig(nil),
+		RunE: func(cmd *cobra.Command, args []string) error {
+			dryRun, err := cmd.Flags().GetBool("dry-run")
+			if err != nil {
+				return err
+			}
+			if dryRun {
+				return fmt.Errorf("--dry-run is not yet implemented")
+			}
+
+			strategy, err := cmd.Flags().GetString("merge")
+			if err != nil {
+				return err
+			}
+
+			src, err := snapshot.ParseSource(args[0])
+			if err != nil {
+				return err
+			}
+
+			if err := snapshot.ValidateMergeStrategy(strategy); err != nil {
+				return err
+			}
+
+			rt, client, host, containers, appConfig, err := resolveSnapshotDeps(cmd.Context(), cfg)
+			if err != nil {
+				return err
+			}
+
+			starter := buildStarter(cfg, rt, appConfig, logger, tel)
+
+			if isInteractiveMode(cfg) {
+				return ui.RunSnapshotLoad(cmd.Context(), rt, containers, client, host, src, cfg.AuthToken, strategy, starter)
+			}
+			sink := output.NewPlainSink(os.Stdout)
+			switch src.Kind {
+			case snapshot.KindPod:
+				return snapshot.LoadPod(cmd.Context(), rt, containers, client, host, src.Value, cfg.AuthToken, strategy, starter, sink)
+			default:
+				return snapshot.LoadLocal(cmd.Context(), rt, containers, client, host, src.Value, strategy, starter, sink)
+			}
+		},
+	}
+	cmd.Flags().String("merge", snapshot.MergeStrategyAccountRegion, "Merge strategy: overwrite, account-region-merge, service-merge")
+	cmd.Flags().Bool("dry-run", false, "Preview changes without applying (not yet implemented)")
+	return cmd
+}
+
+func resolveSnapshotDeps(ctx context.Context, cfg *env.Env) (rt runtime.Runtime, client *aws.Client, host string, containers []config.ContainerConfig, appConfig *config.Config, err error) {
+	appConfig, err = config.Get()
+	if err != nil {
+		return nil, nil, "", nil, nil, fmt.Errorf("failed to get config: %w", err)
+	}
+
+	var awsContainer config.ContainerConfig
+	var found bool
+	for _, c := range appConfig.Containers {
+		if c.Type == config.EmulatorAWS {
+			awsContainer = c
+			found = true
+			break
+		}
+	}
+	if !found {
+		return nil, nil, "", nil, nil, fmt.Errorf("snapshot is only supported for the AWS emulator")
+	}
+
+	rt, err = runtime.NewDockerRuntime(cfg.DockerHost)
+	if err != nil {
+		return nil, nil, "", nil, nil, err
+	}
+	host, _ = endpoint.ResolveHost(ctx, awsContainer.Port, cfg.LocalStackHost)
+	return rt, aws.NewClient(), host, []config.ContainerConfig{awsContainer}, appConfig, nil
+}
+
 func newSnapshotSaveCmd(cfg *env.Env) *cobra.Command {
 	return &cobra.Command{
 		Use:   "save [destination]",
@@ -53,31 +158,10 @@ To save to a remote pod on the LocalStack platform, use the pod: prefix:
 				return err
 			}
 
-			appConfig, err := config.Get()
-			if err != nil {
-				return fmt.Errorf("failed to get config: %w", err)
-			}
-
-			var awsContainer config.ContainerConfig
-			var found bool
-			for _, c := range appConfig.Containers {
-				if c.Type == config.EmulatorAWS {
-					awsContainer = c
-					found = true
-					break
-				}
-			}
-			if !found {
-				return fmt.Errorf("snapshot is only supported for the AWS emulator")
-			}
-
-			rt, err := runtime.NewDockerRuntime(cfg.DockerHost)
+			rt, client, host, containers, _, err := resolveSnapshotDeps(cmd.Context(), cfg)
 			if err != nil {
 				return err
 			}
-			host, _ := endpoint.ResolveHost(cmd.Context(), awsContainer.Port, cfg.LocalStackHost)
-			client := aws.NewClient()
-			containers := []config.ContainerConfig{awsContainer}
 
 			if isInteractiveMode(cfg) {
 				return ui.RunSnapshotSave(cmd.Context(), rt, containers, client, host, dest, cfg.AuthToken)

internal/emulator/aws/client.go

@@ -178,6 +178,121 @@ func (c *Client) ExportState(ctx context.Context, host string, dst io.Writer) er
 	return nil
 }
 
+func (c *Client) ImportState(ctx context.Context, host string, src io.Reader, strategy string) error {
+	url := fmt.Sprintf("http://%s/_localstack/pods", host)
+	if strategy != "" {
+		url += "?merge=" + strategy
+	}
+	req, err := http.NewRequestWithContext(ctx, http.MethodPost, url, src)
+	if err != nil {
+		return fmt.Errorf("create request: %w", err)
+	}
+	req.Header.Set("Content-Type", "application/octet-stream")
+
+	resp, err := c.http.Do(req)
+	if err != nil {
+		return fmt.Errorf("connect to LocalStack: %w", err)
+	}
+	defer func() { _ = resp.Body.Close() }()
+
+	if resp.StatusCode == http.StatusUnprocessableEntity {
+		body, _ := io.ReadAll(resp.Body)
+		return fmt.Errorf("snapshot is incompatible with the running LocalStack version: %s", strings.TrimSpace(string(body)))
+	}
+	if resp.StatusCode != http.StatusCreated && resp.StatusCode != http.StatusOK {
+		body, _ := io.ReadAll(resp.Body)
+		return fmt.Errorf("LocalStack returned status %d: %s", resp.StatusCode, strings.TrimSpace(string(body)))
+	}
+
+	scanner := bufio.NewScanner(resp.Body)
+	buf := make([]byte, 1024*1024)
+	scanner.Buffer(buf, 1024*1024)
+	for scanner.Scan() {
+		line := strings.TrimSpace(scanner.Text())
+		if line == "" {
+			continue
+		}
+		var event struct {
+			Service string `json:"service"`
+			Status  string `json:"status"`
+			Message string `json:"message"`
+		}
+		if err := json.Unmarshal([]byte(line), &event); err != nil {
+			continue
+		}
+		if event.Status == "error" && event.Message != "" {
+			return fmt.Errorf("load failed for service %s: %s", event.Service, event.Message)
+		}
+	}
+	return scanner.Err()
+}
+
+func (c *Client) LoadPodSnapshot(ctx context.Context, host, podName, authToken, strategy string) ([]string, error) {
+	url := fmt.Sprintf("http://%s/_localstack/pods/%s", host, podName)
+	if strategy != "" {
+		url += "?merge=" + strategy
+	}
+	req, err := http.NewRequestWithContext(ctx, http.MethodPut, url, bytes.NewReader([]byte("{}")))
+	if err != nil {
+		return nil, fmt.Errorf("create request: %w", err)
+	}
+	req.Header.Set("Content-Type", "application/json")
+	req.Header.Set("Authorization", "Basic "+base64.StdEncoding.EncodeToString([]byte(":"+authToken)))
+
+	resp, err := c.http.Do(req)
+	if err != nil {
+		return nil, fmt.Errorf("connect to LocalStack: %w", err)
+	}
+	defer func() { _ = resp.Body.Close() }()
+
+	if resp.StatusCode == http.StatusUnprocessableEntity {
+		body, _ := io.ReadAll(resp.Body)
+		return nil, fmt.Errorf("snapshot is incompatible with the running LocalStack version: %s", strings.TrimSpace(string(body)))
+	}
+	if resp.StatusCode != http.StatusOK {
+		body, _ := io.ReadAll(resp.Body)
+		return nil, fmt.Errorf("pod load failed (HTTP %d): %s", resp.StatusCode, strings.TrimSpace(string(body)))
+	}
+
+	var services []string
+	scanner := bufio.NewScanner(resp.Body)
+	buf := make([]byte, 1024*1024)
+	scanner.Buffer(buf, 1024*1024)
+	for scanner.Scan() {
+		line := strings.TrimSpace(scanner.Text())
+		if line == "" {
+			continue
+		}
+		var event struct {
+			Event   string `json:"event"`
+			Service string `json:"service"`
+			Status  string `json:"status"`
+			Message string `json:"message"`
+		}
+		if err := json.Unmarshal([]byte(line), &event); err != nil {
+			continue
+		}
+		switch event.Event {
+		case "service":
+			switch event.Status {
+			case "ok":
+				services = append(services, event.Service)
+			case "error":
+				return nil, fmt.Errorf("load failed for service %s: %s", event.Service, event.Message)
+			}
+		case "completion":
+			if event.Status != "ok" {
+				return nil, fmt.Errorf("pod load failed: %s", event.Message)
+			}
+			return services, nil
+		}
+	}
+	if err := scanner.Err(); err != nil {
+		return nil, fmt.Errorf("reading response: %w", err)
+	}
+	return services, nil
+}
+
 func (c *Client) SavePodSnapshot(ctx context.Context, host, podName, authToken string) (snapshot.PodSaveResult, error) {
 	url := fmt.Sprintf("http://%s/_localstack/pods/%s", host, podName)
 	req, err := http.NewRequestWithContext(ctx, http.MethodPost, url, bytes.NewReader([]byte("{}")))

internal/output/events.go

@@ -84,6 +84,11 @@ type PodSnapshotSavedEvent struct {
 	Size     int64
 }
 
+type SnapshotLoadedEvent struct {
+	Source   string   // display source shown to the user (e.g. "./snap.zip" or "pod:my-baseline")
+	Services []string // services restored
+}
+
 type AuthCompleteEvent struct{}
 
 // Event is a sealed marker — only event types in this package implement it,
@@ -98,7 +103,8 @@ func (AuthCompleteEvent) sealedEvent()     {}
 func (InstanceInfoEvent) sealedEvent()     {}
 func (TableEvent) sealedEvent()            {}
 func (ResourceSummaryEvent) sealedEvent()  {}
-func (PodSnapshotSavedEvent) sealedEvent() {}
+func (PodSnapshotSavedEvent) sealedEvent()  {}
+func (SnapshotLoadedEvent) sealedEvent()    {}
 func (ContainerStatusEvent) sealedEvent()  {}
 func (ProgressEvent) sealedEvent()         {}
 func (UserInputRequestEvent) sealedEvent() {}

internal/output/plain_format.go

@@ -42,6 +42,8 @@ func FormatEventLine(event Event) (string, bool) {
 		return formatResourceSummary(e), true
 	case PodSnapshotSavedEvent:
 		return formatPodSnapshotSaved(e), true
+	case SnapshotLoadedEvent:
+		return formatSnapshotLoaded(e), true
 	case AuthCompleteEvent:
 		return "", false
 	default:
@@ -198,6 +200,15 @@ func formatResourceSummary(e ResourceSummaryEvent) string {
 	return fmt.Sprintf("~ %d resources · %d services", e.Resources, e.Services)
 }
 
+func formatSnapshotLoaded(e SnapshotLoadedEvent) string {
+	var sb strings.Builder
+	sb.WriteString(SuccessMarker() + fmt.Sprintf(" Snapshot loaded from %s", e.Source))
+	if len(e.Services) > 0 {
+		sb.WriteString("\n• Services: " + strings.Join(e.Services, ", "))
+	}
+	return sb.String()
+}
+
 func formatPodSnapshotSaved(e PodSnapshotSavedEvent) string {
 	var sb strings.Builder
 	sb.WriteString(SuccessMarker() + fmt.Sprintf(" Snapshot saved to pod:%s", e.PodName))

internal/output/plain_format_test.go

@@ -210,6 +210,32 @@ func TestFormatEventLine(t *testing.T) {
 			want:   SuccessMarker() + " Snapshot saved to pod:minimal-pod",
 			wantOK: true,
 		},
+
+		// snapshot load events
+		{
+			name:   "snapshot loaded with services",
+			event:  SnapshotLoadedEvent{Source: "./my-baseline.zip", Services: []string{"s3", "dynamodb"}},
+			want:   SuccessMarker() + " Snapshot loaded from ./my-baseline.zip\n• Services: s3, dynamodb",
+			wantOK: true,
+		},
+		{
+			name:   "snapshot loaded no services",
+			event:  SnapshotLoadedEvent{Source: "./snap.zip"},
+			want:   SuccessMarker() + " Snapshot loaded from ./snap.zip",
+			wantOK: true,
+		},
+		{
+			name:   "pod snapshot loaded with services",
+			event:  SnapshotLoadedEvent{Source: "pod:my-baseline", Services: []string{"s3", "lambda"}},
+			want:   SuccessMarker() + " Snapshot loaded from pod:my-baseline\n• Services: s3, lambda",
+			wantOK: true,
+		},
+		{
+			name:   "pod snapshot loaded no services",
+			event:  SnapshotLoadedEvent{Source: "pod:empty-pod"},
+			want:   SuccessMarker() + " Snapshot loaded from pod:empty-pod",
+			wantOK: true,
+		},
 	}
 
 	for _, tt := range tests {