Skip to content

Commit bc30c9e

Browse files
andrewnesterandrewnester
authored
Added --restart flag for bundle run command (#1191)
## Changes Added `--restart` flag for `bundle run` command When running with this flag, `bundle run` will cancel all existing runs before starting a new one ## Tests Manually
1 parent cac112c commit bc30c9e

6 files changed

Lines changed: 197 additions & 0 deletions

File tree

‎bundle/run/job.go‎

Lines changed: 40 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -15,6 +15,7 @@ import (
1515
"github.com/databricks/cli/libs/log"
1616
"github.com/databricks/databricks-sdk-go/service/jobs"
1717
"github.com/fatih/color"
18+
"golang.org/x/sync/errgroup"
1819
)
1920

2021
// Default timeout for waiting for a job run to complete.
@@ -275,3 +276,42 @@ func (r *jobRunner) convertPythonParams(opts *Options) error {
275276

276277
return nil
277278
}
279+
280+
func (r *jobRunner) Cancel(ctx context.Context) error {
281+
w := r.bundle.WorkspaceClient()
282+
jobID, err := strconv.ParseInt(r.job.ID, 10, 64)
283+
if err != nil {
284+
return fmt.Errorf("job ID is not an integer: %s", r.job.ID)
285+
}
286+
287+
runs, err := w.Jobs.ListRunsAll(ctx, jobs.ListRunsRequest{
288+
ActiveOnly: true,
289+
JobId: jobID,
290+
})
291+
292+
if err != nil {
293+
return err
294+
}
295+
296+
if len(runs) == 0 {
297+
return nil
298+
}
299+
300+
errGroup, errCtx := errgroup.WithContext(ctx)
301+
for _, run := range runs {
302+
runId := run.RunId
303+
errGroup.Go(func() error {
304+
wait, err := w.Jobs.CancelRun(errCtx, jobs.CancelRun{
305+
RunId: runId,
306+
})
307+
if err != nil {
308+
return err
309+
}
310+
// Waits for the Terminated or Skipped state
311+
_, err = wait.GetWithTimeout(jobRunTimeout)
312+
return err
313+
})
314+
}
315+
316+
return errGroup.Wait()
317+
}

‎bundle/run/job_test.go‎

Lines changed: 79 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,12 +1,16 @@
11
package run
22

33
import (
4+
"context"
45
"testing"
6+
"time"
57

68
"github.com/databricks/cli/bundle"
79
"github.com/databricks/cli/bundle/config"
810
"github.com/databricks/cli/bundle/config/resources"
11+
"github.com/databricks/databricks-sdk-go/experimental/mocks"
912
"github.com/databricks/databricks-sdk-go/service/jobs"
13+
"github.com/stretchr/testify/mock"
1014
"github.com/stretchr/testify/require"
1115
)
1216

@@ -47,3 +51,78 @@ func TestConvertPythonParams(t *testing.T) {
4751
require.Contains(t, opts.Job.notebookParams, "__python_params")
4852
require.Equal(t, opts.Job.notebookParams["__python_params"], `["param1","param2","param3"]`)
4953
}
54+
55+
func TestJobRunnerCancel(t *testing.T) {
56+
job := &resources.Job{
57+
ID: "123",
58+
}
59+
b := &bundle.Bundle{
60+
Config: config.Root{
61+
Resources: config.Resources{
62+
Jobs: map[string]*resources.Job{
63+
"test_job": job,
64+
},
65+
},
66+
},
67+
}
68+
69+
runner := jobRunner{key: "test", bundle: b, job: job}
70+
71+
m := mocks.NewMockWorkspaceClient(t)
72+
b.SetWorkpaceClient(m.WorkspaceClient)
73+
74+
jobApi := m.GetMockJobsAPI()
75+
jobApi.EXPECT().ListRunsAll(mock.Anything, jobs.ListRunsRequest{
76+
ActiveOnly: true,
77+
JobId: 123,
78+
}).Return([]jobs.BaseRun{
79+
{RunId: 1},
80+
{RunId: 2},
81+
}, nil)
82+
83+
mockWait := &jobs.WaitGetRunJobTerminatedOrSkipped[struct{}]{
84+
Poll: func(time time.Duration, f func(j *jobs.Run)) (*jobs.Run, error) {
85+
return nil, nil
86+
},
87+
}
88+
jobApi.EXPECT().CancelRun(mock.Anything, jobs.CancelRun{
89+
RunId: 1,
90+
}).Return(mockWait, nil)
91+
jobApi.EXPECT().CancelRun(mock.Anything, jobs.CancelRun{
92+
RunId: 2,
93+
}).Return(mockWait, nil)
94+
95+
err := runner.Cancel(context.Background())
96+
require.NoError(t, err)
97+
}
98+
99+
func TestJobRunnerCancelWithNoActiveRuns(t *testing.T) {
100+
job := &resources.Job{
101+
ID: "123",
102+
}
103+
b := &bundle.Bundle{
104+
Config: config.Root{
105+
Resources: config.Resources{
106+
Jobs: map[string]*resources.Job{
107+
"test_job": job,
108+
},
109+
},
110+
},
111+
}
112+
113+
runner := jobRunner{key: "test", bundle: b, job: job}
114+
115+
m := mocks.NewMockWorkspaceClient(t)
116+
b.SetWorkpaceClient(m.WorkspaceClient)
117+
118+
jobApi := m.GetMockJobsAPI()
119+
jobApi.EXPECT().ListRunsAll(mock.Anything, jobs.ListRunsRequest{
120+
ActiveOnly: true,
121+
JobId: 123,
122+
}).Return([]jobs.BaseRun{}, nil)
123+
124+
jobApi.AssertNotCalled(t, "CancelRun")
125+
126+
err := runner.Cancel(context.Background())
127+
require.NoError(t, err)
128+
}

‎bundle/run/pipeline.go‎

Lines changed: 15 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -166,3 +166,18 @@ func (r *pipelineRunner) Run(ctx context.Context, opts *Options) (output.RunOutp
166166
time.Sleep(time.Second)
167167
}
168168
}
169+
170+
func (r *pipelineRunner) Cancel(ctx context.Context) error {
171+
w := r.bundle.WorkspaceClient()
172+
wait, err := w.Pipelines.Stop(ctx, pipelines.StopRequest{
173+
PipelineId: r.pipeline.ID,
174+
})
175+
176+
if err != nil {
177+
return err
178+
}
179+
180+
// Waits for the Idle state of the pipeline
181+
_, err = wait.GetWithTimeout(jobRunTimeout)
182+
return err
183+
}

‎bundle/run/pipeline_test.go‎

Lines changed: 49 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,49 @@
1+
package run
2+
3+
import (
4+
"context"
5+
"testing"
6+
"time"
7+
8+
"github.com/databricks/cli/bundle"
9+
"github.com/databricks/cli/bundle/config"
10+
"github.com/databricks/cli/bundle/config/resources"
11+
"github.com/databricks/databricks-sdk-go/experimental/mocks"
12+
"github.com/databricks/databricks-sdk-go/service/pipelines"
13+
"github.com/stretchr/testify/require"
14+
)
15+
16+
func TestPipelineRunnerCancel(t *testing.T) {
17+
pipeline := &resources.Pipeline{
18+
ID: "123",
19+
}
20+
21+
b := &bundle.Bundle{
22+
Config: config.Root{
23+
Resources: config.Resources{
24+
Pipelines: map[string]*resources.Pipeline{
25+
"test_pipeline": pipeline,
26+
},
27+
},
28+
},
29+
}
30+
31+
runner := pipelineRunner{key: "test", bundle: b, pipeline: pipeline}
32+
33+
m := mocks.NewMockWorkspaceClient(t)
34+
b.SetWorkpaceClient(m.WorkspaceClient)
35+
36+
mockWait := &pipelines.WaitGetPipelineIdle[struct{}]{
37+
Poll: func(time.Duration, func(*pipelines.GetPipelineResponse)) (*pipelines.GetPipelineResponse, error) {
38+
return nil, nil
39+
},
40+
}
41+
42+
pipelineApi := m.GetMockPipelinesAPI()
43+
pipelineApi.EXPECT().Stop(context.Background(), pipelines.StopRequest{
44+
PipelineId: "123",
45+
}).Return(mockWait, nil)
46+
47+
err := runner.Cancel(context.Background())
48+
require.NoError(t, err)
49+
}

‎bundle/run/runner.go‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -26,6 +26,9 @@ type Runner interface {
2626

2727
// Run the underlying worklow.
2828
Run(ctx context.Context, opts *Options) (output.RunOutput, error)
29+
30+
// Cancel the underlying workflow.
31+
Cancel(ctx context.Context) error
2932
}
3033

3134
// Find locates a runner matching the specified argument.

‎cmd/bundle/run.go‎

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -27,7 +27,9 @@ func newRunCommand() *cobra.Command {
2727
runOptions.Define(cmd)
2828

2929
var noWait bool
30+
var restart bool
3031
cmd.Flags().BoolVar(&noWait, "no-wait", false, "Don't wait for the run to complete.")
32+
cmd.Flags().BoolVar(&restart, "restart", false, "Restart the run if it is already running.")
3133

3234
cmd.RunE = func(cmd *cobra.Command, args []string) error {
3335
ctx := cmd.Context()
@@ -68,6 +70,15 @@ func newRunCommand() *cobra.Command {
6870
}
6971

7072
runOptions.NoWait = noWait
73+
if restart {
74+
s := cmdio.Spinner(ctx)
75+
s <- "Cancelling all runs"
76+
err := runner.Cancel(ctx)
77+
close(s)
78+
if err != nil {
79+
return err
80+
}
81+
}
7182
output, err := runner.Run(ctx, &runOptions)
7283
if err != nil {
7384
return err

0 commit comments

Comments
 (0)