Add snapshot save command

Author: anisaoshafi

cmd/root.go

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

cmd/snapshot.go

@@ -0,0 +1,78 @@
+package cmd
+
+import (
+	"fmt"
+	"os"
+
+	"github.com/localstack/lstk/internal/config"
+	"github.com/localstack/lstk/internal/endpoint"
+	"github.com/localstack/lstk/internal/env"
+	"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, tel *telemetry.Client) *cobra.Command {
+	cmd := &cobra.Command{
+		Use:   "snapshot",
+		Short: "Manage emulator snapshots",
+	}
+	cmd.AddCommand(newSnapshotSaveCmd(cfg, tel))
+	return cmd
+}
+
+func newSnapshotSaveCmd(cfg *env.Env, tel *telemetry.Client) *cobra.Command {
+	return &cobra.Command{
+		Use:   "save [destination]",
+		Short: "Save a snapshot of the emulator state",
+		Long: `Save a snapshot of the running emulator's state to a local file.
+
+The destination must be a file path. Use a path prefix to save locally:
+
+  lstk snapshot save                  # saves to ./ls-state-export
+  lstk snapshot save ./my-snapshot    # saves to ./my-snapshot
+  lstk snapshot save /tmp/my-state    # saves to /tmp/my-state
+
+Cloud destinations are not yet supported.`,
+		Args:    cobra.MaximumNArgs(1),
+		PreRunE: initConfig,
+		RunE: commandWithTelemetry("snapshot save", tel, func(cmd *cobra.Command, args []string) error {
+			var destArg string
+			if len(args) > 0 {
+				destArg = args[0]
+			}
+
+			dest, err := snapshot.ParseDestination(destArg)
+			if err != nil {
+				return err
+			}
+
+			rt, err := runtime.NewDockerRuntime(cfg.DockerHost)
+			if err != nil {
+				return err
+			}
+
+			appConfig, err := config.Get()
+			if err != nil {
+				return fmt.Errorf("failed to get config: %w", err)
+			}
+			if len(appConfig.Containers) == 0 {
+				return fmt.Errorf("no emulator configured")
+			}
+
+			c := appConfig.Containers[0]
+			host, _ := endpoint.ResolveHost(c.Port, cfg.LocalStackHost)
+			exporter := snapshot.NewStateClient("http://" + host)
+
+			opts := snapshot.SaveOptions{}
+
+			if isInteractiveMode(cfg) {
+				return ui.RunSnapshotSave(cmd.Context(), rt, appConfig.Containers, exporter, dest, opts)
+			}
+			return snapshot.Save(cmd.Context(), rt, appConfig.Containers, exporter, dest, output.NewPlainSink(os.Stdout), opts)
+		}),
+	}
+}

internal/snapshot/client.go

@@ -0,0 +1,45 @@
+package snapshot
+
+import (
+	"context"
+	"fmt"
+	"io"
+	"net/http"
+)
+
+// StateExporter retrieves state from the running LocalStack instance.
+type StateExporter interface {
+	ExportState(ctx context.Context) (io.ReadCloser, error)
+}
+
+// StateClient calls the LocalStack state API.
+type StateClient struct {
+	baseURL    string
+	httpClient *http.Client
+}
+
+func NewStateClient(baseURL string) *StateClient {
+	return &StateClient{
+		baseURL:    baseURL,
+		httpClient: &http.Client{},
+	}
+}
+
+// ExportState calls GET /_localstack/pods/state and returns the response body.
+// The caller is responsible for closing the returned ReadCloser.
+func (c *StateClient) ExportState(ctx context.Context) (io.ReadCloser, error) {
+	req, err := http.NewRequestWithContext(ctx, http.MethodGet, c.baseURL+"/_localstack/pods/state", nil)
+	if err != nil {
+		return nil, fmt.Errorf("create request: %w", err)
+	}
+
+	resp, err := c.httpClient.Do(req)
+	if err != nil {
+		return nil, fmt.Errorf("connect to LocalStack: %w", err)
+	}
+	if resp.StatusCode != http.StatusOK {
+		_ = resp.Body.Close()
+		return nil, fmt.Errorf("LocalStack returned status %d", resp.StatusCode)
+	}
+	return resp.Body, nil
+}

internal/snapshot/client_test.go

@@ -0,0 +1,114 @@
+package snapshot_test
+
+import (
+	"context"
+	"io"
+	"net/http"
+	"net/http/httptest"
+	"strings"
+	"testing"
+
+	"github.com/localstack/lstk/internal/snapshot"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+)
+
+func TestStateClient_ExportState_OK(t *testing.T) {
+	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		assert.Equal(t, "/_localstack/pods/state", r.URL.Path)
+		assert.Equal(t, http.MethodGet, r.Method)
+		w.WriteHeader(http.StatusOK)
+		_, _ = w.Write([]byte("ZIP_DATA"))
+	}))
+	defer srv.Close()
+
+	client := snapshot.NewStateClient(srv.URL)
+	body, err := client.ExportState(context.Background())
+	require.NoError(t, err)
+	defer func() { _ = body.Close() }()
+
+	data, err := io.ReadAll(body)
+	require.NoError(t, err)
+	assert.Equal(t, "ZIP_DATA", string(data))
+}
+
+func TestStateClient_ExportState_ServerError(t *testing.T) {
+	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		w.WriteHeader(http.StatusInternalServerError)
+	}))
+	defer srv.Close()
+
+	client := snapshot.NewStateClient(srv.URL)
+	_, err := client.ExportState(context.Background())
+	require.Error(t, err)
+	assert.Contains(t, err.Error(), "500")
+}
+
+func TestStateClient_ExportState_NotFound(t *testing.T) {
+	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		w.WriteHeader(http.StatusNotFound)
+	}))
+	defer srv.Close()
+
+	client := snapshot.NewStateClient(srv.URL)
+	_, err := client.ExportState(context.Background())
+	require.Error(t, err)
+	assert.Contains(t, err.Error(), "404")
+}
+
+func TestStateClient_ExportState_ConnectionRefused(t *testing.T) {
+	// Bind then immediately close to get a port that refuses connections.
+	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {}))
+	addr := srv.URL
+	srv.Close()
+
+	client := snapshot.NewStateClient(addr)
+	_, err := client.ExportState(context.Background())
+	require.Error(t, err)
+	assert.Contains(t, err.Error(), "connect to LocalStack")
+}
+
+func TestStateClient_ExportState_ContextCancelled(t *testing.T) {
+	started := make(chan struct{})
+	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		close(started)
+		// block until the client cancels
+		<-r.Context().Done()
+	}))
+	defer srv.Close()
+
+	ctx, cancel := context.WithCancel(context.Background())
+	client := snapshot.NewStateClient(srv.URL)
+
+	errCh := make(chan error, 1)
+	go func() {
+		_, err := client.ExportState(ctx)
+		errCh <- err
+	}()
+
+	<-started
+	cancel()
+
+	err := <-errCh
+	require.Error(t, err)
+}
+
+func TestStateClient_ExportState_LargeBody(t *testing.T) {
+	const size = 1 << 20 // 1 MB
+	payload := strings.Repeat("X", size)
+
+	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		w.WriteHeader(http.StatusOK)
+		_, _ = w.Write([]byte(payload))
+	}))
+	defer srv.Close()
+
+	client := snapshot.NewStateClient(srv.URL)
+	body, err := client.ExportState(context.Background())
+	require.NoError(t, err)
+	defer func() { _ = body.Close() }()
+
+	data, err := io.ReadAll(body)
+	require.NoError(t, err)
+	assert.Equal(t, size, len(data))
+}

internal/snapshot/destination.go

@@ -0,0 +1,49 @@
+package snapshot
+
+import (
+	"fmt"
+	"io"
+	"os"
+	"strings"
+)
+
+// Destination is where snapshot state is written.
+// LocalDestination writes to a local file; future implementations handle cloud storage.
+type Destination interface {
+	Writer() (io.WriteCloser, error)
+	String() string
+}
+
+// ParseDestination parses the user-supplied destination string.
+//
+// Disambiguation rules:
+//   - empty string          → LocalDestination at default path "ls-state-export"
+//   - contains "://"        → unsupported cloud URI error
+//   - starts with ".", "/", or "~", or contains "/" → LocalDestination
+//   - bare name (no path chars, no scheme) → reserved for future cloud pod names; returns error
+func ParseDestination(dest string) (Destination, error) {
+	if dest == "" {
+		return LocalDestination{Path: "ls-state-export"}, nil
+	}
+	if strings.Contains(dest, "://") {
+		return nil, fmt.Errorf("cloud destinations are not yet supported — use a file path like ./my-snapshot")
+	}
+	if strings.HasPrefix(dest, ".") || strings.HasPrefix(dest, "/") || strings.HasPrefix(dest, "~") || strings.Contains(dest, "/") {
+		return LocalDestination{Path: dest}, nil
+	}
+	// bare name with no path separators: reserved for future cloud pod names
+	return nil, fmt.Errorf("cloud destinations are not yet supported — use a file path like ./my-snapshot")
+}
+
+// LocalDestination writes snapshot state to a local file.
+type LocalDestination struct {
+	Path string
+}
+
+func (d LocalDestination) Writer() (io.WriteCloser, error) {
+	return os.Create(d.Path)
+}
+
+func (d LocalDestination) String() string {
+	return d.Path
+}

internal/snapshot/destination_test.go

@@ -0,0 +1,64 @@
+package snapshot_test
+
+import (
+	"testing"
+
+	"github.com/localstack/lstk/internal/snapshot"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+)
+
+func TestParseDestination(t *testing.T) {
+	tests := []struct {
+		input   string
+		want    snapshot.Destination
+		wantErr string
+	}{
+		{
+			input: "",
+			want:  snapshot.LocalDestination{Path: "ls-state-export"},
+		},
+		{
+			input: "./my-state",
+			want:  snapshot.LocalDestination{Path: "./my-state"},
+		},
+		{
+			input: "/tmp/state",
+			want:  snapshot.LocalDestination{Path: "/tmp/state"},
+		},
+		{
+			input: "~/snapshots/s",
+			want:  snapshot.LocalDestination{Path: "~/snapshots/s"},
+		},
+		{
+			input: "subdir/state",
+			want:  snapshot.LocalDestination{Path: "subdir/state"},
+		},
+		{
+			input:   "my-pod",
+			wantErr: "cloud destinations are not yet supported",
+		},
+		{
+			input:   "cloud://my-pod",
+			wantErr: "cloud destinations are not yet supported",
+		},
+		{
+			input:   "s3://bucket/key",
+			wantErr: "cloud destinations are not yet supported",
+		},
+	}
+
+	for _, tc := range tests {
+		t.Run(tc.input, func(t *testing.T) {
+			got, err := snapshot.ParseDestination(tc.input)
+			if tc.wantErr != "" {
+				require.Error(t, err)
+				assert.Contains(t, err.Error(), tc.wantErr)
+				assert.Contains(t, err.Error(), "./my-snapshot")
+				return
+			}
+			require.NoError(t, err)
+			assert.Equal(t, tc.want, got)
+		})
+	}
+}