Add environments to project configuration (#68)

cmd/init/legacy-cli.go

@@ -29,7 +29,14 @@ func loadCliProfiles() (profiles []prompt.Answer, err error) {
 			Value:   v.Name(),
 			Details: fmt.Sprintf(`Connecting to "%s" workspace`, host),
 			Callback: func(ans prompt.Answer, config *project.Config, _ prompt.Results) {
-				config.Profile = ans.Value
+				if config.Environments == nil {
+					config.Environments = make(map[string]project.Environment)
+				}
+				config.Environments[project.DefaultEnvironment] = project.Environment{
+					Workspace: project.Workspace{
+						Profile: ans.Value,
+					},
+				}
 			},
 		})
 	}

project/config.go

@@ -50,6 +50,13 @@ type Config struct {
 	// created by administrator users or admin-level automation, like Terraform
 	// and/or SCIM provisioning.
 	Assertions *Assertions `json:"assertions,omitempty"`
+
+	// Environments contain this project's defined environments.
+	// They can be used to differentiate settings and resources between
+	// development, staging, production, etc.
+	// If not specified, the code below initializes this field with a
+	// single default-initialized environment called "development".
+	Environments map[string]Environment `json:"environments"`
 }
 
 func (c Config) IsDevClusterDefined() bool {
@@ -83,7 +90,7 @@ func loadProjectConf(root string) (c Config, err error) {
 		baseDir := filepath.Base(root)
 		// If bricks config file is missing we assume the project root dir name
 		// as the name of the project
-		return Config{Name: baseDir}, nil
+		return validateAndApplyProjectDefaults(Config{Name: baseDir})
 	}
 
 	config, err := os.Open(configFilePath)
@@ -102,6 +109,11 @@ func loadProjectConf(root string) (c Config, err error) {
 }
 
 func validateAndApplyProjectDefaults(c Config) (Config, error) {
+	// If no environments are specified, define default environment under default name.
+	if c.Environments == nil {
+		c.Environments = make(map[string]Environment)
+		c.Environments[DefaultEnvironment] = Environment{}
+	}
 	// defaultCluster := clusters.ClusterInfo{
 	// 	NodeTypeID: "smallest",
 	// 	SparkVersion: "latest",

project/environment.go

@@ -0,0 +1,41 @@
+package project
+
+import (
+	"os"
+
+	"github.com/spf13/cobra"
+)
+
+const bricksEnv = "BRICKS_ENV"
+
+const DefaultEnvironment = "development"
+
+// Workspace defines configurables at the workspace level.
+type Workspace struct {
+	Profile string `json:"profile,omitempty"`
+}
+
+// Environment defines all configurables for a single environment.
+type Environment struct {
+	Workspace Workspace `json:"workspace"`
+}
+
+// getEnvironment returns the name of the environment to operate in.
+func getEnvironment(cmd *cobra.Command) (value string) {
+	// The command line flag takes precedence.
+	flag := cmd.Flag("environment")
+	if flag != nil {
+		value = flag.Value.String()
+		if value != "" {
+			return
+		}
+	}
+
+	// If it's not set, use the environment variable.
+	value = os.Getenv(bricksEnv)
+	if value != "" {
+		return
+	}
+
+	return DefaultEnvironment
+}

project/environment_test.go

@@ -0,0 +1,38 @@
+package project
+
+import (
+	"testing"
+
+	"github.com/spf13/cobra"
+	"github.com/stretchr/testify/assert"
+)
+
+func TestEnvironmentFromCommand(t *testing.T) {
+	var cmd cobra.Command
+	cmd.Flags().String("environment", "", "specify environment")
+	cmd.Flags().Set("environment", "env-from-arg")
+	t.Setenv(bricksEnv, "")
+
+	value := getEnvironment(&cmd)
+	assert.Equal(t, "env-from-arg", value)
+}
+
+func TestEnvironmentFromEnvironment(t *testing.T) {
+	var cmd cobra.Command
+	cmd.Flags().String("environment", "", "specify environment")
+	cmd.Flags().Set("environment", "")
+	t.Setenv(bricksEnv, "env-from-env")
+
+	value := getEnvironment(&cmd)
+	assert.Equal(t, "env-from-env", value)
+}
+
+func TestEnvironmentDefault(t *testing.T) {
+	var cmd cobra.Command
+	cmd.Flags().String("environment", "", "specify environment")
+	cmd.Flags().Set("environment", "")
+	t.Setenv(bricksEnv, "")
+
+	value := getEnvironment(&cmd)
+	assert.Equal(t, DefaultEnvironment, value)
+}

project/project.go

@@ -17,10 +17,12 @@ type project struct {
 	mu sync.Mutex
 
 	root string
+	env  string
 
-	config *Config
-	wsc    *workspaces.WorkspacesClient
-	me     *scim.User
+	config      *Config
+	environment *Environment
+	wsc         *workspaces.WorkspacesClient
+	me          *scim.User
 }
 
 // Configure is used as a PreRunE function for all commands that
@@ -32,7 +34,7 @@ func Configure(cmd *cobra.Command, args []string) error {
 		return err
 	}
 
-	ctx, err := Initialize(cmd.Context(), root)
+	ctx, err := Initialize(cmd.Context(), root, getEnvironment(cmd))
 	if err != nil {
 		return err
 	}
@@ -44,33 +46,46 @@ func Configure(cmd *cobra.Command, args []string) error {
 // Placeholder to use as unique key in context.Context.
 var projectKey int
 
-// Initialize loads a project configuration given a root.
+// Initialize loads a project configuration given a root and environment.
 // It stores the project on a new context.
 // The project is available through the `Get()` function.
-func Initialize(ctx context.Context, root string) (context.Context, error) {
+func Initialize(ctx context.Context, root, env string) (context.Context, error) {
 	config, err := loadProjectConf(root)
 	if err != nil {
 		return nil, err
 	}
 
+	// Confirm that the specified environment is valid.
+	environment, ok := config.Environments[env]
+	if !ok {
+		return nil, fmt.Errorf("environment [%s] not defined", env)
+	}
+
 	p := project{
-		root:   root,
-		config: &config,
-	}
-
-	if config.Profile == "" {
-		// Bricks config doesn't define the profile to use, so go sdk will figure
-		// out the auth credentials based on the enviroment.
-		// eg. DATABRICKS_CONFIG_PROFILE can be used to select which profile to use or
-		// DATABRICKS_HOST and DATABRICKS_TOKEN can be used to set the workspace auth creds
-		p.wsc = workspaces.New()
-	} else {
-		p.wsc = workspaces.New(&databricks.Config{Profile: config.Profile})
+		root: root,
+		env:  env,
+
+		config:      &config,
+		environment: &environment,
 	}
 
+	p.initializeWorkspacesClient(ctx)
 	return context.WithValue(ctx, &projectKey, &p), nil
 }
 
+func (p *project) initializeWorkspacesClient(ctx context.Context) {
+	var config databricks.Config
+
+	// If the config specifies a profile, or other authentication related properties,
+	// pass them along to the SDK here. If nothing is defined, the SDK will figure
+	// out which autentication mechanism to use using enviroment variables.
+	if p.environment.Workspace.Profile != "" {
+		config.Profile = p.environment.Workspace.Profile
+	}
+
+	p.wsc = workspaces.New(&config)
+}
+
 // Get returns the project as configured on the context.
 // It panics if it isn't configured.
 func Get(ctx context.Context) *project {
@@ -90,6 +105,14 @@ func (p *project) Root() string {
 	return p.root
 }
 
+func (p *project) Config() Config {
+	return *p.config
+}
+
+func (p *project) Environment() Environment {
+	return *p.environment
+}
+
 func (p *project) Me() (*scim.User, error) {
 	p.mu.Lock()
 	defer p.mu.Unlock()

project/project_test.go

@@ -9,7 +9,7 @@ import (
 )
 
 func TestProjectInitialize(t *testing.T) {
-	ctx, err := Initialize(context.Background(), "./testdata")
+	ctx, err := Initialize(context.Background(), "./testdata", DefaultEnvironment)
 	require.NoError(t, err)
 	assert.Equal(t, Get(ctx).config.Name, "dev")
 }

project/testdata/databricks.yml

@@ -1,4 +1,4 @@
 name: dev
 profile: demo
 dev_cluster:
-  cluster_name: Shared Autoscaling
+  cluster_name: Shared Autoscaling