Added --restart flag for bundle run command

bundle/run/job.go

@@ -15,6 +15,7 @@ import (
 	"github.com/databricks/cli/libs/log"
 	"github.com/databricks/databricks-sdk-go/service/jobs"
 	"github.com/fatih/color"
+	"golang.org/x/sync/errgroup"
 )
 
 // Default timeout for waiting for a job run to complete.
@@ -275,3 +276,42 @@ func (r *jobRunner) convertPythonParams(opts *Options) error {
 
 	return nil
 }
+
+func (r *jobRunner) Cancel(ctx context.Context) error {
+	w := r.bundle.WorkspaceClient()
+	jobID, err := strconv.ParseInt(r.job.ID, 10, 64)
+	if err != nil {
+		return fmt.Errorf("job ID is not an integer: %s", r.job.ID)
+	}
+
+	runs, err := w.Jobs.ListRunsAll(ctx, jobs.ListRunsRequest{
+		ActiveOnly: true,
+		JobId:      jobID,
+	})
+
+	if err != nil {
+		return err
+	}
+
+	if len(runs) == 0 {
+		return nil
+	}
+
+	errGroup, errCtx := errgroup.WithContext(ctx)
+	for _, run := range runs {
+		runId := run.RunId
+		errGroup.Go(func() error {
+			wait, err := w.Jobs.CancelRun(errCtx, jobs.CancelRun{
+				RunId: runId,
+			})
+			if err != nil {
+				return err
+			}
+			// Waits for the Terminated or Skipped state
+			_, err = wait.GetWithTimeout(jobRunTimeout)
+			return err
+		})
+	}
+
+	return errGroup.Wait()
+}

bundle/run/pipeline.go

@@ -166,3 +166,18 @@ func (r *pipelineRunner) Run(ctx context.Context, opts *Options) (output.RunOutp
 		time.Sleep(time.Second)
 	}
 }
+
+func (r *pipelineRunner) Cancel(ctx context.Context) error {
+	w := r.bundle.WorkspaceClient()
+	wait, err := w.Pipelines.Stop(ctx, pipelines.StopRequest{
+		PipelineId: r.pipeline.ID,
+	})
+
+	if err != nil {
+		return err
+	}
+
+	// Waits for the Idle state of the pipeline
+	_, err = wait.GetWithTimeout(jobRunTimeout)
+	return err
+}

bundle/run/runner.go

@@ -26,6 +26,9 @@ type Runner interface {
 
 	// Run the underlying worklow.
 	Run(ctx context.Context, opts *Options) (output.RunOutput, error)
+
+	// Cancel the underlying workflow.
+	Cancel(ctx context.Context) error
 }
 
 // Find locates a runner matching the specified argument.

cmd/bundle/run.go

@@ -27,7 +27,9 @@ func newRunCommand() *cobra.Command {
 	runOptions.Define(cmd)
 
 	var noWait bool
+	var restart bool
 	cmd.Flags().BoolVar(&noWait, "no-wait", false, "Don't wait for the run to complete.")
+	cmd.Flags().BoolVar(&restart, "restart", false, "Restart the run if it is already running.")
 
 	cmd.RunE = func(cmd *cobra.Command, args []string) error {
 		ctx := cmd.Context()
@@ -68,6 +70,15 @@ func newRunCommand() *cobra.Command {
 		}
 
 		runOptions.NoWait = noWait
+		if restart {
+			s := cmdio.Spinner(ctx)
+			s <- "Cancelling all runs"
+			err := runner.Cancel(ctx)
+			close(s)
+			if err != nil {
+				return err
+			}
+		}
 		output, err := runner.Run(ctx, &runOptions)
 		if err != nil {
 			return err