Add support for "debug runs"

bundle/config/bundle.go

@@ -28,4 +28,6 @@ type Bundle struct {
 	// Contains Git information like current commit, current branch and
 	// origin url. Automatically loaded by reading .git directory if not specified
 	Git Git `json:"git,omitempty"`
+
+	Mode Mode `json:"mode,omitempty"`
 }

bundle/config/environment.go

@@ -1,12 +1,22 @@
 package config
 
+type Mode string
+
+const (
+	Debug Mode = "debug"
+	Default Mode = "default"
+	PullRequest Mode = "pull-request"
+)
+
 // Environment defines overrides for a single environment.
 // This structure is recursively merged into the root configuration.
 type Environment struct {
 	// Default marks that this environment must be used if one isn't specified
 	// by the user (through environment variable or command line argument).
 	Default bool `json:"default,omitempty"`
 
+	Mode Mode `json:"mode,omitempty"`
+
 	Bundle *Bundle `json:"bundle,omitempty"`
 
 	Workspace *Workspace `json:"workspace,omitempty"`

bundle/config/mutator/override_compute.go

@@ -0,0 +1,50 @@
+package mutator
+
+import (
+	"context"
+	"fmt"
+
+	"github.com/databricks/cli/bundle"
+	"github.com/databricks/cli/bundle/config"
+	"github.com/databricks/cli/bundle/config/resources"
+)
+
+type overrideCompute struct {
+	compute string
+}
+
+func OverrideCompute(compute string) bundle.Mutator {
+	return &overrideCompute{compute: compute}
+}
+
+func (m *overrideCompute) Name() string {
+	return "OverrideCompute"
+}
+
+func (m *overrideCompute) overrideJobCompute(j *resources.Job) {
+	for i := range j.Tasks {
+ task := &j.Tasks[i]
+ if task.NewCluster != nil {
+ task.NewCluster = nil
+ task.ExistingClusterId = m.compute
+ } else if task.ExistingClusterId != "" {
+ task.ExistingClusterId = m.compute
+ }
+	}
+}
+
+func (m *overrideCompute) Apply(ctx context.Context, b *bundle.Bundle) error {
+	if m.compute == "" {
+ return nil
+	}
+	if b.Config.Bundle.Mode != config.Debug {
+ return fmt.Errorf("cannot override compute for an environment that does not use 'mode: debug'")
+	}
+
+	r := b.Config.Resources
+	for i := range r.Jobs {
+ m.overrideJobCompute(r.Jobs[i])
+	}
+
+	return nil
+}

bundle/config/mutator/override_compute_test.go

@@ -0,0 +1,47 @@
+package mutator_test
+
+import (
+	"context"
+	"testing"
+
+	"github.com/databricks/cli/bundle"
+	"github.com/databricks/cli/bundle/config"
+	"github.com/databricks/cli/bundle/config/mutator"
+	"github.com/databricks/cli/bundle/config/resources"
+	"github.com/databricks/databricks-sdk-go/service/compute"
+	"github.com/databricks/databricks-sdk-go/service/jobs"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+)
+
+func TestOverrideCompute(t *testing.T) {
+	bundle := &bundle.Bundle{
+ Config: config.Root{
+ Bundle: config.Bundle{
+ Mode: config.Debug,
+ },
+ Resources: config.Resources{
+ Jobs: map[string]*resources.Job{
+ "job1": {JobSettings: &jobs.JobSettings{
+ Name: "job1",
+ Tasks: []jobs.JobTaskSettings{
+ {
+ NewCluster: &compute.BaseClusterInfo{},
+ },
+ {
+ ExistingClusterId: "cluster2",
+ },
+ },
+ }},
+ },
+ },
+ },
+	}
+
+	m := mutator.OverrideCompute("newClusterID")
+	err := m.Apply(context.Background(), bundle)
+	require.NoError(t, err)
+	assert.Nil(t, bundle.Config.Resources.Jobs["job1"].Tasks[0].NewCluster)
+	assert.Equal(t, "newClusterID", bundle.Config.Resources.Jobs["job1"].Tasks[0].ExistingClusterId)
+	assert.Equal(t, "newClusterID", bundle.Config.Resources.Jobs["job1"].Tasks[1].ExistingClusterId)
+}

bundle/config/mutator/populate_current_user.go

@@ -18,6 +18,10 @@ func (m *populateCurrentUser) Name() string {
 }
 
 func (m *populateCurrentUser) Apply(ctx context.Context, b *bundle.Bundle) error {
+	if b.Config.Workspace.CurrentUser != nil {
+ return nil
+	}
+
 	w := b.WorkspaceClient()
 	me, err := w.CurrentUser.Me(ctx)
 	if err != nil {

bundle/config/mutator/process_environment_mode.go

@@ -0,0 +1,82 @@
+package mutator
+
+import (
+	"context"
+	"fmt"
+
+	"github.com/databricks/cli/bundle"
+	"github.com/databricks/cli/bundle/config"
+	"github.com/databricks/databricks-sdk-go/service/ml"
+)
+
+type processEnvironmentMode struct{}
+
+const debugConcurrentRuns = 4
+
+func ProcessEnvironmentMode() bundle.Mutator {
+	return &processEnvironmentMode{}
+}
+
+func (m *processEnvironmentMode) Name() string {
+	return "ProcessEnvironmentMode"
+}
+
+// Mark all resources as being for 'debug' purposes, i.e.
+// changing their their name, adding tags, and (in the future)
+// marking them as 'hidden' in the UI.
+func processDebugMode(b *bundle.Bundle) error {
+	r := b.Config.Resources
+
+	for i := range r.Jobs {
+ r.Jobs[i].Name = "[debug] " + r.Jobs[i].Name
+ if r.Jobs[i].Tags == nil {
+ r.Jobs[i].Tags = make(map[string]string)
+ }
+ r.Jobs[i].Tags["debug"] = ""
+ if r.Jobs[i].MaxConcurrentRuns == 0 {
+ r.Jobs[i].MaxConcurrentRuns = debugConcurrentRuns
+ }
+ if r.Jobs[i].Schedule != nil {
+ r.Jobs[i].Schedule.PauseStatus = "PAUSED"
+ }
+ if r.Jobs[i].Continuous != nil {
+ r.Jobs[i].Continuous.PauseStatus = "PAUSED"
+ }
+ if r.Jobs[i].Trigger != nil {
+ r.Jobs[i].Trigger.PauseStatus = "PAUSED"
+ }
+	}
+
+	for i := range r.Pipelines {
+ r.Pipelines[i].Name = "[debug] " + r.Pipelines[i].Name
+ r.Pipelines[i].Development = true
+ // (pipelines don't yet support tags)
+	}
+
+	for i := range r.Models {
+ r.Models[i].Name = "[debug] " + r.Models[i].Name
+ r.Models[i].Tags = append(r.Models[i].Tags, ml.ModelTag{Key: "debug", Value: ""})
+	}
+
+	for i := range r.Experiments {
+ r.Experiments[i].Name = "[debug] " + r.Experiments[i].Name
+ r.Experiments[i].Tags = append(r.Experiments[i].Tags, ml.ExperimentTag{Key: "debug", Value: ""})
+	}
+
+	return nil
+}
+
+func (m *processEnvironmentMode) Apply(ctx context.Context, b *bundle.Bundle) error {
+	switch b.Config.Bundle.Mode {
+	case config.Debug:
+ return processDebugMode(b)
+	case config.Default, "":
+ // No action
+	case config.PullRequest:
+ return fmt.Errorf("not implemented")
+	default:
+ return fmt.Errorf("unsupported value specified for 'mode': %s", b.Config.Bundle.Mode)
+	}
+
+	return nil
+}

bundle/config/mutator/process_environment_mode_test.go

@@ -0,0 +1,75 @@
+package mutator_test
+
+import (
+	"context"
+	"testing"
+
+	"github.com/databricks/cli/bundle"
+	"github.com/databricks/cli/bundle/config"
+	"github.com/databricks/cli/bundle/config/mutator"
+	"github.com/databricks/cli/bundle/config/resources"
+	"github.com/databricks/databricks-sdk-go/service/jobs"
+	"github.com/databricks/databricks-sdk-go/service/ml"
+	"github.com/databricks/databricks-sdk-go/service/pipelines"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+)
+
+func TestProcessEnvironmentModeApplyDebug(t *testing.T) {
+	bundle := &bundle.Bundle{
+ Config: config.Root{
+ Bundle: config.Bundle{
+ Mode: config.Debug,
+ },
+ Resources: config.Resources{
+ Jobs: map[string]*resources.Job{
+ "job1": {JobSettings: &jobs.JobSettings{Name: "job1"}},
+ },
+ Pipelines: map[string]*resources.Pipeline{
+ "pipeline1": {PipelineSpec: &pipelines.PipelineSpec{Name: "pipeline1"}},
+ },
+ Experiments: map[string]*resources.MlflowExperiment{
+ "experiment1": {Experiment: &ml.Experiment{Name: "experiment1"}},
+ },
+ Models: map[string]*resources.MlflowModel{
+ "model1": {Model: &ml.Model{Name: "model1"}},
+ },
+ },
+ },
+	}
+
+	m := mutator.ProcessEnvironmentMode()
+	err := m.Apply(context.Background(), bundle)
+	require.NoError(t, err)
+	assert.Equal(t, "[debug] job1", bundle.Config.Resources.Jobs["job1"].Name)
+	assert.Equal(t, "[debug] pipeline1", bundle.Config.Resources.Pipelines["pipeline1"].Name)
+	assert.Equal(t, "[debug] experiment1", bundle.Config.Resources.Experiments["experiment1"].Name)
+	assert.Equal(t, "[debug] model1", bundle.Config.Resources.Models["model1"].Name)
+	assert.Equal(t, "debug", bundle.Config.Resources.Experiments["experiment1"].Experiment.Tags[0].Key)
+	assert.True(t, bundle.Config.Resources.Pipelines["pipeline1"].PipelineSpec.Development)
+}
+
+func TestProcessEnvironmentModeApplyDefault(t *testing.T) {
+	bundle := &bundle.Bundle{
+ Config: config.Root{
+ Bundle: config.Bundle{
+ Mode: config.Default,
+ },
+ Resources: config.Resources{
+ Jobs: map[string]*resources.Job{
+ "job1": {JobSettings: &jobs.JobSettings{Name: "job1"}},
+ },
+ Pipelines: map[string]*resources.Pipeline{
+ "pipeline1": {PipelineSpec: &pipelines.PipelineSpec{Name: "pipeline1"}},
+ },
+ },
+ },
+	}
+
+	m := mutator.ProcessEnvironmentMode()
+	err := m.Apply(context.Background(), bundle)
+	require.NoError(t, err)
+	assert.Equal(t, "job1", bundle.Config.Resources.Jobs["job1"].Name)
+	assert.Equal(t, "pipeline1", bundle.Config.Resources.Pipelines["pipeline1"].Name)
+	assert.False(t, bundle.Config.Resources.Pipelines["pipeline1"].PipelineSpec.Development)
+}

bundle/config/root.go

@@ -190,5 +190,9 @@ func (r *Root) MergeEnvironment(env *Environment) error {
  }
 	}
 
+	if env.Mode != "" {
+ r.Bundle.Mode = env.Mode
+	}
+
 	return nil
 }

bundle/config/root_test.go

@@ -154,3 +154,12 @@ func TestInitializeVariablesUndefinedVariables(t *testing.T) {
 	err := root.InitializeVariables([]string{"bar=567"})
 	assert.ErrorContains(t, err, "variable bar has not been defined")
 }
+
+func TestRootMergeEnvironmentWithMode(t *testing.T) {
+	root := &Root{
+ Bundle: Bundle{},
+	}
+	env := &Environment{Mode: Debug}
+	require.NoError(t, root.MergeEnvironment(env))
+	assert.Equal(t, Debug, root.Bundle.Mode)
+}

bundle/phases/initialize.go

@@ -11,7 +11,7 @@ import (
 // The initialize phase fills in defaults and connects to the workspace.
 // Interpolation of fields referring to the "bundle" and "workspace" keys
 // happens upon completion of this phase.
-func Initialize() bundle.Mutator {
+func Initialize(overrideCompute string) bundle.Mutator {
 	return newPhase(
  "initialize",
  []bundle.Mutator{
@@ -25,6 +25,8 @@ func Initialize() bundle.Mutator {
  interpolation.IncludeLookupsInPath("workspace"),
  interpolation.IncludeLookupsInPath(variable.VariableReferencePrefix),
  ),
+ mutator.OverrideCompute(overrideCompute),
+ mutator.ProcessEnvironmentMode(),
  mutator.TranslatePaths(),
  terraform.Initialize(),
  },

bundle/run/job.go

@@ -243,6 +243,15 @@ func (r *jobRunner) Run(ctx context.Context, opts *Options) (output.RunOutput, e
 	if err != nil {
  return nil, fmt.Errorf("cannot start job")
 	}
+
+	if opts.NoWait {
+ details, err := w.Jobs.GetRun(ctx, jobs.GetRunRequest{
+ RunId: waiter.RunId,
+ })
+ progressLogger.Log(progress.NewJobRunUrlEvent(details.RunPageUrl))
+ return nil, err
+	}
+
 	run, err := waiter.OnProgress(func(r *jobs.Run) {
  pullRunId(r)
  logDebug(r)

bundle/run/options.go

@@ -7,6 +7,7 @@ import (
 type Options struct {
 	Job JobOptions
 	Pipeline PipelineOptions
+	NoWait bool
 }
 
 func (o *Options) Define(fs *flag.FlagSet) {

bundle/run/pipeline.go

@@ -167,6 +167,10 @@ func (r *pipelineRunner) Run(ctx context.Context, opts *Options) (output.RunOutp
  return nil, fmt.Errorf("no progress logger found")
 	}
 
+	if opts.NoWait {
+ log.Warnf(ctx, "--no-wait is not yet implemented for pipelines")
+	}
+
 	// Log the pipeline update URL as soon as it is available.
 	progressLogger.Log(progress.NewPipelineUpdateUrlEvent(w.Config.Host, updateID, pipelineID))
 

bundle/tests/job_and_pipeline/bundle.yml

@@ -8,6 +8,7 @@ resources:
 
 environments:
  development:
+ mode: debug
  resources:
  pipelines:
  nyc_taxi_pipeline:

bundle/tests/job_and_pipeline_test.go

@@ -4,6 +4,7 @@ import (
 	"path/filepath"
 	"testing"
 
+	"github.com/databricks/cli/bundle/config"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 )
@@ -15,6 +16,7 @@ func TestJobAndPipelineDevelopment(t *testing.T) {
 
 	p := b.Config.Resources.Pipelines["nyc_taxi_pipeline"]
 	assert.Equal(t, "job_and_pipeline/bundle.yml", filepath.ToSlash(p.ConfigFilePath))
+	assert.Equal(t, b.Config.Bundle.Mode, config.Debug)
 	assert.True(t, p.Development)
 	require.Len(t, p.Libraries, 1)
 	assert.Equal(t, "./dlt/nyc_taxi_loader", p.Libraries[0].Notebook.Path)