diff --git a/bundle/permissions/filter.go b/bundle/permissions/filter.go new file mode 100644 index 00000000..f4834a65 --- /dev/null +++ b/bundle/permissions/filter.go @@ -0,0 +1,80 @@ +package permissions + +import ( + "context" + + "github.com/databricks/cli/bundle" + "github.com/databricks/cli/libs/dyn" +) + +type filterCurrentUser struct{} + +// The databricks terraform provider does not allow changing the permissions of +// current user. The current user is implied to be the owner of all deployed resources. +// This mutator removes the current user from the permissions of all resources. +func FilterCurrentUser() bundle.Mutator { + return &filterCurrentUser{} +} + +func (m *filterCurrentUser) Name() string { + return "FilterCurrentUserFromPermissions" +} + +func filter(currentUser string) dyn.WalkValueFunc { + return func(p dyn.Path, v dyn.Value) (dyn.Value, error) { + // Permissions are defined at top level of a resource. We can skip walking + // after a depth of 4. + // [resource_type].[resource_name].[permissions].[array_index] + // Example: pipelines.foo.permissions.0 + if len(p) > 4 { + return v, dyn.ErrSkip + } + + // We can skip walking at a depth of 3 if the key is not "permissions". + // Example: pipelines.foo.libraries + if len(p) == 3 && p[2] != dyn.Key("permissions") { + return v, dyn.ErrSkip + } + + // We want to be at the level of an individual permission to check it's + // user_name and service_principal_name fields. + if len(p) != 4 || p[2] != dyn.Key("permissions") { + return v, nil + } + + // Filter if the user_name matches the current user + userName, ok := v.Get("user_name").AsString() + if ok && userName == currentUser { + return v, dyn.ErrDrop + } + + // Filter if the service_principal_name matches the current user + servicePrincipalName, ok := v.Get("service_principal_name").AsString() + if ok && servicePrincipalName == currentUser { + return v, dyn.ErrDrop + } + + return v, nil + + } +} + +func (m *filterCurrentUser) Apply(ctx context.Context, b *bundle.Bundle) error { + currentUser := b.Config.Workspace.CurrentUser.UserName + + return b.Config.Mutate(func(v dyn.Value) (dyn.Value, error) { + rv, err := dyn.Get(v, "resources") + if err != nil { + return dyn.InvalidValue, err + } + + // Walk the resources and filter out the current user from the permissions + nv, err := dyn.Walk(rv, filter(currentUser)) + if err != nil { + return dyn.InvalidValue, err + } + + // Set the resources with the filtered permissions back into the bundle + return dyn.Set(v, "resources", nv) + }) +} diff --git a/bundle/permissions/filter_test.go b/bundle/permissions/filter_test.go new file mode 100644 index 00000000..07f5ae77 --- /dev/null +++ b/bundle/permissions/filter_test.go @@ -0,0 +1,174 @@ +package permissions + +import ( + "context" + "testing" + + "github.com/databricks/cli/bundle" + "github.com/databricks/cli/bundle/config" + "github.com/databricks/cli/bundle/config/resources" + "github.com/databricks/databricks-sdk-go/service/iam" + "github.com/stretchr/testify/assert" +) + +var alice = resources.Permission{ + Level: CAN_MANAGE, + UserName: "alice@databricks.com", +} + +var bob = resources.Permission{ + Level: CAN_VIEW, + UserName: "bob@databricks.com", +} + +var robot = resources.Permission{ + Level: CAN_RUN, + ServicePrincipalName: "i-Robot", +} + +func testFixture(userName string) *bundle.Bundle { + p := []resources.Permission{ + alice, + bob, + robot, + } + + return &bundle.Bundle{ + Config: config.Root{ + Workspace: config.Workspace{ + CurrentUser: &config.User{ + User: &iam.User{ + UserName: userName, + }, + }, + }, + Resources: config.Resources{ + Jobs: map[string]*resources.Job{ + "job1": { + Permissions: p, + }, + "job2": { + Permissions: p, + }, + }, + Pipelines: map[string]*resources.Pipeline{ + "pipeline1": { + Permissions: p, + }, + }, + Experiments: map[string]*resources.MlflowExperiment{ + "experiment1": { + Permissions: p, + }, + }, + Models: map[string]*resources.MlflowModel{ + "model1": { + Permissions: p, + }, + }, + ModelServingEndpoints: map[string]*resources.ModelServingEndpoint{ + "endpoint1": { + Permissions: p, + }, + }, + RegisteredModels: map[string]*resources.RegisteredModel{ + "registered_model1": { + Grants: []resources.Grant{ + { + Principal: "abc", + }, + }, + }, + }, + }, + }, + } + +} + +func TestFilterCurrentUser(t *testing.T) { + b := testFixture("alice@databricks.com") + + err := bundle.Apply(context.Background(), b, FilterCurrentUser()) + assert.NoError(t, err) + + // Assert current user is filtered out. + assert.Equal(t, 2, len(b.Config.Resources.Jobs["job1"].Permissions)) + assert.Contains(t, b.Config.Resources.Jobs["job1"].Permissions, robot) + assert.Contains(t, b.Config.Resources.Jobs["job1"].Permissions, bob) + + assert.Equal(t, 2, len(b.Config.Resources.Jobs["job2"].Permissions)) + assert.Contains(t, b.Config.Resources.Jobs["job2"].Permissions, robot) + assert.Contains(t, b.Config.Resources.Jobs["job2"].Permissions, bob) + + assert.Equal(t, 2, len(b.Config.Resources.Pipelines["pipeline1"].Permissions)) + assert.Contains(t, b.Config.Resources.Pipelines["pipeline1"].Permissions, robot) + assert.Contains(t, b.Config.Resources.Pipelines["pipeline1"].Permissions, bob) + + assert.Equal(t, 2, len(b.Config.Resources.Experiments["experiment1"].Permissions)) + assert.Contains(t, b.Config.Resources.Experiments["experiment1"].Permissions, robot) + assert.Contains(t, b.Config.Resources.Experiments["experiment1"].Permissions, bob) + + assert.Equal(t, 2, len(b.Config.Resources.Models["model1"].Permissions)) + assert.Contains(t, b.Config.Resources.Models["model1"].Permissions, robot) + assert.Contains(t, b.Config.Resources.Models["model1"].Permissions, bob) + + assert.Equal(t, 2, len(b.Config.Resources.ModelServingEndpoints["endpoint1"].Permissions)) + assert.Contains(t, b.Config.Resources.ModelServingEndpoints["endpoint1"].Permissions, robot) + assert.Contains(t, b.Config.Resources.ModelServingEndpoints["endpoint1"].Permissions, bob) + + // Assert there's no change to the grant. + assert.Equal(t, 1, len(b.Config.Resources.RegisteredModels["registered_model1"].Grants)) +} + +func TestFilterCurrentServicePrincipal(t *testing.T) { + b := testFixture("i-Robot") + + err := bundle.Apply(context.Background(), b, FilterCurrentUser()) + assert.NoError(t, err) + + // Assert current user is filtered out. + assert.Equal(t, 2, len(b.Config.Resources.Jobs["job1"].Permissions)) + assert.Contains(t, b.Config.Resources.Jobs["job1"].Permissions, alice) + assert.Contains(t, b.Config.Resources.Jobs["job1"].Permissions, bob) + + assert.Equal(t, 2, len(b.Config.Resources.Jobs["job2"].Permissions)) + assert.Contains(t, b.Config.Resources.Jobs["job2"].Permissions, alice) + assert.Contains(t, b.Config.Resources.Jobs["job2"].Permissions, bob) + + assert.Equal(t, 2, len(b.Config.Resources.Pipelines["pipeline1"].Permissions)) + assert.Contains(t, b.Config.Resources.Pipelines["pipeline1"].Permissions, alice) + assert.Contains(t, b.Config.Resources.Pipelines["pipeline1"].Permissions, bob) + + assert.Equal(t, 2, len(b.Config.Resources.Experiments["experiment1"].Permissions)) + assert.Contains(t, b.Config.Resources.Experiments["experiment1"].Permissions, alice) + assert.Contains(t, b.Config.Resources.Experiments["experiment1"].Permissions, bob) + + assert.Equal(t, 2, len(b.Config.Resources.Models["model1"].Permissions)) + assert.Contains(t, b.Config.Resources.Models["model1"].Permissions, alice) + assert.Contains(t, b.Config.Resources.Models["model1"].Permissions, bob) + + assert.Equal(t, 2, len(b.Config.Resources.ModelServingEndpoints["endpoint1"].Permissions)) + assert.Contains(t, b.Config.Resources.ModelServingEndpoints["endpoint1"].Permissions, alice) + assert.Contains(t, b.Config.Resources.ModelServingEndpoints["endpoint1"].Permissions, bob) + + // Assert there's no change to the grant. + assert.Equal(t, 1, len(b.Config.Resources.RegisteredModels["registered_model1"].Grants)) +} + +func TestFilterCurrentUserDoesNotErrorWhenNoResources(t *testing.T) { + b := &bundle.Bundle{ + Config: config.Root{ + Workspace: config.Workspace{ + CurrentUser: &config.User{ + User: &iam.User{ + UserName: "abc", + }, + }, + }, + }, + } + + err := bundle.Apply(context.Background(), b, FilterCurrentUser()) + assert.NoError(t, err) +} diff --git a/bundle/phases/initialize.go b/bundle/phases/initialize.go index 2c401c6b..6761ffab 100644 --- a/bundle/phases/initialize.go +++ b/bundle/phases/initialize.go @@ -41,6 +41,7 @@ func Initialize() bundle.Mutator { mutator.TranslatePaths(), python.WrapperWarning(), permissions.ApplyBundlePermissions(), + permissions.FilterCurrentUser(), metadata.AnnotateJobs(), terraform.Initialize(), scripts.Execute(config.ScriptPostInit),