Added support for top-level permissions (#928)

## Changes Now it's possible to define top level `permissions` section in bundle configuration and permissions defined there will be applied to all resources defined in the bundle. Supported top-level permission levels: CAN_MANAGE, CAN_VIEW, CAN_RUN. Permissions are applied to: Jobs, DLT Pipelines, ML Models, ML Experiments and Model Service Endpoints ``` bundle: name: permissions workspace: host: *** permissions: - level: CAN_VIEW group_name: test-group - level: CAN_MANAGE user_name: [email protected] - level: CAN_RUN service_principal_name: 123456-abcdef ``` ## Tests Added corresponding unit tests + ran `bundle validate` and `bundle deploy` manually
databricks · Nov 13, 2023 · f3db42e · f3db42e
1 parent 14d2d0a
commit f3db42e
Show file tree

Hide file tree

Showing 11 changed files with 678 additions and 1 deletion.
diff --git a/bundle/config/root.go b/bundle/config/root.go
@@ -6,6 +6,7 @@ import (
 	"path/filepath"
 	"strings"
 
+	"github.com/databricks/cli/bundle/config/resources"
 	"github.com/databricks/cli/bundle/config/variable"
 	"github.com/databricks/databricks-sdk-go/service/jobs"
 	"github.com/ghodss/yaml"
@@ -56,6 +57,10 @@ type Root struct {
 	RunAs *jobs.JobRunAs `json:"run_as,omitempty"`
 
 	Experimental *Experimental `json:"experimental,omitempty"`
+
+	// Permissions section allows to define permissions which will be
+	// applied to all resources defined in bundle
+	Permissions []resources.Permission `json:"permissions,omitempty"`
 }
 
 // Load loads the bundle configuration file at the specified path.
@@ -237,5 +242,12 @@ func (r *Root) MergeTargetOverrides(target *Target) error {
 		}
 	}
 
+	if target.Permissions != nil {
+		err = mergo.Merge(&r.Permissions, target.Permissions, mergo.WithAppendSlice)
+		if err != nil {
+			return err
+		}
+	}
+
 	return nil
 }
diff --git a/bundle/config/target.go b/bundle/config/target.go
@@ -1,6 +1,9 @@
 package config
 
-import "github.com/databricks/databricks-sdk-go/service/jobs"
+import (
+	"github.com/databricks/cli/bundle/config/resources"
+	"github.com/databricks/databricks-sdk-go/service/jobs"
+)
 
 type Mode string
 
@@ -37,6 +40,8 @@ type Target struct {
 	RunAs *jobs.JobRunAs `json:"run_as,omitempty"`
 
 	Sync *Sync `json:"sync,omitempty"`
+
+	Permissions []resources.Permission `json:"permissions,omitempty"`
 }
 
 const (

diff --git a/bundle/permissions/mutator.go b/bundle/permissions/mutator.go
@@ -0,0 +1,136 @@
+package permissions
+
+import (
+	"context"
+	"fmt"
+	"slices"
+	"strings"
+
+	"github.com/databricks/cli/bundle"
+)
+
+const CAN_MANAGE = "CAN_MANAGE"
+const CAN_VIEW = "CAN_VIEW"
+const CAN_RUN = "CAN_RUN"
+
+var allowedLevels = []string{CAN_MANAGE, CAN_VIEW, CAN_RUN}
+var levelsMap = map[string](map[string]string){
+	"jobs": {
+		CAN_MANAGE: "CAN_MANAGE",
+		CAN_VIEW:   "CAN_VIEW",
+		CAN_RUN:    "CAN_MANAGE_RUN",
+	},
+	"pipelines": {
+		CAN_MANAGE: "CAN_MANAGE",
+		CAN_VIEW:   "CAN_VIEW",
+		CAN_RUN:    "CAN_RUN",
+	},
+	"mlflow_experiments": {
+		CAN_MANAGE: "CAN_MANAGE",
+		CAN_VIEW:   "CAN_READ",
+	},
+	"mlflow_models": {
+		CAN_MANAGE: "CAN_MANAGE",
+		CAN_VIEW:   "CAN_READ",
+	},
+	"model_serving_endpoints": {
+		CAN_MANAGE: "CAN_MANAGE",
+		CAN_VIEW:   "CAN_VIEW",
+		CAN_RUN:    "CAN_QUERY",
+	},
+}
+
+type bundlePermissions struct{}
+
+func ApplyBundlePermissions() bundle.Mutator {
+	return &bundlePermissions{}
+}
+
+func (m *bundlePermissions) Apply(ctx context.Context, b *bundle.Bundle) error {
+	err := validate(b)
+	if err != nil {
+		return err
+	}
+
+	applyForJobs(ctx, b)
+	applyForPipelines(ctx, b)
+	applyForMlModels(ctx, b)
+	applyForMlExperiments(ctx, b)
+	applyForModelServiceEndpoints(ctx, b)
+
+	return nil
+}
+
+func validate(b *bundle.Bundle) error {
+	for _, p := range b.Config.Permissions {
+		if !slices.Contains(allowedLevels, p.Level) {
+			return fmt.Errorf("invalid permission level: %s, allowed values: [%s]", p.Level, strings.Join(allowedLevels, ", "))
+		}
+	}
+
+	return nil
+}
+
+func applyForJobs(ctx context.Context, b *bundle.Bundle) {
+	for _, job := range b.Config.Resources.Jobs {
+		job.Permissions = append(job.Permissions, convert(
+			ctx,
+			b.Config.Permissions,
+			job.Permissions,
+			job.Name,
+			levelsMap["jobs"],
+		)...)
+	}
+}
+
+func applyForPipelines(ctx context.Context, b *bundle.Bundle) {
+	for _, pipeline := range b.Config.Resources.Pipelines {
+		pipeline.Permissions = append(pipeline.Permissions, convert(
+			ctx,
+			b.Config.Permissions,
+			pipeline.Permissions,
+			pipeline.Name,
+			levelsMap["pipelines"],
+		)...)
+	}
+}
+
+func applyForMlExperiments(ctx context.Context, b *bundle.Bundle) {
+	for _, experiment := range b.Config.Resources.Experiments {
+		experiment.Permissions = append(experiment.Permissions, convert(
+			ctx,
+			b.Config.Permissions,
+			experiment.Permissions,
+			experiment.Name,
+			levelsMap["mlflow_experiments"],
+		)...)
+	}
+}
+
+func applyForMlModels(ctx context.Context, b *bundle.Bundle) {
+	for _, model := range b.Config.Resources.Models {
+		model.Permissions = append(model.Permissions, convert(
+			ctx,
+			b.Config.Permissions,
+			model.Permissions,
+			model.Name,
+			levelsMap["mlflow_models"],
+		)...)
+	}
+}
+
+func applyForModelServiceEndpoints(ctx context.Context, b *bundle.Bundle) {
+	for _, model := range b.Config.Resources.ModelServingEndpoints {
+		model.Permissions = append(model.Permissions, convert(
+			ctx,
+			b.Config.Permissions,
+			model.Permissions,
+			model.Name,
+			levelsMap["model_serving_endpoints"],
+		)...)
+	}
+}
+
+func (m *bundlePermissions) Name() string {
+	return "ApplyBundlePermissions"
+}
diff --git a/bundle/permissions/mutator_test.go b/bundle/permissions/mutator_test.go
@@ -0,0 +1,141 @@
+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/jobs"
+	"github.com/databricks/databricks-sdk-go/service/ml"
+	"github.com/databricks/databricks-sdk-go/service/pipelines"
+	"github.com/databricks/databricks-sdk-go/service/serving"
+	"github.com/stretchr/testify/require"
+)
+
+func TestApplyBundlePermissions(t *testing.T) {
+	b := &bundle.Bundle{
+		Config: config.Root{
+			Workspace: config.Workspace{
+				RootPath: "/Users/[email protected]",
+			},
+			Permissions: []resources.Permission{
+				{Level: CAN_MANAGE, UserName: "TestUser"},
+				{Level: CAN_VIEW, GroupName: "TestGroup"},
+				{Level: CAN_RUN, ServicePrincipalName: "TestServicePrincipal"},
+			},
+			Resources: config.Resources{
+				Jobs: map[string]*resources.Job{
+					"job_1": {JobSettings: &jobs.JobSettings{}},
+					"job_2": {JobSettings: &jobs.JobSettings{}},
+				},
+				Pipelines: map[string]*resources.Pipeline{
+					"pipeline_1": {PipelineSpec: &pipelines.PipelineSpec{}},
+					"pipeline_2": {PipelineSpec: &pipelines.PipelineSpec{}},
+				},
+				Models: map[string]*resources.MlflowModel{
+					"model_1": {Model: &ml.Model{}},
+					"model_2": {Model: &ml.Model{}},
+				},
+				Experiments: map[string]*resources.MlflowExperiment{
+					"experiment_1": {Experiment: &ml.Experiment{}},
+					"experiment_2": {Experiment: &ml.Experiment{}},
+				},
+				ModelServingEndpoints: map[string]*resources.ModelServingEndpoint{
+					"endpoint_1": {CreateServingEndpoint: &serving.CreateServingEndpoint{}},
+					"endpoint_2": {CreateServingEndpoint: &serving.CreateServingEndpoint{}},
+				},
+			},
+		},
+	}
+
+	err := bundle.Apply(context.Background(), b, ApplyBundlePermissions())
+	require.NoError(t, err)
+
+	require.Len(t, b.Config.Resources.Jobs["job_1"].Permissions, 3)
+	require.Contains(t, b.Config.Resources.Jobs["job_1"].Permissions, resources.Permission{Level: "CAN_MANAGE", UserName: "TestUser"})
+	require.Contains(t, b.Config.Resources.Jobs["job_1"].Permissions, resources.Permission{Level: "CAN_VIEW", GroupName: "TestGroup"})
+	require.Contains(t, b.Config.Resources.Jobs["job_1"].Permissions, resources.Permission{Level: "CAN_MANAGE_RUN", ServicePrincipalName: "TestServicePrincipal"})
+
+	require.Len(t, b.Config.Resources.Jobs["job_2"].Permissions, 3)
+	require.Contains(t, b.Config.Resources.Jobs["job_2"].Permissions, resources.Permission{Level: "CAN_MANAGE", UserName: "TestUser"})
+	require.Contains(t, b.Config.Resources.Jobs["job_2"].Permissions, resources.Permission{Level: "CAN_VIEW", GroupName: "TestGroup"})
+	require.Contains(t, b.Config.Resources.Jobs["job_2"].Permissions, resources.Permission{Level: "CAN_MANAGE_RUN", ServicePrincipalName: "TestServicePrincipal"})
+
+	require.Len(t, b.Config.Resources.Pipelines["pipeline_1"].Permissions, 3)
+	require.Contains(t, b.Config.Resources.Pipelines["pipeline_1"].Permissions, resources.Permission{Level: "CAN_MANAGE", UserName: "TestUser"})
+	require.Contains(t, b.Config.Resources.Pipelines["pipeline_1"].Permissions, resources.Permission{Level: "CAN_VIEW", GroupName: "TestGroup"})
+	require.Contains(t, b.Config.Resources.Pipelines["pipeline_1"].Permissions, resources.Permission{Level: "CAN_RUN", ServicePrincipalName: "TestServicePrincipal"})
+
+	require.Len(t, b.Config.Resources.Pipelines["pipeline_2"].Permissions, 3)
+	require.Contains(t, b.Config.Resources.Pipelines["pipeline_2"].Permissions, resources.Permission{Level: "CAN_MANAGE", UserName: "TestUser"})
+	require.Contains(t, b.Config.Resources.Pipelines["pipeline_2"].Permissions, resources.Permission{Level: "CAN_VIEW", GroupName: "TestGroup"})
+	require.Contains(t, b.Config.Resources.Pipelines["pipeline_2"].Permissions, resources.Permission{Level: "CAN_RUN", ServicePrincipalName: "TestServicePrincipal"})
+
+	require.Len(t, b.Config.Resources.Models["model_1"].Permissions, 2)
+	require.Contains(t, b.Config.Resources.Models["model_1"].Permissions, resources.Permission{Level: "CAN_MANAGE", UserName: "TestUser"})
+	require.Contains(t, b.Config.Resources.Models["model_1"].Permissions, resources.Permission{Level: "CAN_READ", GroupName: "TestGroup"})
+
+	require.Len(t, b.Config.Resources.Models["model_2"].Permissions, 2)
+	require.Contains(t, b.Config.Resources.Models["model_2"].Permissions, resources.Permission{Level: "CAN_MANAGE", UserName: "TestUser"})
+	require.Contains(t, b.Config.Resources.Models["model_2"].Permissions, resources.Permission{Level: "CAN_READ", GroupName: "TestGroup"})
+
+	require.Len(t, b.Config.Resources.Experiments["experiment_1"].Permissions, 2)
+	require.Contains(t, b.Config.Resources.Experiments["experiment_1"].Permissions, resources.Permission{Level: "CAN_MANAGE", UserName: "TestUser"})
+	require.Contains(t, b.Config.Resources.Experiments["experiment_1"].Permissions, resources.Permission{Level: "CAN_READ", GroupName: "TestGroup"})
+
+	require.Len(t, b.Config.Resources.Experiments["experiment_2"].Permissions, 2)
+	require.Contains(t, b.Config.Resources.Experiments["experiment_2"].Permissions, resources.Permission{Level: "CAN_MANAGE", UserName: "TestUser"})
+	require.Contains(t, b.Config.Resources.Experiments["experiment_2"].Permissions, resources.Permission{Level: "CAN_READ", GroupName: "TestGroup"})
+
+	require.Len(t, b.Config.Resources.ModelServingEndpoints["endpoint_1"].Permissions, 3)
+	require.Contains(t, b.Config.Resources.ModelServingEndpoints["endpoint_1"].Permissions, resources.Permission{Level: "CAN_MANAGE", UserName: "TestUser"})
+	require.Contains(t, b.Config.Resources.ModelServingEndpoints["endpoint_1"].Permissions, resources.Permission{Level: "CAN_VIEW", GroupName: "TestGroup"})
+	require.Contains(t, b.Config.Resources.ModelServingEndpoints["endpoint_1"].Permissions, resources.Permission{Level: "CAN_QUERY", ServicePrincipalName: "TestServicePrincipal"})
+
+	require.Len(t, b.Config.Resources.ModelServingEndpoints["endpoint_2"].Permissions, 3)
+	require.Contains(t, b.Config.Resources.ModelServingEndpoints["endpoint_2"].Permissions, resources.Permission{Level: "CAN_MANAGE", UserName: "TestUser"})
+	require.Contains(t, b.Config.Resources.ModelServingEndpoints["endpoint_2"].Permissions, resources.Permission{Level: "CAN_VIEW", GroupName: "TestGroup"})
+	require.Contains(t, b.Config.Resources.ModelServingEndpoints["endpoint_2"].Permissions, resources.Permission{Level: "CAN_QUERY", ServicePrincipalName: "TestServicePrincipal"})
+}
+
+func TestWarningOnOverlapPermission(t *testing.T) {
+	b := &bundle.Bundle{
+		Config: config.Root{
+			Workspace: config.Workspace{
+				RootPath: "/Users/[email protected]",
+			},
+			Permissions: []resources.Permission{
+				{Level: CAN_MANAGE, UserName: "TestUser"},
+				{Level: CAN_VIEW, GroupName: "TestGroup"},
+			},
+			Resources: config.Resources{
+				Jobs: map[string]*resources.Job{
+					"job_1": {
+						Permissions: []resources.Permission{
+							{Level: CAN_VIEW, UserName: "TestUser"},
+						},
+						JobSettings: &jobs.JobSettings{},
+					},
+					"job_2": {
+						Permissions: []resources.Permission{
+							{Level: CAN_VIEW, UserName: "TestUser2"},
+						},
+						JobSettings: &jobs.JobSettings{},
+					},
+				},
+			},
+		},
+	}
+
+	err := bundle.Apply(context.Background(), b, ApplyBundlePermissions())
+	require.NoError(t, err)
+
+	require.Contains(t, b.Config.Resources.Jobs["job_1"].Permissions, resources.Permission{Level: "CAN_VIEW", UserName: "TestUser"})
+	require.Contains(t, b.Config.Resources.Jobs["job_1"].Permissions, resources.Permission{Level: "CAN_VIEW", GroupName: "TestGroup"})
+	require.Contains(t, b.Config.Resources.Jobs["job_2"].Permissions, resources.Permission{Level: "CAN_VIEW", UserName: "TestUser2"})
+	require.Contains(t, b.Config.Resources.Jobs["job_2"].Permissions, resources.Permission{Level: "CAN_MANAGE", UserName: "TestUser"})
+	require.Contains(t, b.Config.Resources.Jobs["job_2"].Permissions, resources.Permission{Level: "CAN_VIEW", GroupName: "TestGroup"})
+
+}