Skip to content

Commit

Permalink
Added env.UserHomeDir(ctx) for parallel-friendly tests (#955)
Browse files Browse the repository at this point in the history
## 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
  • Loading branch information
nfx authored Nov 8, 2023
1 parent 7847388 commit e68a88e
Show file tree
Hide file tree
Showing 11 changed files with 190 additions and 56 deletions.
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

0 comments on commit e68a88e

Please sign in to comment.