Added run_as section for bundle configuration

bundle/config/mutator/mutator.go

@@ -4,14 +4,16 @@ import (
 	"github.com/databricks/cli/bundle"
 )
 
+var defaultMutators []bundle.Mutator = []bundle.Mutator{
+	ProcessRootIncludes(),
+	DefineDefaultTarget(),
+	LoadGitDetails(),
+}
+
 func DefaultMutators() []bundle.Mutator {
-	return []bundle.Mutator{
-		ProcessRootIncludes(),
-		DefineDefaultTarget(),
-		LoadGitDetails(),
-	}
+	return append(defaultMutators, SetRunAs())
 }
 
 func DefaultMutatorsForTarget(env string) []bundle.Mutator {
-	return append(DefaultMutators(), SelectTarget(env))
+	return append(defaultMutators, SelectTarget(env), SetRunAs())
 }

bundle/config/mutator/run_as.go

@@ -0,0 +1,49 @@
+package mutator
+
+import (
+	"context"
+
+	"github.com/databricks/cli/bundle"
+	"github.com/databricks/cli/bundle/config/resources"
+	"github.com/databricks/databricks-sdk-go/service/jobs"
+)
+
+type setRunAs struct {
+}
+
+func SetRunAs() bundle.Mutator {
+	return &setRunAs{}
+}
+
+func (m *setRunAs) Name() string {
+	return "SetRunAs"
+}
+
+func (m *setRunAs) Apply(_ context.Context, b *bundle.Bundle) error {
+	runAs := b.Config.RunAs
+	if runAs == nil {
+		return nil
+	}
+
+	for i := range b.Config.Resources.Jobs {
+		job := b.Config.Resources.Jobs[i]
+		if job.RunAs != nil {
+			continue
+		}
+		job.RunAs = &jobs.JobRunAs{
+			ServicePrincipalName: runAs.ServicePrincipalName,
+			UserName:             runAs.UserName,
+		}
+	}
+
+	for i := range b.Config.Resources.Pipelines {
+		pipeline := b.Config.Resources.Pipelines[i]
+		pipeline.Permissions = append(pipeline.Permissions, resources.Permission{
+			Level:                "IS_OWNER",
+			ServicePrincipalName: runAs.ServicePrincipalName,
+			UserName:             runAs.UserName,
+		})
+	}
+
+	return nil
+}

bundle/config/root.go

@@ -7,6 +7,7 @@ import (
 	"strings"
 
 	"github.com/databricks/cli/bundle/config/variable"
+	"github.com/databricks/databricks-sdk-go/service/jobs"
 	"github.com/ghodss/yaml"
 	"github.com/imdario/mergo"
 )
@@ -80,6 +81,9 @@ type Root struct {
 
 	// Sync section specifies options for files synchronization
 	Sync Sync `json:"sync"`
+
+	// RunAs section allows to define an execution identity for jobs and pipelines runs
+	RunAs *jobs.JobRunAs `json:"run_as,omitempty"`
 }
 
 func Load(path string) (*Root, error) {
@@ -237,6 +241,11 @@ func (r *Root) MergeTargetOverrides(target *Target) error {
 		}
 	}
 
+	if target.RunAs != nil {
+		r.RunAs = target.RunAs
+		fmt.Println("Merging run as")
+	}
+
 	if target.Mode != "" {
 		r.Bundle.Mode = target.Mode
 	}

bundle/config/target.go

@@ -1,5 +1,7 @@
 package config
 
+import "github.com/databricks/databricks-sdk-go/service/jobs"
+
 type Mode string
 
 // Target defines overrides for a single target.
@@ -31,6 +33,8 @@ type Target struct {
 	Variables map[string]string `json:"variables,omitempty"`
 
 	Git Git `json:"git,omitempty"`
+
+	RunAs *jobs.JobRunAs `json:"run_as,omitempty"`
 }
 
 const (

bundle/tests/loader.go

@@ -9,18 +9,19 @@ import (
 	"github.com/stretchr/testify/require"
 )
 
-func load(t *testing.T, path string) *bundle.Bundle {
+func loadBundle(t *testing.T, path string, mutators []bundle.Mutator) *bundle.Bundle {
 	ctx := context.Background()
 	b, err := bundle.Load(ctx, path)
 	require.NoError(t, err)
-	err = bundle.Apply(ctx, b, bundle.Seq(mutator.DefaultMutators()...))
+	err = bundle.Apply(ctx, b, bundle.Seq(mutators...))
 	require.NoError(t, err)
 	return b
 }
 
+func load(t *testing.T, path string) *bundle.Bundle {
+	return loadBundle(t, path, mutator.DefaultMutators())
+}
+
 func loadTarget(t *testing.T, path, env string) *bundle.Bundle {
-	b := load(t, path)
-	err := bundle.Apply(context.Background(), b, mutator.SelectTarget(env))
-	require.NoError(t, err)
-	return b
+	return loadBundle(t, path, mutator.DefaultMutatorsForTarget(env))
 }

bundle/tests/run_as/databricks.yml

@@ -0,0 +1,37 @@
+bundle:
+  name: "run_as"
+
+run_as:
+  service_principal_name: "my_service_principal"
+
+targets:
+  development:
+    mode: development
+    run_as:
+      user_name: "my_user_name"
+
+resources:
+  pipelines:
+    nyc_taxi_pipeline:
+      name: "nyc taxi loader"
+      libraries:
+        - notebook:
+            path: ./dlt/nyc_taxi_loader
+  jobs:
+    job_one:
+      name: Job One
+      tasks:
+        - task:
+            notebook_path: "./test.py"
+    job_two:
+      name: Job Two
+      tasks:
+        - task:
+            notebook_path: "./test.py"
+    job_three:
+      name: Job Three
+      run_as:
+        service_principal_name: "my_service_principal_for_job"
+      tasks:
+        - task:
+            notebook_path: "./test.py"

bundle/tests/run_as_test.go

@@ -0,0 +1,54 @@
+package config_tests
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+)
+
+func TestRunAsDefault(t *testing.T) {
+	b := load(t, "./run_as")
+	assert.Len(t, b.Config.Resources.Jobs, 3)
+	jobs := b.Config.Resources.Jobs
+
+	assert.NotNil(t, jobs["job_one"].RunAs)
+	assert.Equal(t, "my_service_principal", jobs["job_one"].RunAs.ServicePrincipalName)
+	assert.Equal(t, "", jobs["job_one"].RunAs.UserName)
+
+	assert.NotNil(t, jobs["job_two"].RunAs)
+	assert.Equal(t, "my_service_principal", jobs["job_two"].RunAs.ServicePrincipalName)
+	assert.Equal(t, "", jobs["job_two"].RunAs.UserName)
+
+	assert.NotNil(t, jobs["job_three"].RunAs)
+	assert.Equal(t, "my_service_principal_for_job", jobs["job_three"].RunAs.ServicePrincipalName)
+	assert.Equal(t, "", jobs["job_three"].RunAs.UserName)
+
+	pipelines := b.Config.Resources.Pipelines
+	assert.NotNil(t, pipelines["nyc_taxi_pipeline"].Permissions)
+	assert.Equal(t, pipelines["nyc_taxi_pipeline"].Permissions[0].Level, "IS_OWNER")
+	assert.Equal(t, pipelines["nyc_taxi_pipeline"].Permissions[0].ServicePrincipalName, "my_service_principal")
+}
+
+func TestRunAsDevelopment(t *testing.T) {
+	b := loadTarget(t, "./run_as", "development")
+	assert.Len(t, b.Config.Resources.Jobs, 3)
+
+	jobs := b.Config.Resources.Jobs
+
+	assert.NotNil(t, jobs["job_one"].RunAs)
+	assert.Equal(t, "", jobs["job_one"].RunAs.ServicePrincipalName)
+	assert.Equal(t, "my_user_name", jobs["job_one"].RunAs.UserName)
+
+	assert.NotNil(t, jobs["job_two"].RunAs)
+	assert.Equal(t, "", jobs["job_two"].RunAs.ServicePrincipalName)
+	assert.Equal(t, "my_user_name", jobs["job_two"].RunAs.UserName)
+
+	assert.NotNil(t, jobs["job_three"].RunAs)
+	assert.Equal(t, "my_service_principal_for_job", jobs["job_three"].RunAs.ServicePrincipalName)
+	assert.Equal(t, "", jobs["job_three"].RunAs.UserName)
+
+	pipelines := b.Config.Resources.Pipelines
+	assert.NotNil(t, pipelines["nyc_taxi_pipeline"].Permissions)
+	assert.Equal(t, pipelines["nyc_taxi_pipeline"].Permissions[0].Level, "IS_OWNER")
+	assert.Equal(t, pipelines["nyc_taxi_pipeline"].Permissions[0].UserName, "my_user_name")
+}