Extend deployment mode support

bundle/config/environment.go

@@ -32,7 +32,13 @@ type Environment struct {
 }
 
 const (
-	// Right now, we just have a default / "" mode and a "development" mode.
-	// Additional modes are expected to come for pull-requests and production.
+	// Development mode: deployments done purely for running things in development.
+	// Any deployed resources will be marked as "dev" and might be hidden or cleaned up.
 	Development Mode = "development"
+
+	// Production mode: deployments done for production purposes.
+	// Any deployed resources will not be changed but this mode will enable
+	// various strictness checks to make sure that a deployment is correctly setup
+	// for production purposes.
+	Production Mode = "production"
 )

bundle/config/mutator/expand_workspace_root_test.go

@@ -16,8 +16,10 @@ func TestExpandWorkspaceRoot(t *testing.T) {
 	bundle := &bundle.Bundle{
  Config: config.Root{
  Workspace: config.Workspace{
- CurrentUser: &iam.User{
- UserName: "jane@doe.com",
+ CurrentUser: &config.User{
+ User: &iam.User{
+ UserName: "jane@doe.com",
+ },
  },
  RootPath: "~/foo",
  },
@@ -32,8 +34,10 @@ func TestExpandWorkspaceRootDoesNothing(t *testing.T) {
 	bundle := &bundle.Bundle{
  Config: config.Root{
  Workspace: config.Workspace{
- CurrentUser: &iam.User{
- UserName: "jane@doe.com",
+ CurrentUser: &config.User{
+ User: &iam.User{
+ UserName: "jane@doe.com",
+ },
  },
  RootPath: "/Users/charly@doe.com/foo",
  },
@@ -48,8 +52,10 @@ func TestExpandWorkspaceRootWithoutRoot(t *testing.T) {
 	bundle := &bundle.Bundle{
  Config: config.Root{
  Workspace: config.Workspace{
- CurrentUser: &iam.User{
- UserName: "jane@doe.com",
+ CurrentUser: &config.User{
+ User: &iam.User{
+ UserName: "jane@doe.com",
+ },
  },
  },
  },

bundle/config/mutator/populate_current_user.go

@@ -2,8 +2,11 @@ package mutator
 
 import (
 	"context"
+	"strings"
+	"unicode"
 
 	"github.com/databricks/cli/bundle"
+	"github.com/databricks/cli/bundle/config"
 )
 
 type populateCurrentUser struct{}
@@ -24,6 +27,22 @@ func (m *populateCurrentUser) Apply(ctx context.Context, b *bundle.Bundle) error
  return err
 	}
 
-	b.Config.Workspace.CurrentUser = me
+	b.Config.Workspace.CurrentUser = &config.User{
+ ShortName: getShortUserName(me.UserName),
+ User: me,
+	}
 	return nil
 }
+
+// Get a short-form username, based on the user's primary email address.
+// We leave the full range of unicode letters in tact, but remove all "special" characters,
+// including dots, which are not supported in e.g. experiment names.
+func getShortUserName(emailAddress string) string {
+	r := []rune(strings.Split(emailAddress, "@")[0])
+	for i := 0; i < len(r); i++ {
+ if !unicode.IsLetter(r[i]) {
+ r[i] = '_'
+ }
+	}
+	return string(r)
+}

bundle/config/mutator/populate_current_user_test.go

@@ -1,3 +1,40 @@
 package mutator
 
-// We need to implement workspace client mocking to implement this test.
+import "testing"
+
+func TestPopulateCurrentUser(t *testing.T) {
+	// We need to implement workspace client mocking to implement this test.
+}
+
+func TestGetShortUserName(t *testing.T) {
+	tests := []struct {
+ name string
+ email string
+ expected string
+	}{
+ {
+ name: "test alphanumeric characters",
+ email: "test.user@example.com",
+ expected: "test_user",
+ },
+ {
+ name: "test unicode characters",
+ email: "tést.üser@example.com",
+ expected: "tést_üser",
+ },
+ {
+ name: "test special characters",
+ email: "test$.user@example.com",
+ expected: "test__user",
+ },
+	}
+
+	for _, tt := range tests {
+ t.Run(tt.name, func(t *testing.T) {
+ result := getShortUserName(tt.email)
+ if result != tt.expected {
+ t.Errorf("getShortUserName(%q) = %q; expected %q", tt.email, result, tt.expected)
+ }
+ })
+	}
+}

bundle/config/mutator/process_environment_mode.go

@@ -4,9 +4,11 @@ import (
 	"context"
 	"fmt"
 	"path"
+	"strings"
 
 	"github.com/databricks/cli/bundle"
 	"github.com/databricks/cli/bundle/config"
+	"github.com/databricks/databricks-sdk-go/service/iam"
 	"github.com/databricks/databricks-sdk-go/service/jobs"
 	"github.com/databricks/databricks-sdk-go/service/ml"
 )
@@ -26,15 +28,17 @@ func (m *processEnvironmentMode) Name() string {
 // Mark all resources as being for 'development' purposes, i.e.
 // changing their their name, adding tags, and (in the future)
 // marking them as 'hidden' in the UI.
-func processDevelopmentMode(b *bundle.Bundle) error {
+func transformDevelopmentMode(b *bundle.Bundle) error {
 	r := b.Config.Resources
 
+	prefix := "[dev " + b.Config.Workspace.CurrentUser.ShortName + "] "
+
 	for i := range r.Jobs {
- r.Jobs[i].Name = "[dev] " + r.Jobs[i].Name
+ r.Jobs[i].Name = prefix + r.Jobs[i].Name
  if r.Jobs[i].Tags == nil {
  r.Jobs[i].Tags = make(map[string]string)
  }
- r.Jobs[i].Tags["dev"] = ""
+ r.Jobs[i].Tags["dev"] = b.Config.Workspace.CurrentUser.DisplayName
  if r.Jobs[i].MaxConcurrentRuns == 0 {
  r.Jobs[i].MaxConcurrentRuns = developmentConcurrentRuns
  }
@@ -50,13 +54,13 @@ func processDevelopmentMode(b *bundle.Bundle) error {
 	}
 
 	for i := range r.Pipelines {
- r.Pipelines[i].Name = "[dev] " + r.Pipelines[i].Name
+ r.Pipelines[i].Name = prefix + r.Pipelines[i].Name
  r.Pipelines[i].Development = true
  // (pipelines don't yet support tags)
 	}
 
 	for i := range r.Models {
- r.Models[i].Name = "[dev] " + r.Models[i].Name
+ r.Models[i].Name = prefix + r.Models[i].Name
  r.Models[i].Tags = append(r.Models[i].Tags, ml.ModelTag{Key: "dev", Value: ""})
 	}
 
@@ -65,20 +69,111 @@ func processDevelopmentMode(b *bundle.Bundle) error {
  dir := path.Dir(filepath)
  base := path.Base(filepath)
  if dir == "." {
- r.Experiments[i].Name = "[dev] " + base
+ r.Experiments[i].Name = prefix + base
  } else {
- r.Experiments[i].Name = dir + "/[dev] " + base
+ r.Experiments[i].Name = dir + "/" + prefix + base
+ }
+ r.Experiments[i].Tags = append(r.Experiments[i].Tags, ml.ExperimentTag{Key: "dev", Value: b.Config.Workspace.CurrentUser.DisplayName})
+	}
+
+	return nil
+}
+
+func validateDevelopmentMode(b *bundle.Bundle) error {
+	if path := findIncorrectPath(b, config.Development); path != "" {
+ return fmt.Errorf("%s must start with '~/' or contain the current username when using 'mode: development'", path)
+	}
+	return nil
+}
+
+func findIncorrectPath(b *bundle.Bundle, mode config.Mode) string {
+	username := b.Config.Workspace.CurrentUser.UserName
+	containsExpected := true
+	if mode == config.Production {
+ containsExpected = false
+	}
+
+	if strings.Contains(b.Config.Workspace.RootPath, username) != containsExpected && b.Config.Workspace.RootPath != "" {
+ return "root_path"
+	}
+	if strings.Contains(b.Config.Workspace.StatePath, username) != containsExpected {
+ return "state_path"
+	}
+	if strings.Contains(b.Config.Workspace.FilesPath, username) != containsExpected {
+ return "files_path"
+	}
+	if strings.Contains(b.Config.Workspace.ArtifactsPath, username) != containsExpected {
+ return "artifacts_path"
+	}
+	return ""
+}
+
+func validateProductionMode(ctx context.Context, b *bundle.Bundle, isPrincipalUsed bool) error {
+	r := b.Config.Resources
+	for i := range r.Pipelines {
+ if r.Pipelines[i].Development {
+ return fmt.Errorf("environment with 'mode: production' cannot specify a pipeline with 'development: true'")
  }
- r.Experiments[i].Tags = append(r.Experiments[i].Tags, ml.ExperimentTag{Key: "dev", Value: ""})
 	}
 
+	if !isPrincipalUsed {
+ if path := findIncorrectPath(b, config.Production); path != "" {
+ message := "%s must not contain the current username when using 'mode: production'"
+ if path == "root_path" {
+ return fmt.Errorf(message+"\n tip: set workspace.root_path to a shared path such as /Shared/.bundle/${bundle.name}/${bundle.environment}", path)
+ } else {
+ return fmt.Errorf(message, path)
+ }
+ }
+
+ if !isRunAsSet(r) {
+ return fmt.Errorf("'run_as' must be set for all jobs when using 'mode: production'")
+ }
+	}
 	return nil
 }
 
+// Determines whether a service principal identity is used to run the CLI.
+func isServicePrincipalUsed(ctx context.Context, b *bundle.Bundle) (bool, error) {
+	ws := b.WorkspaceClient()
+
+	// Check if a principal with the current user's ID exists.
+	// We need to use the ListAll method since Get is only usable by admins.
+	matches, err := ws.ServicePrincipals.ListAll(ctx, iam.ListServicePrincipalsRequest{
+ Filter: "id eq " + b.Config.Workspace.CurrentUser.Id,
+	})
+	if err != nil {
+ return false, err
+	}
+	return len(matches) > 0, nil
+}
+
+// Determines whether run_as is explicitly set for all resources.
+// We do this in a best-effort fashion rather than check the top-level
+// 'run_as' field because the latter is not required to be set.
+func isRunAsSet(r config.Resources) bool {
+	for i := range r.Jobs {
+ if r.Jobs[i].RunAs == nil {
+ return false
+ }
+	}
+	return true
+}
+
 func (m *processEnvironmentMode) Apply(ctx context.Context, b *bundle.Bundle) error {
 	switch b.Config.Bundle.Mode {
 	case config.Development:
- return processDevelopmentMode(b)
+ err := validateDevelopmentMode(b)
+ if err != nil {
+ return err
+ }
+ return transformDevelopmentMode(b)
+	case config.Production:
+ isPrincipal, err := isServicePrincipalUsed(ctx, b)
+ if err != nil {
+ return err
+ }
+ return validateProductionMode(ctx, b, isPrincipal)
 	case "":
  // No action
 	default: