From c96d332e3eda3a9bbf3822ce458f8fba7dd47319 Mon Sep 17 00:00:00 2001
From: Pudong Zheng <tianshi8650@gmail.com>
Date: Tue, 7 Dec 2021 10:18:35 +0000
Subject: [PATCH] [gp-cli] allow to create snapshot

---
 components/gitpod-cli/cmd/snapshot.go         | 78 +++++++++++++++++++
 .../gitpod-protocol/go/gitpod-service.go      | 12 +--
 components/gitpod-protocol/go/mock.go         |  4 +-
 3 files changed, 84 insertions(+), 10 deletions(-)
 create mode 100644 components/gitpod-cli/cmd/snapshot.go

diff --git a/components/gitpod-cli/cmd/snapshot.go b/components/gitpod-cli/cmd/snapshot.go
new file mode 100644
index 00000000000000..c9584932cdcb5b
--- /dev/null
+++ b/components/gitpod-cli/cmd/snapshot.go
@@ -0,0 +1,78 @@
+// Copyright (c) 2020 Gitpod GmbH. All rights reserved.
+// Licensed under the GNU Affero General Public License (AGPL).
+// See License-AGPL.txt in the project root for license information.
+
+package cmd
+
+import (
+	"context"
+	"errors"
+	"fmt"
+	"os"
+	"os/signal"
+	"syscall"
+	"time"
+
+	gitpod "github.com/gitpod-io/gitpod/gitpod-cli/pkg/gitpod"
+	protocol "github.com/gitpod-io/gitpod/gitpod-protocol"
+	"github.com/sourcegraph/jsonrpc2"
+	"github.com/spf13/cobra"
+)
+
+const (
+	ErrorCodeSnapshotNotFound = 404
+	ErrorCodeSnapshotError    = 630
+)
+
+// snapshotCmd represents the snapshotCmd command
+var snapshotCmd = &cobra.Command{
+	Use:   "snapshot",
+	Short: "Take a snapshot of the current workspace",
+	Args:  cobra.ArbitraryArgs,
+	Run: func(cmd *cobra.Command, args []string) {
+		ctx, cancel := context.WithCancel(context.Background())
+		sigChan := make(chan os.Signal, 1)
+		signal.Notify(sigChan, os.Interrupt, syscall.SIGTERM)
+		go func() {
+			<-sigChan
+			cancel()
+		}()
+		wsInfo, err := gitpod.GetWSInfo(ctx)
+		if err != nil {
+			fail(err.Error())
+		}
+		client, err := gitpod.ConnectToServer(ctx, wsInfo, []string{
+			"function:takeSnapshot",
+			"function:waitForSnapshot",
+			"resource:workspace::" + wsInfo.WorkspaceId + "::get/update",
+		})
+		if err != nil {
+			fail(err.Error())
+		}
+		snapshotId, err := client.TakeSnapshot(ctx, &protocol.TakeSnapshotOptions{
+			WorkspaceID: wsInfo.WorkspaceId,
+			DontWait:    true,
+		})
+		if err != nil {
+			fail(err.Error())
+		}
+		for ctx.Err() == nil {
+			err := client.WaitForSnapshot(ctx, snapshotId)
+			if err != nil {
+				var responseErr *jsonrpc2.Error
+				if errors.As(err, &responseErr) && (responseErr.Code == ErrorCodeSnapshotNotFound || responseErr.Code == ErrorCodeSnapshotError) {
+					panic(err)
+				}
+				time.Sleep(time.Second * 3)
+			} else {
+				break
+			}
+		}
+		url := fmt.Sprintf("%s/#snapshot/%s", wsInfo.GitpodHost, snapshotId)
+		fmt.Println(url)
+	},
+}
+
+func init() {
+	rootCmd.AddCommand(snapshotCmd)
+}
diff --git a/components/gitpod-protocol/go/gitpod-service.go b/components/gitpod-protocol/go/gitpod-service.go
index 3a8fa6c0c886a4..4978aaa2f54cc5 100644
--- a/components/gitpod-protocol/go/gitpod-service.go
+++ b/components/gitpod-protocol/go/gitpod-service.go
@@ -73,7 +73,7 @@ type APIInterface interface {
 	SendFeedback(ctx context.Context, feedback string) (res string, err error)
 	RegisterGithubApp(ctx context.Context, installationID string) (err error)
 	TakeSnapshot(ctx context.Context, options *TakeSnapshotOptions) (res string, err error)
-	WaitForSnapshot(ctx context.Context, options *WaitForSnapshotOptions) (err error)
+	WaitForSnapshot(ctx context.Context, snapshotId string) (err error)
 	GetSnapshots(ctx context.Context, workspaceID string) (res []*string, err error)
 	StoreLayout(ctx context.Context, workspaceID string, layoutData string) (err error)
 	GetLayout(ctx context.Context, workspaceID string) (res string, err error)
@@ -1279,14 +1279,14 @@ func (gp *APIoverJSONRPC) TakeSnapshot(ctx context.Context, options *TakeSnapsho
 }
 
 // WaitForSnapshot calls waitForSnapshot on the server
-func (gp *APIoverJSONRPC) WaitForSnapshot(ctx context.Context, options *WaitForSnapshotOptions) (err error) {
+func (gp *APIoverJSONRPC) WaitForSnapshot(ctx context.Context, snapshotId string) (err error) {
 	if gp == nil {
 		err = errNotConnected
 		return
 	}
 	var _params []interface{}
 
-	_params = append(_params, options)
+	_params = append(_params, snapshotId)
 
 	var result string
 	err = gp.C.Call(ctx, "waitForSnapshot", _params, &result)
@@ -1939,11 +1939,7 @@ type GenerateNewGitpodTokenOptions struct {
 type TakeSnapshotOptions struct {
 	LayoutData  string `json:"layoutData,omitempty"`
 	WorkspaceID string `json:"workspaceId,omitempty"`
-}
-
-// WaitForSnapshotOptions is the WaitForSnapshotOptions message type
-type WaitForSnapshotOptions struct {
-	SnapshotID string `json:"snapshotID,omitempty"`
+	DontWait    bool   `json:"dontWait",omitempty`
 }
 
 // PreparePluginUploadParams is the PreparePluginUploadParams message type
diff --git a/components/gitpod-protocol/go/mock.go b/components/gitpod-protocol/go/mock.go
index 0aad6982722fd2..a846d47bcd0acd 100644
--- a/components/gitpod-protocol/go/mock.go
+++ b/components/gitpod-protocol/go/mock.go
@@ -805,9 +805,9 @@ func (mr *MockAPIInterfaceMockRecorder) TakeSnapshot(ctx, options interface{}) *
 }
 
 // WaitForSnapshot mocks base method.
-func (m *MockAPIInterface) WaitForSnapshot(ctx context.Context, options *WaitForSnapshotOptions) error {
+func (m *MockAPIInterface) WaitForSnapshot(ctx context.Context, snapshotId string) error {
 	m.ctrl.T.Helper()
-	ret := m.ctrl.Call(m, "WaitForSnapshot", ctx, options)
+	ret := m.ctrl.Call(m, "WaitForSnapshot", ctx, snapshotId)
 	ret0, _ := ret[0].(error)
 	return ret0
 }