databricks-cli/bundle/tests/run_as_test.go

package config_tests

import (
	"context"
	"fmt"
	"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/databricks/databricks-sdk-go/service/serving"
	"github.com/stretchr/testify/assert"
)

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 {
		b.Config.Workspace.CurrentUser = &config.User{
			User: &iam.User{
				UserName: "jane@doe.com",
			},
		}
		return nil
	})

	diags := bundle.Apply(ctx, b, mutator.SetRunAs())
	assert.NoError(t, diags.Error())

	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)

	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)

	// 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)

	// 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 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 {
		b.Config.Workspace.CurrentUser = &config.User{
			User: &iam.User{
				UserName: "jane@doe.com",
			},
		}
		return nil
	})

	diags := bundle.Apply(ctx, b, mutator.SetRunAs())
	assert.NoError(t, diags.Error())

	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)

	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)

	// 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)

	// 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 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()

	assert.ErrorContains(t, err, "pipelines do not support a setting a run_as user that is different from the owner.\n"+
		"Current identity: jane@doe.com. Run as identity: my_service_principal.\n"+
		"See https://docs")
}

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()

	assert.ErrorContains(t, err, "model_serving_endpoints do not support a setting a run_as user that is different from the owner.\n"+
		"Current identity: jane@doe.com. Run as identity: my_service_principal.\n"+
		"See https://docs")
}

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()

	assert.ErrorContains(t, err, "run_as section cannot specify both user_name and service_principal_name")
}

func TestRunAsErrorNeitherUserOrSpSpecified(t *testing.T) {
	tcases := []struct {
		name string
		err  string
	}{
		{
			name: "empty_run_as",
			err:  "run_as section must specify exactly one identity. Neither service_principal_name nor user_name is specified",
		},
		{
			name: "empty_sp",
			err:  "run_as section must specify exactly one identity. Neither service_principal_name nor user_name is specified",
		},
		{
			name: "empty_user",
			err:  "run_as section must specify exactly one identity. Neither service_principal_name nor user_name is specified",
		},
		{
			name: "empty_user_and_sp",
			err:  "run_as section must specify exactly one identity. Neither service_principal_name nor user_name is specified",
		},
	}

	for _, tc := range tcases {
		t.Run(tc.name, func(t *testing.T) {

			bundlePath := fmt.Sprintf("./run_as/not_allowed/neither_sp_nor_user/%s", tc.name)
			b := load(t, bundlePath)

			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()
			assert.EqualError(t, err, tc.err)
		})
	}
}

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()

	assert.EqualError(t, err, "run_as section must specify exactly one identity. Neither service_principal_name nor user_name is specified")
}

func TestLegacyRunAs(t *testing.T) {
	b := load(t, "./run_as/legacy")

	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())
	assert.NoError(t, diags.Error())

	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)

	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)

	// job_three should retain it's 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)

	// Assert owner permissions for pipelines are set.
	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)
	assert.Equal(t, serving.CreateServingEndpoint{Name: "skynet"}, *b.Config.Resources.ModelServingEndpoints["model_serving_one"].CreateServingEndpoint)
}