Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
7 changes: 4 additions & 3 deletions cmd/auth/env.go
Original file line number Diff line number Diff line change
@@ -1,6 +1,7 @@
package auth

import (
"context"
"encoding/json"
"errors"
"fmt"
Expand Down Expand Up @@ -68,8 +69,8 @@ func resolveSection(cfg *config.Config, iniFile *config.File) (*ini.Section, err
return candidates[0], nil
}

func loadFromDatabricksCfg(cfg *config.Config) error {
iniFile, err := databrickscfg.Get()
func loadFromDatabricksCfg(ctx context.Context, cfg *config.Config) error {
iniFile, err := databrickscfg.Get(ctx)
if errors.Is(err, fs.ErrNotExist) {
// it's fine not to have ~/.databrickscfg
return nil
Expand Down Expand Up @@ -110,7 +111,7 @@ func newEnvCommand() *cobra.Command {
cfg.Profile = profile
} else if cfg.Host == "" {
cfg.Profile = "DEFAULT"
} else if err := loadFromDatabricksCfg(cfg); err != nil {
} else if err := loadFromDatabricksCfg(cmd.Context(), cfg); err != nil {
return err
}
// Go SDK is lazy loaded because of Terraform semantics,
Expand Down
2 changes: 1 addition & 1 deletion cmd/auth/login.go
Original file line number Diff line number Diff line change
Expand Up @@ -128,7 +128,7 @@ func newLoginCommand(persistentAuth *auth.PersistentAuth) *cobra.Command {

func setHost(ctx context.Context, profileName string, persistentAuth *auth.PersistentAuth, args []string) error {
// If the chosen profile has a hostname and the user hasn't specified a host, infer the host from the profile.
_, profiles, err := databrickscfg.LoadProfiles(func(p databrickscfg.Profile) bool {
_, profiles, err := databrickscfg.LoadProfiles(ctx, func(p databrickscfg.Profile) bool {
return p.Name == profileName
})
if err != nil {
Expand Down
2 changes: 1 addition & 1 deletion cmd/auth/profiles.go
Original file line number Diff line number Diff line change
Expand Up @@ -95,7 +95,7 @@ func newProfilesCommand() *cobra.Command {

cmd.RunE = func(cmd *cobra.Command, args []string) error {
var profiles []*profileMetadata
iniFile, err := databrickscfg.Get()
iniFile, err := databrickscfg.Get(cmd.Context())
if os.IsNotExist(err) {
// return empty list for non-configured machines
iniFile = &config.File{
Expand Down
30 changes: 11 additions & 19 deletions cmd/root/auth.go
Original file line number Diff line number Diff line change
Expand Up @@ -5,7 +5,6 @@ import (
"errors"
"fmt"
"net/http"
"os"

"github.com/databricks/cli/bundle"
"github.com/databricks/cli/libs/cmdio"
Expand Down Expand Up @@ -55,7 +54,7 @@ func accountClientOrPrompt(ctx context.Context, cfg *config.Config, allowPrompt
}

// Try picking a profile dynamically if the current configuration is not valid.
profile, err := askForAccountProfile(ctx)
profile, err := AskForAccountProfile(ctx)
if err != nil {
return nil, err
}
Expand Down Expand Up @@ -83,7 +82,7 @@ func MustAccountClient(cmd *cobra.Command, args []string) error {
// 1. only admins will have account configured
// 2. 99% of admins will have access to just one account
// hence, we don't need to create a special "DEFAULT_ACCOUNT" profile yet
_, profiles, err := databrickscfg.LoadProfiles(databrickscfg.MatchAccountProfiles)
_, profiles, err := databrickscfg.LoadProfiles(cmd.Context(), databrickscfg.MatchAccountProfiles)
if err != nil {
return err
}
Expand Down Expand Up @@ -123,7 +122,7 @@ func workspaceClientOrPrompt(ctx context.Context, cfg *config.Config, allowPromp
}

// Try picking a profile dynamically if the current configuration is not valid.
profile, err := askForWorkspaceProfile(ctx)
profile, err := AskForWorkspaceProfile(ctx)
if err != nil {
return nil, err
}
Expand Down Expand Up @@ -173,21 +172,14 @@ func SetWorkspaceClient(ctx context.Context, w *databricks.WorkspaceClient) cont
return context.WithValue(ctx, &workspaceClient, w)
}

func transformLoadError(path string, err error) error {
if os.IsNotExist(err) {
return fmt.Errorf("no configuration file found at %s; please create one first", path)
}
return err
}

func askForWorkspaceProfile(ctx context.Context) (string, error) {
path, err := databrickscfg.GetPath()
func AskForWorkspaceProfile(ctx context.Context) (string, error) {
path, err := databrickscfg.GetPath(ctx)
if err != nil {
return "", fmt.Errorf("cannot determine Databricks config file path: %w", err)
}
file, profiles, err := databrickscfg.LoadProfiles(databrickscfg.MatchWorkspaceProfiles)
file, profiles, err := databrickscfg.LoadProfiles(ctx, databrickscfg.MatchWorkspaceProfiles)
if err != nil {
return "", transformLoadError(path, err)
return "", err
}
switch len(profiles) {
case 0:
Expand All @@ -213,14 +205,14 @@ func askForWorkspaceProfile(ctx context.Context) (string, error) {
return profiles[i].Name, nil
}

func askForAccountProfile(ctx context.Context) (string, error) {
path, err := databrickscfg.GetPath()
func AskForAccountProfile(ctx context.Context) (string, error) {
path, err := databrickscfg.GetPath(ctx)
if err != nil {
return "", fmt.Errorf("cannot determine Databricks config file path: %w", err)
}
file, profiles, err := databrickscfg.LoadProfiles(databrickscfg.MatchAccountProfiles)
file, profiles, err := databrickscfg.LoadProfiles(ctx, databrickscfg.MatchAccountProfiles)
if err != nil {
return "", transformLoadError(path, err)
return "", err
}
switch len(profiles) {
case 0:
Expand Down
43 changes: 24 additions & 19 deletions libs/databrickscfg/profiles.go
Original file line number Diff line number Diff line change
@@ -1,11 +1,14 @@
package databrickscfg

import (
"context"
"errors"
"fmt"
"os"
"io/fs"
"path/filepath"
"strings"

"github.com/databricks/cli/libs/env"
"github.com/databricks/databricks-sdk-go/config"
"github.com/spf13/cobra"
)
Expand Down Expand Up @@ -67,43 +70,45 @@ func MatchAllProfiles(p Profile) bool {
}

// Get the path to the .databrickscfg file, falling back to the default in the current user's home directory.
func GetPath() (string, error) {
configFile := os.Getenv("DATABRICKS_CONFIG_FILE")
func GetPath(ctx context.Context) (string, error) {
configFile := env.Get(ctx, "DATABRICKS_CONFIG_FILE")
if configFile == "" {
configFile = "~/.databrickscfg"
}
if strings.HasPrefix(configFile, "~") {
homedir, err := os.UserHomeDir()
if err != nil {
return "", fmt.Errorf("cannot find homedir: %w", err)
}
homedir := env.UserHomeDir(ctx)
configFile = filepath.Join(homedir, configFile[1:])
}
return configFile, nil
}

func Get() (*config.File, error) {
configFile, err := GetPath()
var ErrNoConfiguration = errors.New("no configuration file found")

func Get(ctx context.Context) (*config.File, error) {
path, err := GetPath(ctx)
if err != nil {
return nil, fmt.Errorf("cannot determine Databricks config file path: %w", err)
}
return config.LoadFile(configFile)
configFile, err := config.LoadFile(path)
if errors.Is(err, fs.ErrNotExist) {
// downstreams depend on ErrNoConfiguration. TODO: expose this error through SDK
return nil, fmt.Errorf("%w at %s; please create one first", ErrNoConfiguration, path)
} else if err != nil {
return nil, err
}
return configFile, nil
}

func LoadProfiles(fn ProfileMatchFunction) (file string, profiles Profiles, err error) {
f, err := Get()
func LoadProfiles(ctx context.Context, fn ProfileMatchFunction) (file string, profiles Profiles, err error) {
f, err := Get(ctx)
if err != nil {
return "", nil, fmt.Errorf("cannot load Databricks config file: %w", err)
}

homedir, err := os.UserHomeDir()
if err != nil {
return
}

// Replace homedir with ~ if applicable.
// This is to make the output more readable.
file = f.Path()
file = filepath.Clean(f.Path())
homedir := filepath.Clean(env.UserHomeDir(ctx))
if strings.HasPrefix(file, homedir) {
file = "~" + file[len(homedir):]
}
Expand All @@ -130,7 +135,7 @@ func LoadProfiles(fn ProfileMatchFunction) (file string, profiles Profiles, err
}

func ProfileCompletion(cmd *cobra.Command, args []string, toComplete string) ([]string, cobra.ShellCompDirective) {
_, profiles, err := LoadProfiles(MatchAllProfiles)
_, profiles, err := LoadProfiles(cmd.Context(), MatchAllProfiles)
if err != nil {
return nil, cobra.ShellCompDirectiveError
}
Expand Down
51 changes: 38 additions & 13 deletions libs/databrickscfg/profiles_test.go
Original file line number Diff line number Diff line change
@@ -1,9 +1,11 @@
package databrickscfg

import (
"runtime"
"context"
"path/filepath"
"testing"

"github.com/databricks/cli/libs/env"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
)
Expand All @@ -27,27 +29,50 @@ func TestProfilesSearchCaseInsensitive(t *testing.T) {
}

func TestLoadProfilesReturnsHomedirAsTilde(t *testing.T) {
if runtime.GOOS == "windows" {
t.Setenv("USERPROFILE", "./testdata")
} else {
t.Setenv("HOME", "./testdata")
}
t.Setenv("DATABRICKS_CONFIG_FILE", "./testdata/databrickscfg")
file, _, err := LoadProfiles(func(p Profile) bool { return true })
ctx := context.Background()
ctx = env.WithUserHomeDir(ctx, "testdata")
ctx = env.Set(ctx, "DATABRICKS_CONFIG_FILE", "./testdata/databrickscfg")
file, _, err := LoadProfiles(ctx, func(p Profile) bool { return true })
require.NoError(t, err)
assert.Equal(t, "~/databrickscfg", file)
require.Equal(t, filepath.Clean("~/databrickscfg"), file)
}

func TestLoadProfilesReturnsHomedirAsTildeExoticFile(t *testing.T) {
ctx := context.Background()
ctx = env.WithUserHomeDir(ctx, "testdata")
ctx = env.Set(ctx, "DATABRICKS_CONFIG_FILE", "~/databrickscfg")
file, _, err := LoadProfiles(ctx, func(p Profile) bool { return true })
require.NoError(t, err)
require.Equal(t, filepath.Clean("~/databrickscfg"), file)
}

func TestLoadProfilesReturnsHomedirAsTildeDefaultFile(t *testing.T) {
ctx := context.Background()
ctx = env.WithUserHomeDir(ctx, "testdata/sample-home")
file, _, err := LoadProfiles(ctx, func(p Profile) bool { return true })
require.NoError(t, err)
require.Equal(t, filepath.Clean("~/.databrickscfg"), file)
}

func TestLoadProfilesNoConfiguration(t *testing.T) {
ctx := context.Background()
ctx = env.WithUserHomeDir(ctx, "testdata")
_, _, err := LoadProfiles(ctx, func(p Profile) bool { return true })
require.ErrorIs(t, err, ErrNoConfiguration)
}

func TestLoadProfilesMatchWorkspace(t *testing.T) {
t.Setenv("DATABRICKS_CONFIG_FILE", "./testdata/databrickscfg")
_, profiles, err := LoadProfiles(MatchWorkspaceProfiles)
ctx := context.Background()
ctx = env.Set(ctx, "DATABRICKS_CONFIG_FILE", "./testdata/databrickscfg")
_, profiles, err := LoadProfiles(ctx, MatchWorkspaceProfiles)
require.NoError(t, err)
assert.Equal(t, []string{"DEFAULT", "query", "foo1", "foo2"}, profiles.Names())
}

func TestLoadProfilesMatchAccount(t *testing.T) {
t.Setenv("DATABRICKS_CONFIG_FILE", "./testdata/databrickscfg")
_, profiles, err := LoadProfiles(MatchAccountProfiles)
ctx := context.Background()
ctx = env.Set(ctx, "DATABRICKS_CONFIG_FILE", "./testdata/databrickscfg")
_, profiles, err := LoadProfiles(ctx, MatchAccountProfiles)
require.NoError(t, err)
assert.Equal(t, []string{"acc"}, profiles.Names())
}
7 changes: 7 additions & 0 deletions libs/databrickscfg/testdata/sample-home/.databrickscfg
Original file line number Diff line number Diff line change
@@ -0,0 +1,7 @@
[DEFAULT]
host = https://default
token = default

[acc]
host = https://accounts.cloud.databricks.com
account_id = abc
21 changes: 21 additions & 0 deletions libs/env/context.go
Original file line number Diff line number Diff line change
Expand Up @@ -2,7 +2,9 @@ package env

import (
"context"
"fmt"
"os"
"runtime"
"strings"
)

Expand Down Expand Up @@ -63,6 +65,25 @@ func Set(ctx context.Context, key, value string) context.Context {
return setMap(ctx, m)
}

func homeEnvVar() string {
if runtime.GOOS == "windows" {
return "USERPROFILE"
}
return "HOME"
}

func WithUserHomeDir(ctx context.Context, value string) context.Context {
return Set(ctx, homeEnvVar(), value)
}

func UserHomeDir(ctx context.Context) string {
home := Get(ctx, homeEnvVar())
if home == "" {
panic(fmt.Errorf("$HOME is not set"))
}
return home
}

// All returns environment variables that are defined in both os.Environ
// and this package. `env.Set(ctx, x, y)` will override x from os.Environ.
func All(ctx context.Context) map[string]string {
Expand Down
7 changes: 7 additions & 0 deletions libs/env/context_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -47,3 +47,10 @@ func TestContext(t *testing.T) {
assert.Equal(t, "x=y", all["BAR"])
assert.NotEmpty(t, all["PATH"])
}

func TestHome(t *testing.T) {
ctx := context.Background()
ctx = WithUserHomeDir(ctx, "...")
home := UserHomeDir(ctx)
assert.Equal(t, "...", home)
}
50 changes: 50 additions & 0 deletions libs/env/loader.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,50 @@
package env

import (
"context"

"github.com/databricks/databricks-sdk-go/config"
)

// NewConfigLoader creates Databricks SDK Config loader that is aware of env.Set variables:
//
// ctx = env.Set(ctx, "DATABRICKS_WAREHOUSE_ID", "...")
//
// Usage:
//
// &config.Config{
// Loaders: []config.Loader{
// env.NewConfigLoader(ctx),
// config.ConfigAttributes,
// config.ConfigFile,
// },
// }
func NewConfigLoader(ctx context.Context) *configLoader {
return &configLoader{
ctx: ctx,
}
}

type configLoader struct {
ctx context.Context
}

func (le *configLoader) Name() string {
return "cli-env"
}

func (le *configLoader) Configure(cfg *config.Config) error {
for _, a := range config.ConfigAttributes {
if !a.IsZero(cfg) {
continue
}
for _, k := range a.EnvVars {
v := Get(le.ctx, k)
if v == "" {
continue
}
a.Set(cfg, v)
}
}
return nil
}
Loading