Add allow list for resources when bundle `run_as` is set (#1233)

## Changes
This PR introduces an allow list for resource types that are allowed
when the run_as for the bundle is not the same as the current deployment
user.

This PR also adds a test to ensure that any new resources added to DABs
will have to add the resource to either the allow list or add an error
to fail when run_as identity is not the same as deployment user.

## Tests
Unit tests

bundle/config/mutator/run_as.go

@@ -2,20 +2,24 @@ package mutator
 
 import (
 	"context"
-	"slices"
+	"fmt"
 
 	"github.com/databricks/cli/bundle"
-	"github.com/databricks/cli/bundle/config/resources"
 	"github.com/databricks/cli/libs/diag"
+	"github.com/databricks/cli/libs/dyn"
 	"github.com/databricks/databricks-sdk-go/service/jobs"
 )
 
 type setRunAs struct {
 }
 
-// SetRunAs mutator is used to go over defined resources such as Jobs and DLT Pipelines
-// And set correct execution identity ("run_as" for a job or "is_owner" permission for DLT)
-// if top-level "run-as" section is defined in the configuration.
+// This mutator does two things:
+//
+//  1. Sets the run_as field for jobs to the value of the run_as field in the bundle.
+//
+//  2. Validates that the bundle run_as configuration is valid in the context of the bundle.
+//     If the run_as user is different from the current deployment user, DABs only
+//     supports a subset of resources.
 func SetRunAs() bundle.Mutator {
 	return &setRunAs{}
 }
@@ -24,12 +28,94 @@ func (m *setRunAs) Name() string {
 	return "SetRunAs"
 }
 
+type errUnsupportedResourceTypeForRunAs struct {
+	resourceType     string
+	resourceLocation dyn.Location
+	currentUser      string
+	runAsUser        string
+}
+
+// TODO(6 March 2024): Link the docs page describing run_as semantics in the error below
+// once the page is ready.
+func (e errUnsupportedResourceTypeForRunAs) Error() string {
+	return fmt.Sprintf("%s are not supported when the current deployment user is different from the bundle's run_as identity. Please deploy as the run_as identity. Location of the unsupported resource: %s. Current identity: %s. Run as identity: %s", e.resourceType, e.resourceLocation, e.currentUser, e.runAsUser)
+}
+
+type errBothSpAndUserSpecified struct {
+	spName   string
+	spLoc    dyn.Location
+	userName string
+	userLoc  dyn.Location
+}
+
+func (e errBothSpAndUserSpecified) Error() string {
+	return fmt.Sprintf("run_as section must specify exactly one identity. A service_principal_name %q is specified at %s. A user_name %q is defined at %s", e.spName, e.spLoc, e.userName, e.userLoc)
+}
+
+func validateRunAs(b *bundle.Bundle) error {
+	runAs := b.Config.RunAs
+
+	// Error if neither service_principal_name nor user_name are specified
+	if runAs.ServicePrincipalName == "" && runAs.UserName == "" {
+		return fmt.Errorf("run_as section must specify exactly one identity. Neither service_principal_name nor user_name is specified at %s", b.Config.GetLocation("run_as"))
+	}
+
+	// Error if both service_principal_name and user_name are specified
+	if runAs.UserName != "" && runAs.ServicePrincipalName != "" {
+		return errBothSpAndUserSpecified{
+			spName:   runAs.ServicePrincipalName,
+			userName: runAs.UserName,
+			spLoc:    b.Config.GetLocation("run_as.service_principal_name"),
+			userLoc:  b.Config.GetLocation("run_as.user_name"),
+		}
+	}
+
+	identity := runAs.ServicePrincipalName
+	if identity == "" {
+		identity = runAs.UserName
+	}
+
+	// All resources are supported if the run_as identity is the same as the current deployment identity.
+	if identity == b.Config.Workspace.CurrentUser.UserName {
+		return nil
+	}
+
+	// DLT pipelines do not support run_as in the API.
+	if len(b.Config.Resources.Pipelines) > 0 {
+		return errUnsupportedResourceTypeForRunAs{
+			resourceType:     "pipelines",
+			resourceLocation: b.Config.GetLocation("resources.pipelines"),
+			currentUser:      b.Config.Workspace.CurrentUser.UserName,
+			runAsUser:        identity,
+		}
+	}
+
+	// Model serving endpoints do not support run_as in the API.
+	if len(b.Config.Resources.ModelServingEndpoints) > 0 {
+		return errUnsupportedResourceTypeForRunAs{
+			resourceType:     "model_serving_endpoints",
+			resourceLocation: b.Config.GetLocation("resources.model_serving_endpoints"),
+			currentUser:      b.Config.Workspace.CurrentUser.UserName,
+			runAsUser:        identity,
+		}
+	}
+
+	return nil
+}
+
 func (m *setRunAs) Apply(_ context.Context, b *bundle.Bundle) diag.Diagnostics {
+	// Mutator is a no-op if run_as is not specified in the bundle
 	runAs := b.Config.RunAs
 	if runAs == nil {
 		return nil
 	}
 
+	// Assert the run_as configuration is valid in the context of the bundle
+	if err := validateRunAs(b); err != nil {
+		return diag.FromErr(err)
+	}
+
+	// Set run_as for jobs
 	for i := range b.Config.Resources.Jobs {
 		job := b.Config.Resources.Jobs[i]
 		if job.RunAs != nil {
@@ -41,26 +127,5 @@ func (m *setRunAs) Apply(_ context.Context, b *bundle.Bundle) diag.Diagnostics {
 		}
 	}
 
-	me := b.Config.Workspace.CurrentUser.UserName
-	// If user deploying the bundle and the one defined in run_as are the same
-	// Do not add IS_OWNER permission. Current user is implied to be an owner in this case.
-	// Otherwise, it will fail due to this bug https://github.com/databricks/terraform-provider-databricks/issues/2407
-	if runAs.UserName == me || runAs.ServicePrincipalName == me {
-		return nil
-	}
-
-	for i := range b.Config.Resources.Pipelines {
-		pipeline := b.Config.Resources.Pipelines[i]
-		pipeline.Permissions = slices.DeleteFunc(pipeline.Permissions, func(p resources.Permission) bool {
-			return (runAs.ServicePrincipalName != "" && p.ServicePrincipalName == runAs.ServicePrincipalName) ||
-				(runAs.UserName != "" && p.UserName == runAs.UserName)
-		})
-		pipeline.Permissions = append(pipeline.Permissions, resources.Permission{
-			Level:                "IS_OWNER",
-			ServicePrincipalName: runAs.ServicePrincipalName,
-			UserName:             runAs.UserName,
-		})
-	}
-
 	return nil
 }

bundle/config/mutator/run_as_test.go

@@ -0,0 +1,188 @@
+package mutator
+
+import (
+	"context"
+	"slices"
+	"testing"
+
+	"github.com/databricks/cli/bundle"
+	"github.com/databricks/cli/bundle/config"
+	"github.com/databricks/cli/bundle/config/resources"
+	"github.com/databricks/cli/libs/dyn"
+	"github.com/databricks/cli/libs/dyn/convert"
+	"github.com/databricks/databricks-sdk-go/service/iam"
+	"github.com/databricks/databricks-sdk-go/service/jobs"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+)
+
+func allResourceTypes(t *testing.T) []string {
+	// Compute supported resource types based on the `Resources{}` struct.
+	r := config.Resources{}
+	rv, err := convert.FromTyped(r, dyn.NilValue)
+	require.NoError(t, err)
+	normalized, _ := convert.Normalize(r, rv, convert.IncludeMissingFields)
+	resourceTypes := []string{}
+	for _, k := range normalized.MustMap().Keys() {
+		resourceTypes = append(resourceTypes, k.MustString())
+	}
+	slices.Sort(resourceTypes)
+
+	// Assert the total list of resource supported, as a sanity check that using
+	// the dyn library gives us the correct list of all resources supported. Please
+	// also update this check when adding a new resource
+	require.Equal(t, []string{
+		"experiments",
+		"jobs",
+		"model_serving_endpoints",
+		"models",
+		"pipelines",
+		"registered_models",
+	},
+		resourceTypes,
+	)
+
+	return resourceTypes
+}
+
+func TestRunAsWorksForAllowedResources(t *testing.T) {
+	config := config.Root{
+		Workspace: config.Workspace{
+			CurrentUser: &config.User{
+				User: &iam.User{
+					UserName: "alice",
+				},
+			},
+		},
+		RunAs: &jobs.JobRunAs{
+			UserName: "bob",
+		},
+		Resources: config.Resources{
+			Jobs: map[string]*resources.Job{
+				"job_one": {
+					JobSettings: &jobs.JobSettings{
+						Name: "foo",
+					},
+				},
+				"job_two": {
+					JobSettings: &jobs.JobSettings{
+						Name: "bar",
+					},
+				},
+				"job_three": {
+					JobSettings: &jobs.JobSettings{
+						Name: "baz",
+					},
+				},
+			},
+			Models: map[string]*resources.MlflowModel{
+				"model_one": {},
+			},
+			RegisteredModels: map[string]*resources.RegisteredModel{
+				"registered_model_one": {},
+			},
+			Experiments: map[string]*resources.MlflowExperiment{
+				"experiment_one": {},
+			},
+		},
+	}
+
+	b := &bundle.Bundle{
+		Config: config,
+	}
+
+	diags := bundle.Apply(context.Background(), b, SetRunAs())
+	assert.NoError(t, diags.Error())
+
+	for _, job := range b.Config.Resources.Jobs {
+		assert.Equal(t, "bob", job.RunAs.UserName)
+	}
+}
+
+func TestRunAsErrorForUnsupportedResources(t *testing.T) {
+	// Bundle "run_as" has two modes of operation, each with a different set of
+	// resources that are supported.
+	// Cases:
+	//   1. When the bundle "run_as" identity is same as the current deployment
+	//      identity. In this case all resources are supported.
+	//   2. When the bundle "run_as" identity is different from the current
+	//      deployment identity. In this case only a subset of resources are
+	//      supported. This subset of resources are defined in the allow list below.
+	//
+	// To be a part of the allow list, the resource must satisfy one of the following
+	// two conditions:
+	//   1. The resource supports setting a run_as identity to a different user
+	//      from the owner/creator of the resource. For example, jobs.
+	//   2. Run as semantics do not apply to the resource. We do not plan to add
+	//      platform side support for `run_as` for these resources. For example,
+	//      experiments or registered models.
+	//
+	// Any resource that is not on the allow list cannot be used when the bundle
+	// run_as is different from the current deployment user. "bundle validate" must
+	// return an error if such a resource has been defined, and the run_as identity
+	// is different from the current deployment identity.
+	//
+	// Action Item: If you are adding a new resource to DABs, please check in with
+	// the relevant owning team whether the resource should be on the allow list or (implicitly) on
+	// the deny list. Any resources that could have run_as semantics in the future
+	// should be on the deny list.
+	// For example: Teams for pipelines, model serving endpoints or Lakeview dashboards
+	// are planning to add platform side support for `run_as` for these resources at
+	// some point in the future. These resources are (implicitly) on the deny list, since
+	// they are not on the allow list below.
+	allowList := []string{
+		"jobs",
+		"models",
+		"registered_models",
+		"experiments",
+	}
+
+	base := config.Root{
+		Workspace: config.Workspace{
+			CurrentUser: &config.User{
+				User: &iam.User{
+					UserName: "alice",
+				},
+			},
+		},
+		RunAs: &jobs.JobRunAs{
+			UserName: "bob",
+		},
+	}
+
+	v, err := convert.FromTyped(base, dyn.NilValue)
+	require.NoError(t, err)
+
+	for _, rt := range allResourceTypes(t) {
+		// Skip allowed resources
+		if slices.Contains(allowList, rt) {
+			continue
+		}
+
+		// Add an instance of the resource type that is not on the allow list to
+		// the bundle configuration.
+		nv, err := dyn.SetByPath(v, dyn.NewPath(dyn.Key("resources"), dyn.Key(rt)), dyn.V(map[string]dyn.Value{
+			"foo": dyn.V(map[string]dyn.Value{
+				"path": dyn.V("bar"),
+			}),
+		}))
+		require.NoError(t, err)
+
+		// Get back typed configuration from the newly created invalid bundle configuration.
+		r := &config.Root{}
+		err = convert.ToTyped(r, nv)
+		require.NoError(t, err)
+
+		// Assert this invalid bundle configuration fails validation.
+		b := &bundle.Bundle{
+			Config: *r,
+		}
+		diags := bundle.Apply(context.Background(), b, SetRunAs())
+		assert.Equal(t, diags.Error().Error(), errUnsupportedResourceTypeForRunAs{
+			resourceType:     rt,
+			resourceLocation: dyn.Location{},
+			currentUser:      "alice",
+			runAsUser:        "bob",
+		}.Error(), "expected run_as with a different identity than the current deployment user to not supported for resources of type: %s", rt)
+	}
+}

bundle/config/root.go

@@ -448,3 +448,14 @@ func validateVariableOverrides(root, target dyn.Value) (err error) {
 
 	return nil
 }
+
+// Best effort to get the location of configuration value at the specified path.
+// This function is useful to annotate error messages with the location, because
+// we don't want to fail with a different error message if we cannot retrieve the location.
+func (r *Root) GetLocation(path string) dyn.Location {
+	v, err := dyn.Get(r.value, path)
+	if err != nil {
+		return dyn.Location{}
+	}
+	return v.Location()
+}

bundle/tests/run_as/allowed/databricks.yml

@@ -11,20 +11,6 @@ targets:
       user_name: "my_user_name"
 
 resources:
-  pipelines:
-    nyc_taxi_pipeline:
-      name: "nyc taxi loader"
-
-      permissions:
-        - level: CAN_VIEW
-          service_principal_name: my_service_principal
-        - level: CAN_VIEW
-          user_name: my_user_name
-
-      libraries:
-        - notebook:
-            path: ./dlt/nyc_taxi_loader
-
   jobs:
     job_one:
       name: Job One
@@ -52,3 +38,15 @@ resources:
         - task_key: "task_three"
           notebook_task:
             notebook_path: "./test.py"
+
+  models:
+    model_one:
+      name: "skynet"
+
+  registered_models:
+    model_two:
+      name: "skynet (in UC)"
+
+  experiments:
+    experiment_one:
+      name: "experiment_one"

bundle/tests/run_as/not_allowed/both_sp_and_user/databricks.yml

@@ -0,0 +1,17 @@
+bundle:
+  name: "run_as"
+
+# This is not allowed because both service_principal_name and user_name are set
+run_as:
+  service_principal_name: "my_service_principal"
+  user_name: "my_user_name"
+
+resources:
+  jobs:
+    job_one:
+      name: Job One
+
+      tasks:
+        - task_key: "task_one"
+          notebook_task:
+            notebook_path: "./test.py"

bundle/tests/run_as/not_allowed/model_serving/databricks.yml

@@ -0,0 +1,15 @@
+bundle:
+  name: "run_as"
+
+run_as:
+  service_principal_name: "my_service_principal"
+
+targets:
+  development:
+    run_as:
+      user_name: "my_user_name"
+
+resources:
+  model_serving_endpoints:
+    foo:
+      name: "skynet"

bundle/tests/run_as/not_allowed/neither_sp_nor_user/databricks.yml

@@ -0,0 +1,4 @@
+bundle:
+  name: "abc"
+
+run_as:

bundle/tests/run_as/not_allowed/neither_sp_nor_user_override/databricks.yml

@@ -0,0 +1,8 @@
+bundle:
+  name: "abc"
+
+run_as:
+  user_name: "my_user_name"
+
+include:
+  - ./override.yml

bundle/tests/run_as/not_allowed/neither_sp_nor_user_override/override.yml

@@ -0,0 +1,4 @@
+targets:
+  development:
+    default: true
+    run_as:

bundle/tests/run_as/not_allowed/pipelines/databricks.yml

@@ -0,0 +1,25 @@
+bundle:
+  name: "run_as"
+
+run_as:
+  service_principal_name: "my_service_principal"
+
+targets:
+  development:
+    run_as:
+      user_name: "my_user_name"
+
+resources:
+  pipelines:
+    nyc_taxi_pipeline:
+      name: "nyc taxi loader"
+
+      permissions:
+        - level: CAN_VIEW
+          service_principal_name: my_service_principal
+        - level: CAN_VIEW
+          user_name: my_user_name
+
+      libraries:
+        - notebook:
+            path: ./dlt/nyc_taxi_loader

bundle/tests/run_as_test.go

@@ -2,18 +2,22 @@ package config_tests
 
 import (
 	"context"
+	"fmt"
+	"path/filepath"
 	"testing"
 
 	"github.com/databricks/cli/bundle"
 	"github.com/databricks/cli/bundle/config"
 	"github.com/databricks/cli/bundle/config/mutator"
 	"github.com/databricks/cli/libs/diag"
+	"github.com/databricks/databricks-sdk-go/service/catalog"
 	"github.com/databricks/databricks-sdk-go/service/iam"
+	"github.com/databricks/databricks-sdk-go/service/ml"
 	"github.com/stretchr/testify/assert"
 )
 
-func TestRunAsDefault(t *testing.T) {
-	b := load(t, "./run_as")
+func TestRunAsForAllowed(t *testing.T) {
+	b := load(t, "./run_as/allowed")
 
 	ctx := context.Background()
 	bundle.ApplyFunc(ctx, b, func(ctx context.Context, b *bundle.Bundle) diag.Diagnostics {
@@ -31,6 +35,7 @@ func TestRunAsDefault(t *testing.T) {
 	assert.Len(t, b.Config.Resources.Jobs, 3)
 	jobs := b.Config.Resources.Jobs
 
+	// job_one and job_two should have the same run_as identity as the bundle.
 	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)
@@ -39,21 +44,19 @@ func TestRunAsDefault(t *testing.T) {
 	assert.Equal(t, "my_service_principal", jobs["job_two"].RunAs.ServicePrincipalName)
 	assert.Equal(t, "", jobs["job_two"].RunAs.UserName)
 
+	// job_three should retain the job level run_as identity.
 	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.Len(t, pipelines["nyc_taxi_pipeline"].Permissions, 2)
-	assert.Equal(t, "CAN_VIEW", pipelines["nyc_taxi_pipeline"].Permissions[0].Level)
-	assert.Equal(t, "my_user_name", pipelines["nyc_taxi_pipeline"].Permissions[0].UserName)
-
-	assert.Equal(t, "IS_OWNER", pipelines["nyc_taxi_pipeline"].Permissions[1].Level)
-	assert.Equal(t, "my_service_principal", pipelines["nyc_taxi_pipeline"].Permissions[1].ServicePrincipalName)
+	// Assert other resources are not affected.
+	assert.Equal(t, ml.Model{Name: "skynet"}, *b.Config.Resources.Models["model_one"].Model)
+	assert.Equal(t, catalog.CreateRegisteredModelRequest{Name: "skynet (in UC)"}, *b.Config.Resources.RegisteredModels["model_two"].CreateRegisteredModelRequest)
+	assert.Equal(t, ml.Experiment{Name: "experiment_one"}, *b.Config.Resources.Experiments["experiment_one"].Experiment)
 }
 
-func TestRunAsDevelopment(t *testing.T) {
-	b := loadTarget(t, "./run_as", "development")
+func TestRunAsForAllowedWithTargetOverride(t *testing.T) {
+	b := loadTarget(t, "./run_as/allowed", "development")
 
 	ctx := context.Background()
 	bundle.ApplyFunc(ctx, b, func(ctx context.Context, b *bundle.Bundle) diag.Diagnostics {
@@ -71,6 +74,8 @@ func TestRunAsDevelopment(t *testing.T) {
 	assert.Len(t, b.Config.Resources.Jobs, 3)
 	jobs := b.Config.Resources.Jobs
 
+	// job_one and job_two should have the same run_as identity as the bundle's
+	// development target.
 	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)
@@ -79,15 +84,152 @@ func TestRunAsDevelopment(t *testing.T) {
 	assert.Equal(t, "", jobs["job_two"].RunAs.ServicePrincipalName)
 	assert.Equal(t, "my_user_name", jobs["job_two"].RunAs.UserName)
 
+	// job_three should retain the job level run_as identity.
 	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.Len(t, pipelines["nyc_taxi_pipeline"].Permissions, 2)
-	assert.Equal(t, "CAN_VIEW", pipelines["nyc_taxi_pipeline"].Permissions[0].Level)
-	assert.Equal(t, "my_service_principal", pipelines["nyc_taxi_pipeline"].Permissions[0].ServicePrincipalName)
+	// Assert other resources are not affected.
+	assert.Equal(t, ml.Model{Name: "skynet"}, *b.Config.Resources.Models["model_one"].Model)
+	assert.Equal(t, catalog.CreateRegisteredModelRequest{Name: "skynet (in UC)"}, *b.Config.Resources.RegisteredModels["model_two"].CreateRegisteredModelRequest)
+	assert.Equal(t, ml.Experiment{Name: "experiment_one"}, *b.Config.Resources.Experiments["experiment_one"].Experiment)
 
-	assert.Equal(t, "IS_OWNER", pipelines["nyc_taxi_pipeline"].Permissions[1].Level)
-	assert.Equal(t, "my_user_name", pipelines["nyc_taxi_pipeline"].Permissions[1].UserName)
+}
+
+func TestRunAsErrorForPipelines(t *testing.T) {
+	b := load(t, "./run_as/not_allowed/pipelines")
+
+	ctx := context.Background()
+	bundle.ApplyFunc(ctx, b, func(ctx context.Context, b *bundle.Bundle) diag.Diagnostics {
+		b.Config.Workspace.CurrentUser = &config.User{
+			User: &iam.User{
+				UserName: "jane@doe.com",
+			},
+		}
+		return nil
+	})
+
+	diags := bundle.Apply(ctx, b, mutator.SetRunAs())
+	err := diags.Error()
+
+	configPath := filepath.FromSlash("run_as/not_allowed/pipelines/databricks.yml")
+	assert.EqualError(t, err, fmt.Sprintf("pipelines are not supported when the current deployment user is different from the bundle's run_as identity. Please deploy as the run_as identity. Location of the unsupported resource: %s:14:5. Current identity: jane@doe.com. Run as identity: my_service_principal", configPath))
+}
+
+func TestRunAsNoErrorForPipelines(t *testing.T) {
+	b := load(t, "./run_as/not_allowed/pipelines")
+
+	// We should not error because the pipeline is being deployed with the same
+	// identity as the bundle run_as identity.
+	ctx := context.Background()
+	bundle.ApplyFunc(ctx, b, func(ctx context.Context, b *bundle.Bundle) diag.Diagnostics {
+		b.Config.Workspace.CurrentUser = &config.User{
+			User: &iam.User{
+				UserName: "my_service_principal",
+			},
+		}
+		return nil
+	})
+
+	diags := bundle.Apply(ctx, b, mutator.SetRunAs())
+	assert.NoError(t, diags.Error())
+}
+
+func TestRunAsErrorForModelServing(t *testing.T) {
+	b := load(t, "./run_as/not_allowed/model_serving")
+
+	ctx := context.Background()
+	bundle.ApplyFunc(ctx, b, func(ctx context.Context, b *bundle.Bundle) diag.Diagnostics {
+		b.Config.Workspace.CurrentUser = &config.User{
+			User: &iam.User{
+				UserName: "jane@doe.com",
+			},
+		}
+		return nil
+	})
+
+	diags := bundle.Apply(ctx, b, mutator.SetRunAs())
+	err := diags.Error()
+
+	configPath := filepath.FromSlash("run_as/not_allowed/model_serving/databricks.yml")
+	assert.EqualError(t, err, fmt.Sprintf("model_serving_endpoints are not supported when the current deployment user is different from the bundle's run_as identity. Please deploy as the run_as identity. Location of the unsupported resource: %s:14:5. Current identity: jane@doe.com. Run as identity: my_service_principal", configPath))
+}
+
+func TestRunAsNoErrorForModelServingEndpoints(t *testing.T) {
+	b := load(t, "./run_as/not_allowed/model_serving")
+
+	// We should not error because the model serving endpoint is being deployed
+	// with the same identity as the bundle run_as identity.
+	ctx := context.Background()
+	bundle.ApplyFunc(ctx, b, func(ctx context.Context, b *bundle.Bundle) diag.Diagnostics {
+		b.Config.Workspace.CurrentUser = &config.User{
+			User: &iam.User{
+				UserName: "my_service_principal",
+			},
+		}
+		return nil
+	})
+
+	diags := bundle.Apply(ctx, b, mutator.SetRunAs())
+	assert.NoError(t, diags.Error())
+}
+
+func TestRunAsErrorWhenBothUserAndSpSpecified(t *testing.T) {
+	b := load(t, "./run_as/not_allowed/both_sp_and_user")
+
+	ctx := context.Background()
+	bundle.ApplyFunc(ctx, b, func(ctx context.Context, b *bundle.Bundle) diag.Diagnostics {
+		b.Config.Workspace.CurrentUser = &config.User{
+			User: &iam.User{
+				UserName: "my_service_principal",
+			},
+		}
+		return nil
+	})
+
+	diags := bundle.Apply(ctx, b, mutator.SetRunAs())
+	err := diags.Error()
+
+	configPath := filepath.FromSlash("run_as/not_allowed/both_sp_and_user/databricks.yml")
+	assert.EqualError(t, err, fmt.Sprintf("run_as section must specify exactly one identity. A service_principal_name \"my_service_principal\" is specified at %s:6:27. A user_name \"my_user_name\" is defined at %s:7:14", configPath, configPath))
+}
+
+func TestRunAsErrorNeitherUserOrSpSpecified(t *testing.T) {
+	b := load(t, "./run_as/not_allowed/neither_sp_nor_user")
+
+	ctx := context.Background()
+	bundle.ApplyFunc(ctx, b, func(ctx context.Context, b *bundle.Bundle) diag.Diagnostics {
+		b.Config.Workspace.CurrentUser = &config.User{
+			User: &iam.User{
+				UserName: "my_service_principal",
+			},
+		}
+		return nil
+	})
+
+	diags := bundle.Apply(ctx, b, mutator.SetRunAs())
+	err := diags.Error()
+
+	configPath := filepath.FromSlash("run_as/not_allowed/neither_sp_nor_user/databricks.yml")
+	assert.EqualError(t, err, fmt.Sprintf("run_as section must specify exactly one identity. Neither service_principal_name nor user_name is specified at %s:4:8", configPath))
+}
+
+func TestRunAsErrorNeitherUserOrSpSpecifiedAtTargetOverride(t *testing.T) {
+	b := loadTarget(t, "./run_as/not_allowed/neither_sp_nor_user_override", "development")
+
+	ctx := context.Background()
+	bundle.ApplyFunc(ctx, b, func(ctx context.Context, b *bundle.Bundle) diag.Diagnostics {
+		b.Config.Workspace.CurrentUser = &config.User{
+			User: &iam.User{
+				UserName: "my_service_principal",
+			},
+		}
+		return nil
+	})
+
+	diags := bundle.Apply(ctx, b, mutator.SetRunAs())
+	err := diags.Error()
+
+	configPath := filepath.FromSlash("run_as/not_allowed/neither_sp_nor_user_override/override.yml")
+	assert.EqualError(t, err, fmt.Sprintf("run_as section must specify exactly one identity. Neither service_principal_name nor user_name is specified at %s:4:12", configPath))
 }