diff --git a/cmd/init/legacy-cli.go b/cmd/init/legacy-cli.go index d2c66cbff0..5391169a45 100644 --- a/cmd/init/legacy-cli.go +++ b/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, + }, + } }, }) } diff --git a/project/config.go b/project/config.go index 4b28f7f7de..f718df3124 100644 --- a/project/config.go +++ b/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", diff --git a/project/environment.go b/project/environment.go new file mode 100644 index 0000000000..77af9a3811 --- /dev/null +++ b/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 +} diff --git a/project/environment_test.go b/project/environment_test.go new file mode 100644 index 0000000000..eb5c018757 --- /dev/null +++ b/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) +} diff --git a/project/project.go b/project/project.go index 3a9dd24b9c..4cd7304bc7 100644 --- a/project/project.go +++ b/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 } - p := project{ - root: root, - config: &config, + // Confirm that the specified environment is valid. + environment, ok := config.Environments[env] + if !ok { + return nil, fmt.Errorf("environment [%s] not defined", env) } - 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}) + p := project{ + 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() diff --git a/project/project_test.go b/project/project_test.go index 8d01fe3dec..f934377791 100644 --- a/project/project_test.go +++ b/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") } diff --git a/project/testdata/databricks.yml b/project/testdata/databricks.yml index 3b8eb81fb7..2b286d17c9 100644 --- a/project/testdata/databricks.yml +++ b/project/testdata/databricks.yml @@ -1,4 +1,4 @@ name: dev profile: demo dev_cluster: - cluster_name: Shared Autoscaling \ No newline at end of file + cluster_name: Shared Autoscaling