action/command/rollout.go

Rama main · réplica de lectura del repositorio original.

package command

import (
	"context"
	"fmt"
	"slices"
	"strings"
	"time"

	"github.com/google/uuid"
	"github.com/pkg/errors"
	"github.com/spf13/cobra"
	"go.uber.org/multierr"

	"github.com/bytebase/bytebase/action/args"
	"github.com/bytebase/bytebase/action/command/output"
	"github.com/bytebase/bytebase/action/world"
	v1pb "github.com/bytebase/bytebase/backend/generated-go/v1"
)

func NewRolloutCommand(w *world.World) *cobra.Command {
	// bytebase-action rollout flags
	cmdRollout := &cobra.Command{
		Use:               "rollout",
		Short:             "Rollout the migrate files",
		Args:              cobra.NoArgs,
		PersistentPreRunE: rolloutPreRun(w),
		RunE:              runRollout(w),
	}
	cmdRollout.Flags().StringVar(&w.TargetStage, "target-stage", "", "Rollout up to the target stage. Format: environments/{environment}.")
	cmdRollout.Flags().StringVar(&w.Plan, "plan", "", "The plan to rollout. Format: projects/{project}/plans/{plan}. Shadows file-pattern and targets.")
	cmdRollout.Flags().StringVar(&w.ReleaseIDTemplate, "release-id-template", "release_{date}-RC{iteration}", "Template for release ID. Available variables: {date}, {time}, {timestamp}, {iteration}")
	cmdRollout.Flags().StringVar(&w.ReleaseIDTimezone, "release-id-timezone", "UTC", "Timezone for {date} and {time} variables (e.g., 'UTC', 'America/Los_Angeles')")
	return cmdRollout
}

func rolloutPreRun(_ *world.World) func(*cobra.Command, []string) error {
	return func(cmd *cobra.Command, args []string) error {
		if p := cmd.Parent(); p != nil {
			if p.PersistentPreRunE != nil {
				if err := p.PersistentPreRunE(cmd, args); err != nil {
					return err
				}
			}
		}
		return nil
	}
}

func runRollout(w *world.World) func(command *cobra.Command, _ []string) error {
	return func(command *cobra.Command, _ []string) error {
		defer func() {
			output.WriteOutput(w)
		}()
		w.IsRollout = true
		ctx := command.Context()
		client, err := newClientFromWorld(w)
		if err != nil {
			return errors.Wrapf(err, "failed to create client")
		}
		defer client.close()

		// Check version compatibility
		if err := checkVersionCompatibility(w, client, args.Version); err != nil {
			return err
		}

		var plan *v1pb.Plan
		if w.Plan != "" {
			planP, err := client.getPlan(ctx, w.Plan)
			if err != nil {
				return errors.Wrapf(err, "failed to get plan")
			}
			plan = planP
			w.Logger.Info("use the provided plan", "url", fmt.Sprintf("%s/%s", client.url, plan.Name))
		} else {
			var release string
			releaseFiles, err := getReleaseFiles(w)
			if err != nil {
				return errors.Wrapf(err, "failed to get release files")
			}
			releaseType := v1pb.Release_VERSIONED
			if w.Declarative {
				releaseType = v1pb.Release_DECLARATIVE
			}
			createReleaseResponse, err := client.createRelease(ctx, w.Project, &v1pb.Release{
				Files:     releaseFiles,
				VcsSource: getVCSSource(w),
				Type:      releaseType,
			}, w.ReleaseIDTemplate, w.ReleaseIDTimezone)
			if err != nil {
				return errors.Wrapf(err, "failed to create release")
			}
			w.Logger.Info("release created", "url", fmt.Sprintf("%s/%s", client.url, createReleaseResponse.Name))
			release = createReleaseResponse.Name
			w.OutputMap.Release = release

			// Extract release ID from release name (format: projects/{project}/releases/{release_id})
			releaseID := strings.TrimPrefix(release, w.Project+"/releases/")

			planCreated, err := client.createPlan(ctx, w.Project, &v1pb.Plan{
				Title: releaseID,
				Specs: []*v1pb.Plan_Spec{
					{
						Id: uuid.New().String(),
						Config: &v1pb.Plan_Spec_ChangeDatabaseConfig{
							ChangeDatabaseConfig: &v1pb.Plan_ChangeDatabaseConfig{
								Targets: w.Targets,
								Release: release,
							},
						},
					},
				},
			})
			if err != nil {
				return errors.Wrapf(err, "failed to create plan")
			}
			plan = planCreated
			w.Logger.Info("plan created", "url", fmt.Sprintf("%s/%s", client.url, plan.Name))
		}
		w.OutputMap.Plan = plan.Name

		if err := runAndWaitForRollout(ctx, w, client, plan.Name); err != nil {
			return errors.Wrapf(err, "failed to run and wait for rollout")
		}

		return nil
	}
}

func runAndWaitForRollout(ctx context.Context, w *world.World, client *client, planName string) error {
	// create rollout with all stages created
	rollout, err := client.createRollout(ctx, &v1pb.CreateRolloutRequest{
		Parent: planName,
		Target: nil, // all stages
	})
	if err != nil {
		return errors.Wrapf(err, "failed to create rollout")
	}
	w.OutputMap.Rollout = rollout.Name
	w.Rollout = rollout

	w.Logger.Info("rollout created", "url", fmt.Sprintf("%s/%s", client.url, rollout.Name))

	return waitForRollout(ctx, w, client, rollout.Name)
}

func waitForRollout(ctx context.Context, w *world.World, client *client, rolloutName string) error {
	if w.TargetStage == "" {
		w.Logger.Info("target stage is not specified, exiting...")
		return nil
	}

	defer func() {
		if ctx.Err() == nil {
			return
		}
		w.Logger.Info("context cancelled, canceling the rollout")
		if err := cancelRollout(context.Background(), client, rolloutName); err != nil {
			w.Logger.Error("failed to cancel rollout", "error", err)
		}
	}()

	rollout, err := client.getRollout(ctx, rolloutName)
	if err != nil {
		return errors.Wrapf(err, "failed to get rollout")
	}

	w.Logger.Info("exit after the target stage is completed", "targetStage", w.TargetStage)
	w.Logger.Info("the rollout has the following stages",
		"createdStages", slices.Collect(func(yield func(string) bool) {
			for _, stage := range rollout.GetStages() {
				if !yield(stage.Environment) {
					return
				}
			}
		}))
	if len(rollout.Stages) == 0 {
		w.Logger.Info("no stages in the rollout. exiting...")
		return nil
	}

	targetStageFound := false
	for _, stage := range rollout.GetStages() {
		if stage.Environment == w.TargetStage {
			targetStageFound = true
			break
		}
	}
	if !targetStageFound {
		w.Logger.Info("the target stage is not found in the rollout. exiting...", "targetStage", w.TargetStage)
		return nil
	}

	for i := range rollout.Stages {
		// Polling loop for the current stage
		for {
			if ctx.Err() != nil {
				return errors.Wrapf(ctx.Err(), "context cancelled")
			}
			// get rollout to refresh status
			rollout, err := client.getRollout(ctx, rolloutName)
			if err != nil {
				return errors.Wrapf(err, "failed to get rollout")
			}
			w.Rollout = rollout
			if i >= len(rollout.Stages) {
				return errors.Errorf("rollout stage index out of bounds")
			}
			stage := rollout.Stages[i]

			// check stage tasks
			done := true
			var foundFailed, foundCanceled bool
			var notStartedTasks []string
			for _, task := range stage.GetTasks() {
				switch task.Status {
				case v1pb.Task_STATUS_UNSPECIFIED:
					done = false
				case v1pb.Task_NOT_STARTED:
					notStartedTasks = append(notStartedTasks, task.Name)
					done = false
				case v1pb.Task_PENDING:
					done = false
				case v1pb.Task_RUNNING:
					done = false
				case v1pb.Task_FAILED:
					foundFailed = true
					done = false
				case v1pb.Task_CANCELED:
					foundCanceled = true
					done = false
				case v1pb.Task_DONE:
				case v1pb.Task_SKIPPED:
				default:
					// Treat unknown task status as not done to be safe
					done = false
				}
			}

			if foundFailed {
				return errors.Errorf("found failed tasks. view on Bytebase")
			}
			if foundCanceled {
				return errors.Errorf("found canceled tasks. view on Bytebase")
			}

			if done {
				w.Logger.Info("stage completed", "stage", stage.Environment)
				if w.TargetStage == stage.Environment {
					return nil
				}
				// Break inner loop to proceed to next stage
				break
			}

			// run stage tasks
			if len(notStartedTasks) > 0 {
				w.Logger.Info("running stage tasks", "stage", stage.Environment, "taskCount", len(notStartedTasks))
				if _, err := client.batchRunTasks(ctx, &v1pb.BatchRunTasksRequest{
					Parent: stage.Name,
					Tasks:  notStartedTasks,
				}); err != nil {
					return errors.Wrapf(err, "failed to batch create tasks")
				}
			}
			time.Sleep(w.RolloutPollInterval)
		}
	}

	return nil
}

func cancelRollout(ctx context.Context, client *client, rolloutName string) error {
	taskRuns, err := client.listAllTaskRuns(ctx, rolloutName)
	if err != nil {
		return errors.Wrapf(err, "failed to list task runs")
	}
	taskRunsToCancelByStage := map[string][]string{}
	for _, taskRun := range taskRuns {
		if taskRun.Status == v1pb.TaskRun_RUNNING || taskRun.Status == v1pb.TaskRun_PENDING {
			stage := strings.Split(taskRun.Name, "/tasks")[0]
			taskRunsToCancelByStage[stage] = append(taskRunsToCancelByStage[stage], taskRun.Name)
		}
	}
	var errs error
	for stage, taskRuns := range taskRunsToCancelByStage {
		_, err := client.batchCancelTaskRuns(ctx, &v1pb.BatchCancelTaskRunsRequest{
			Parent:   stage + "/tasks/-",
			TaskRuns: taskRuns,
		})
		if err != nil {
			err = errors.Wrapf(err, "failed to cancel task runs for stage %v", stage)
			errs = multierr.Append(errs, err)
		}
	}
	return errs
}