add reset branch

docs/generate_doc/ticloud_serverless_branch.md

@@ -23,5 +23,6 @@ Manage TiDB Cloud Serverless branches
 * [ticloud serverless branch delete](ticloud_serverless_branch_delete.md)	 - Delete a branch
 * [ticloud serverless branch describe](ticloud_serverless_branch_describe.md)	 - Describe a branch
 * [ticloud serverless branch list](ticloud_serverless_branch_list.md)	 - List branches
+* [ticloud serverless branch reset](ticloud_serverless_branch_reset.md)	 - Reset a branch to its parent's current state
 * [ticloud serverless branch shell](ticloud_serverless_branch_shell.md)	 - Connect to a branch
 

docs/generate_doc/ticloud_serverless_branch_reset.md

@@ -0,0 +1,39 @@
+## ticloud serverless branch reset
+
+Reset a branch to its parent's current state
+
+```
+ticloud serverless branch reset [flags]
+```
+
+### Examples
+
+```
+  Reset a branch in interactive mode:
+  $ ticloud serverless branch reset
+
+  Reset a branch in non-interactive mode:
+  $ ticloud serverless branch reset -c <cluster-id> -b <branch-id>
+```
+
+### Options
+
+```
+  -b, --branch-id string    The ID of the branch to be reset.
+  -c, --cluster-id string   The cluster ID of the branch to be reset.
+      --force               Reset a branch without confirmation.
+  -h, --help                help for reset
+```
+
+### Options inherited from parent commands
+
+```
+  -D, --debug            Enable debug mode
+      --no-color         Disable color output
+  -P, --profile string   Profile to use from your configuration file
+```
+
+### SEE ALSO
+
+* [ticloud serverless branch](ticloud_serverless_branch.md)	 - Manage TiDB Cloud Serverless branches
+

internal/cli/serverless/branch/reset.go

@@ -15,13 +15,18 @@
 package branch
 
 import (
+	"context"
 	"fmt"
+	"time"
 
+	tea "github.com/charmbracelet/bubbletea"
 	"github.com/tidbcloud/tidbcloud-cli/internal"
 	"github.com/tidbcloud/tidbcloud-cli/internal/config"
 	"github.com/tidbcloud/tidbcloud-cli/internal/flag"
 	"github.com/tidbcloud/tidbcloud-cli/internal/service/cloud"
+	"github.com/tidbcloud/tidbcloud-cli/internal/ui"
 	"github.com/tidbcloud/tidbcloud-cli/internal/util"
+	"github.com/tidbcloud/tidbcloud-cli/pkg/tidbcloud/v1beta1/serverless/branch"
 
 	"github.com/AlecAivazis/survey/v2"
 	"github.com/AlecAivazis/survey/v2/terminal"
@@ -70,7 +75,7 @@ func ResetCmd(h *internal.Helper) *cobra.Command {
 	var force bool
 	var resetCmd = &cobra.Command{
 		Use:   "reset",
-		Short: "Reset a branch",
+		Short: "Reset a branch to its parent's current state",
 		Args:  cobra.NoArgs,
 		Example: fmt.Sprintf(`  Reset a branch in interactive mode:
   $ %[1]s serverless branch reset
@@ -155,12 +160,19 @@ func ResetCmd(h *internal.Helper) *cobra.Command {
 				}
 			}
 
-			_, err = d.ResetBranch(ctx, clusterID, branchID)
-			if err != nil {
-				return errors.Trace(err)
-			}
 			// print success for reset branch is a sync operation
-			fmt.Fprintln(h.IOStreams.Out, color.GreenString("branch %s reset", branchID))
+			if h.IOStreams.CanPrompt {
+				err := ResetAndSpinnerWait(ctx, h, d, clusterID, branchID)
+				if err != nil {
+					return errors.Trace(err)
+				}
+			} else {
+				err := ResetAndWaitReady(ctx, h, d, clusterID, branchID)
+				if err != nil {
+					return err
+				}
+			}
+
 			return nil
 		},
 	}
@@ -171,3 +183,75 @@ func ResetCmd(h *internal.Helper) *cobra.Command {
 
 	return resetCmd
 }
+
+func ResetAndWaitReady(ctx context.Context, h *internal.Helper, d cloud.TiDBCloudClient, clusterID string, branchID string) error {
+	_, err := d.ResetBranch(ctx, clusterID, branchID)
+	if err != nil {
+		return errors.Trace(err)
+	}
+
+	fmt.Fprintln(h.IOStreams.Out, "... Waiting for branch to be ready")
+	ticker := time.NewTicker(WaitInterval)
+	defer ticker.Stop()
+	timer := time.After(WaitTimeout)
+	for {
+		select {
+		case <-timer:
+			return errors.New(fmt.Sprintf("Timeout waiting for branch %s to be ready, please check status on dashboard.", branchID))
+		case <-ticker.C:
+			b, err := d.GetBranch(ctx, clusterID, branchID)
+			if err != nil {
+				return errors.Trace(err)
+			}
+			if *b.State == branch.BRANCHSTATE_ACTIVE {
+				fmt.Fprint(h.IOStreams.Out, color.GreenString("Branch %s is ready.", branchID))
+				return nil
+			}
+		}
+	}
+}
+
+func ResetAndSpinnerWait(ctx context.Context, h *internal.Helper, d cloud.TiDBCloudClient, clusterID string, branchID string) error {
+	// use spinner to indicate that the cluster is resetting
+	task := func() tea.Msg {
+		_, err := d.ResetBranch(ctx, clusterID, branchID)
+		if err != nil {
+			return errors.Trace(err)
+		}
+
+		ticker := time.NewTicker(WaitInterval)
+		defer ticker.Stop()
+		timer := time.After(WaitTimeout)
+		for {
+			select {
+			case <-timer:
+				return ui.Result(fmt.Sprintf("Timeout waiting for branch %s to be ready, please check status on dashboard.", branchID))
+			case <-ticker.C:
+				b, err := d.GetBranch(ctx, clusterID, branchID)
+				if err != nil {
+					return errors.Trace(err)
+				}
+				if *b.State == branch.BRANCHSTATE_ACTIVE {
+					return ui.Result(fmt.Sprintf("Branch %s is ready.", branchID))
+				}
+			case <-ctx.Done():
+				return util.InterruptError
+			}
+		}
+	}
+
+	p := tea.NewProgram(ui.InitialSpinnerModel(task, "Waiting for branch to be ready"))
+	resetModel, err := p.Run()
+	if err != nil {
+		return errors.Trace(err)
+	}
+	if m, _ := resetModel.(ui.SpinnerModel); m.Interrupted {
+		return util.InterruptError
+	}
+	if m, _ := resetModel.(ui.SpinnerModel); m.Err != nil {
+		return m.Err
+	} else {
+		fmt.Fprintln(h.IOStreams.Out, color.GreenString(m.Output))
+	}
+	return nil
+}

internal/cli/serverless/branch/reset_test.go

@@ -0,0 +1,115 @@
+// Copyright 2024 PingCAP, Inc.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//      http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package branch
+
+import (
+	"bytes"
+	"context"
+	"encoding/json"
+	"fmt"
+	"os"
+	"testing"
+
+	"github.com/tidbcloud/tidbcloud-cli/internal"
+	"github.com/tidbcloud/tidbcloud-cli/internal/iostream"
+	"github.com/tidbcloud/tidbcloud-cli/internal/mock"
+	"github.com/tidbcloud/tidbcloud-cli/internal/service/cloud"
+	"github.com/tidbcloud/tidbcloud-cli/pkg/tidbcloud/v1beta1/serverless/branch"
+
+	"github.com/stretchr/testify/require"
+	"github.com/stretchr/testify/suite"
+)
+
+type ResetBranchSuite struct {
+	suite.Suite
+	h          *internal.Helper
+	mockClient *mock.TiDBCloudClient
+}
+
+func (suite *ResetBranchSuite) SetupTest() {
+	if err := os.Setenv("NO_COLOR", "true"); err != nil {
+		suite.T().Error(err)
+	}
+
+	var pageSize int64 = 10
+	suite.mockClient = new(mock.TiDBCloudClient)
+	suite.h = &internal.Helper{
+		Client: func() (cloud.TiDBCloudClient, error) {
+			return suite.mockClient, nil
+		},
+		QueryPageSize: pageSize,
+		IOStreams:     iostream.Test(),
+	}
+}
+
+func (suite *ResetBranchSuite) TestResetBranchArgs() {
+	assert := require.New(suite.T())
+	ctx := context.Background()
+
+	clusterID := "12345"
+	branchID := "12345"
+	suite.mockClient.On("ResetBranch", ctx, clusterID, branchID).Return(&branch.Branch{}, nil)
+
+	body := &branch.Branch{}
+	err := json.Unmarshal([]byte(getBranchResultStr), body)
+	assert.Nil(err)
+	suite.mockClient.On("GetBranch", ctx, clusterID, branchID).Return(body, nil)
+
+	tests := []struct {
+		name         string
+		args         []string
+		err          error
+		stdoutString string
+		stderrString string
+	}{
+		{
+			name:         "reset branch success",
+			args:         []string{"--cluster-id", clusterID, "--branch-id", branchID, "--force"},
+			stdoutString: fmt.Sprintf("... Waiting for branch to be ready\nBranch %s is ready.", branchID),
+		},
+		{
+			name: "reset branch without force",
+			args: []string{"--cluster-id", clusterID, "--branch-id", branchID},
+			err:  fmt.Errorf("the terminal doesn't support prompt, please run with --force to reset the branch"),
+		},
+		{
+			name: "reset branch without required branch id",
+			args: []string{"-c", clusterID, "--force"},
+			err:  fmt.Errorf("required flag(s) \"branch-id\" not set"),
+		},
+	}
+
+	for _, tt := range tests {
+		suite.T().Run(tt.name, func(t *testing.T) {
+			cmd := ResetCmd(suite.h)
+			cmd.SetContext(ctx)
+			suite.h.IOStreams.Out.(*bytes.Buffer).Reset()
+			suite.h.IOStreams.Err.(*bytes.Buffer).Reset()
+			cmd.SetArgs(tt.args)
+			err := cmd.Execute()
+			assert.Equal(tt.err, err)
+
+			assert.Equal(tt.stdoutString, suite.h.IOStreams.Out.(*bytes.Buffer).String())
+			assert.Equal(tt.stderrString, suite.h.IOStreams.Err.(*bytes.Buffer).String())
+			if tt.err == nil {
+				suite.mockClient.AssertExpectations(suite.T())
+			}
+		})
+	}
+}
+
+func TestResetBranchSuite(t *testing.T) {
+	suite.Run(t, new(ResetBranchSuite))
+}