Skip to content

Commit e68a88e

Browse files
nfxnfx
authored
Added env.UserHomeDir(ctx) for parallel-friendly tests (#955)
## Changes `os.Getenv(..)` is not friendly with `libs/env`. This PR makes the relevant changes to places where we need to read user home directory. ## Tests Mainly done in #914
1 parent 7847388 commit e68a88e

File tree

11 files changed

+190
-56
lines changed

11 files changed

+190
-56
lines changed

‎cmd/auth/env.go‎

Lines changed: 4 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,7 @@
11
package auth
22

33
import (
4+
"context"
45
"encoding/json"
56
"errors"
67
"fmt"
@@ -68,8 +69,8 @@ func resolveSection(cfg *config.Config, iniFile *config.File) (*ini.Section, err
6869
return candidates[0], nil
6970
}
7071

71-
func loadFromDatabricksCfg(cfg *config.Config) error {
72-
iniFile, err := databrickscfg.Get()
72+
func loadFromDatabricksCfg(ctx context.Context, cfg *config.Config) error {
73+
iniFile, err := databrickscfg.Get(ctx)
7374
if errors.Is(err, fs.ErrNotExist) {
7475
// it's fine not to have ~/.databrickscfg
7576
return nil
@@ -110,7 +111,7 @@ func newEnvCommand() *cobra.Command {
110111
cfg.Profile = profile
111112
} else if cfg.Host == "" {
112113
cfg.Profile = "DEFAULT"
113-
} else if err := loadFromDatabricksCfg(cfg); err != nil {
114+
} else if err := loadFromDatabricksCfg(cmd.Context(), cfg); err != nil {
114115
return err
115116
}
116117
// Go SDK is lazy loaded because of Terraform semantics,

‎cmd/auth/login.go‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -128,7 +128,7 @@ func newLoginCommand(persistentAuth *auth.PersistentAuth) *cobra.Command {
128128

129129
func setHost(ctx context.Context, profileName string, persistentAuth *auth.PersistentAuth, args []string) error {
130130
// If the chosen profile has a hostname and the user hasn't specified a host, infer the host from the profile.
131-
_, profiles, err := databrickscfg.LoadProfiles(func(p databrickscfg.Profile) bool {
131+
_, profiles, err := databrickscfg.LoadProfiles(ctx, func(p databrickscfg.Profile) bool {
132132
return p.Name == profileName
133133
})
134134
if err != nil {

‎cmd/auth/profiles.go‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -95,7 +95,7 @@ func newProfilesCommand() *cobra.Command {
9595

9696
cmd.RunE = func(cmd *cobra.Command, args []string) error {
9797
var profiles []*profileMetadata
98-
iniFile, err := databrickscfg.Get()
98+
iniFile, err := databrickscfg.Get(cmd.Context())
9999
if os.IsNotExist(err) {
100100
// return empty list for non-configured machines
101101
iniFile = &config.File{

‎cmd/root/auth.go‎

Lines changed: 11 additions & 19 deletions
Original file line numberDiff line numberDiff line change
@@ -5,7 +5,6 @@ import (
55
"errors"
66
"fmt"
77
"net/http"
8-
"os"
98

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

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

125124
// Try picking a profile dynamically if the current configuration is not valid.
126-
profile, err := askForWorkspaceProfile(ctx)
125+
profile, err := AskForWorkspaceProfile(ctx)
127126
if err != nil {
128127
return nil, err
129128
}
@@ -173,21 +172,14 @@ func SetWorkspaceClient(ctx context.Context, w *databricks.WorkspaceClient) cont
173172
return context.WithValue(ctx, &workspaceClient, w)
174173
}
175174

176-
func transformLoadError(path string, err error) error {
177-
if os.IsNotExist(err) {
178-
return fmt.Errorf("no configuration file found at %s; please create one first", path)
179-
}
180-
return err
181-
}
182-
183-
func askForWorkspaceProfile(ctx context.Context) (string, error) {
184-
path, err := databrickscfg.GetPath()
175+
func AskForWorkspaceProfile(ctx context.Context) (string, error) {
176+
path, err := databrickscfg.GetPath(ctx)
185177
if err != nil {
186178
return "", fmt.Errorf("cannot determine Databricks config file path: %w", err)
187179
}
188-
file, profiles, err := databrickscfg.LoadProfiles(databrickscfg.MatchWorkspaceProfiles)
180+
file, profiles, err := databrickscfg.LoadProfiles(ctx, databrickscfg.MatchWorkspaceProfiles)
189181
if err != nil {
190-
return "", transformLoadError(path, err)
182+
return "", err
191183
}
192184
switch len(profiles) {
193185
case 0:
@@ -213,14 +205,14 @@ func askForWorkspaceProfile(ctx context.Context) (string, error) {
213205
return profiles[i].Name, nil
214206
}
215207

216-
func askForAccountProfile(ctx context.Context) (string, error) {
217-
path, err := databrickscfg.GetPath()
208+
func AskForAccountProfile(ctx context.Context) (string, error) {
209+
path, err := databrickscfg.GetPath(ctx)
218210
if err != nil {
219211
return "", fmt.Errorf("cannot determine Databricks config file path: %w", err)
220212
}
221-
file, profiles, err := databrickscfg.LoadProfiles(databrickscfg.MatchAccountProfiles)
213+
file, profiles, err := databrickscfg.LoadProfiles(ctx, databrickscfg.MatchAccountProfiles)
222214
if err != nil {
223-
return "", transformLoadError(path, err)
215+
return "", err
224216
}
225217
switch len(profiles) {
226218
case 0:

‎libs/databrickscfg/profiles.go‎

Lines changed: 24 additions & 19 deletions
Original file line numberDiff line numberDiff line change
@@ -1,11 +1,14 @@
11
package databrickscfg
22

33
import (
4+
"context"
5+
"errors"
46
"fmt"
5-
"os"
7+
"io/fs"
68
"path/filepath"
79
"strings"
810

11+
"github.com/databricks/cli/libs/env"
912
"github.com/databricks/databricks-sdk-go/config"
1013
"github.com/spf13/cobra"
1114
)
@@ -67,43 +70,45 @@ func MatchAllProfiles(p Profile) bool {
6770
}
6871

6972
// Get the path to the .databrickscfg file, falling back to the default in the current user's home directory.
70-
func GetPath() (string, error) {
71-
configFile := os.Getenv("DATABRICKS_CONFIG_FILE")
73+
func GetPath(ctx context.Context) (string, error) {
74+
configFile := env.Get(ctx, "DATABRICKS_CONFIG_FILE")
7275
if configFile == "" {
7376
configFile = "~/.databrickscfg"
7477
}
7578
if strings.HasPrefix(configFile, "~") {
76-
homedir, err := os.UserHomeDir()
77-
if err != nil {
78-
return "", fmt.Errorf("cannot find homedir: %w", err)
79-
}
79+
homedir := env.UserHomeDir(ctx)
8080
configFile = filepath.Join(homedir, configFile[1:])
8181
}
8282
return configFile, nil
8383
}
8484

85-
func Get() (*config.File, error) {
86-
configFile, err := GetPath()
85+
var ErrNoConfiguration = errors.New("no configuration file found")
86+
87+
func Get(ctx context.Context) (*config.File, error) {
88+
path, err := GetPath(ctx)
8789
if err != nil {
8890
return nil, fmt.Errorf("cannot determine Databricks config file path: %w", err)
8991
}
90-
return config.LoadFile(configFile)
92+
configFile, err := config.LoadFile(path)
93+
if errors.Is(err, fs.ErrNotExist) {
94+
// downstreams depend on ErrNoConfiguration. TODO: expose this error through SDK
95+
return nil, fmt.Errorf("%w at %s; please create one first", ErrNoConfiguration, path)
96+
} else if err != nil {
97+
return nil, err
98+
}
99+
return configFile, nil
91100
}
92101

93-
func LoadProfiles(fn ProfileMatchFunction) (file string, profiles Profiles, err error) {
94-
f, err := Get()
102+
func LoadProfiles(ctx context.Context, fn ProfileMatchFunction) (file string, profiles Profiles, err error) {
103+
f, err := Get(ctx)
95104
if err != nil {
96105
return "", nil, fmt.Errorf("cannot load Databricks config file: %w", err)
97106
}
98107

99-
homedir, err := os.UserHomeDir()
100-
if err != nil {
101-
return
102-
}
103-
104108
// Replace homedir with ~ if applicable.
105109
// This is to make the output more readable.
106-
file = f.Path()
110+
file = filepath.Clean(f.Path())
111+
homedir := filepath.Clean(env.UserHomeDir(ctx))
107112
if strings.HasPrefix(file, homedir) {
108113
file = "~" + file[len(homedir):]
109114
}
@@ -130,7 +135,7 @@ func LoadProfiles(fn ProfileMatchFunction) (file string, profiles Profiles, err
130135
}
131136

132137
func ProfileCompletion(cmd *cobra.Command, args []string, toComplete string) ([]string, cobra.ShellCompDirective) {
133-
_, profiles, err := LoadProfiles(MatchAllProfiles)
138+
_, profiles, err := LoadProfiles(cmd.Context(), MatchAllProfiles)
134139
if err != nil {
135140
return nil, cobra.ShellCompDirectiveError
136141
}

‎libs/databrickscfg/profiles_test.go‎

Lines changed: 38 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -1,9 +1,11 @@
11
package databrickscfg
22

33
import (
4-
"runtime"
4+
"context"
5+
"path/filepath"
56
"testing"
67

8+
"github.com/databricks/cli/libs/env"
79
"github.com/stretchr/testify/assert"
810
"github.com/stretchr/testify/require"
911
)
@@ -27,27 +29,50 @@ func TestProfilesSearchCaseInsensitive(t *testing.T) {
2729
}
2830

2931
func TestLoadProfilesReturnsHomedirAsTilde(t *testing.T) {
30-
if runtime.GOOS == "windows" {
31-
t.Setenv("USERPROFILE", "./testdata")
32-
} else {
33-
t.Setenv("HOME", "./testdata")
34-
}
35-
t.Setenv("DATABRICKS_CONFIG_FILE", "./testdata/databrickscfg")
36-
file, _, err := LoadProfiles(func(p Profile) bool { return true })
32+
ctx := context.Background()
33+
ctx = env.WithUserHomeDir(ctx, "testdata")
34+
ctx = env.Set(ctx, "DATABRICKS_CONFIG_FILE", "./testdata/databrickscfg")
35+
file, _, err := LoadProfiles(ctx, func(p Profile) bool { return true })
3736
require.NoError(t, err)
38-
assert.Equal(t, "~/databrickscfg", file)
37+
require.Equal(t, filepath.Clean("~/databrickscfg"), file)
38+
}
39+
40+
func TestLoadProfilesReturnsHomedirAsTildeExoticFile(t *testing.T) {
41+
ctx := context.Background()
42+
ctx = env.WithUserHomeDir(ctx, "testdata")
43+
ctx = env.Set(ctx, "DATABRICKS_CONFIG_FILE", "~/databrickscfg")
44+
file, _, err := LoadProfiles(ctx, func(p Profile) bool { return true })
45+
require.NoError(t, err)
46+
require.Equal(t, filepath.Clean("~/databrickscfg"), file)
47+
}
48+
49+
func TestLoadProfilesReturnsHomedirAsTildeDefaultFile(t *testing.T) {
50+
ctx := context.Background()
51+
ctx = env.WithUserHomeDir(ctx, "testdata/sample-home")
52+
file, _, err := LoadProfiles(ctx, func(p Profile) bool { return true })
53+
require.NoError(t, err)
54+
require.Equal(t, filepath.Clean("~/.databrickscfg"), file)
55+
}
56+
57+
func TestLoadProfilesNoConfiguration(t *testing.T) {
58+
ctx := context.Background()
59+
ctx = env.WithUserHomeDir(ctx, "testdata")
60+
_, _, err := LoadProfiles(ctx, func(p Profile) bool { return true })
61+
require.ErrorIs(t, err, ErrNoConfiguration)
3962
}
4063

4164
func TestLoadProfilesMatchWorkspace(t *testing.T) {
42-
t.Setenv("DATABRICKS_CONFIG_FILE", "./testdata/databrickscfg")
43-
_, profiles, err := LoadProfiles(MatchWorkspaceProfiles)
65+
ctx := context.Background()
66+
ctx = env.Set(ctx, "DATABRICKS_CONFIG_FILE", "./testdata/databrickscfg")
67+
_, profiles, err := LoadProfiles(ctx, MatchWorkspaceProfiles)
4468
require.NoError(t, err)
4569
assert.Equal(t, []string{"DEFAULT", "query", "foo1", "foo2"}, profiles.Names())
4670
}
4771

4872
func TestLoadProfilesMatchAccount(t *testing.T) {
49-
t.Setenv("DATABRICKS_CONFIG_FILE", "./testdata/databrickscfg")
50-
_, profiles, err := LoadProfiles(MatchAccountProfiles)
73+
ctx := context.Background()
74+
ctx = env.Set(ctx, "DATABRICKS_CONFIG_FILE", "./testdata/databrickscfg")
75+
_, profiles, err := LoadProfiles(ctx, MatchAccountProfiles)
5176
require.NoError(t, err)
5277
assert.Equal(t, []string{"acc"}, profiles.Names())
5378
}

‎libs/databrickscfg/testdata/sample-home/.databrickscfg‎

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,7 @@
1+
[DEFAULT]
2+
host = https://default
3+
token = default
4+
5+
[acc]
6+
host = https://accounts.cloud.databricks.com
7+
account_id = abc

‎libs/env/context.go‎

Lines changed: 21 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2,7 +2,9 @@ package env
22

33
import (
44
"context"
5+
"fmt"
56
"os"
7+
"runtime"
68
"strings"
79
)
810

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

68+
func homeEnvVar() string {
69+
if runtime.GOOS == "windows" {
70+
return "USERPROFILE"
71+
}
72+
return "HOME"
73+
}
74+
75+
func WithUserHomeDir(ctx context.Context, value string) context.Context {
76+
return Set(ctx, homeEnvVar(), value)
77+
}
78+
79+
func UserHomeDir(ctx context.Context) string {
80+
home := Get(ctx, homeEnvVar())
81+
if home == "" {
82+
panic(fmt.Errorf("$HOME is not set"))
83+
}
84+
return home
85+
}
86+
6687
// All returns environment variables that are defined in both os.Environ
6788
// and this package. `env.Set(ctx, x, y)` will override x from os.Environ.
6889
func All(ctx context.Context) map[string]string {

‎libs/env/context_test.go‎

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -47,3 +47,10 @@ func TestContext(t *testing.T) {
4747
assert.Equal(t, "x=y", all["BAR"])
4848
assert.NotEmpty(t, all["PATH"])
4949
}
50+
51+
func TestHome(t *testing.T) {
52+
ctx := context.Background()
53+
ctx = WithUserHomeDir(ctx, "...")
54+
home := UserHomeDir(ctx)
55+
assert.Equal(t, "...", home)
56+
}

‎libs/env/loader.go‎

Lines changed: 50 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,50 @@
1+
package env
2+
3+
import (
4+
"context"
5+
6+
"github.com/databricks/databricks-sdk-go/config"
7+
)
8+
9+
// NewConfigLoader creates Databricks SDK Config loader that is aware of env.Set variables:
10+
//
11+
// ctx = env.Set(ctx, "DATABRICKS_WAREHOUSE_ID", "...")
12+
//
13+
// Usage:
14+
//
15+
// &config.Config{
16+
// Loaders: []config.Loader{
17+
// env.NewConfigLoader(ctx),
18+
// config.ConfigAttributes,
19+
// config.ConfigFile,
20+
// },
21+
// }
22+
func NewConfigLoader(ctx context.Context) *configLoader {
23+
return &configLoader{
24+
ctx: ctx,
25+
}
26+
}
27+
28+
type configLoader struct {
29+
ctx context.Context
30+
}
31+
32+
func (le *configLoader) Name() string {
33+
return "cli-env"
34+
}
35+
36+
func (le *configLoader) Configure(cfg *config.Config) error {
37+
for _, a := range config.ConfigAttributes {
38+
if !a.IsZero(cfg) {
39+
continue
40+
}
41+
for _, k := range a.EnvVars {
42+
v := Get(le.ctx, k)
43+
if v == "" {
44+
continue
45+
}
46+
a.Set(cfg, v)
47+
}
48+
}
49+
return nil
50+
}

0 commit comments

Comments
 (0)