mirror of https://github.com/databricks/cli.git
107 lines
2.3 KiB
Go
107 lines
2.3 KiB
Go
package env
|
|
|
|
import (
|
|
"context"
|
|
"errors"
|
|
"os"
|
|
"runtime"
|
|
"strings"
|
|
)
|
|
|
|
var envContextKey int
|
|
|
|
func copyMap(m map[string]string) map[string]string {
|
|
out := make(map[string]string, len(m))
|
|
for k, v := range m {
|
|
out[k] = v
|
|
}
|
|
return out
|
|
}
|
|
|
|
func getMap(ctx context.Context) map[string]string {
|
|
if ctx == nil {
|
|
return nil
|
|
}
|
|
m, ok := ctx.Value(&envContextKey).(map[string]string)
|
|
if !ok {
|
|
return nil
|
|
}
|
|
return m
|
|
}
|
|
|
|
func setMap(ctx context.Context, m map[string]string) context.Context {
|
|
return context.WithValue(ctx, &envContextKey, m)
|
|
}
|
|
|
|
// Lookup key in the context or the the environment.
|
|
// Context has precedence.
|
|
func Lookup(ctx context.Context, key string) (string, bool) {
|
|
m := getMap(ctx)
|
|
|
|
// Return if the key is set in the context.
|
|
v, ok := m[key]
|
|
if ok {
|
|
return v, true
|
|
}
|
|
|
|
// Fall back to the environment.
|
|
return os.LookupEnv(key)
|
|
}
|
|
|
|
// Get key from the context or the environment.
|
|
// Context has precedence.
|
|
func Get(ctx context.Context, key string) string {
|
|
v, _ := Lookup(ctx, key)
|
|
return v
|
|
}
|
|
|
|
// Set key on the context.
|
|
//
|
|
// Note: this does NOT mutate the processes' actual environment variables.
|
|
// It is only visible to other code that uses this package.
|
|
func Set(ctx context.Context, key, value string) context.Context {
|
|
m := copyMap(getMap(ctx))
|
|
m[key] = value
|
|
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)
|
|
}
|
|
|
|
// ErrNoHomeEnv indicates the absence of $HOME env variable
|
|
var ErrNoHomeEnv = errors.New("$HOME is not set")
|
|
|
|
func UserHomeDir(ctx context.Context) (string, error) {
|
|
home := Get(ctx, homeEnvVar())
|
|
if home == "" {
|
|
return "", ErrNoHomeEnv
|
|
}
|
|
return home, nil
|
|
}
|
|
|
|
// 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 {
|
|
m := map[string]string{}
|
|
for _, line := range os.Environ() {
|
|
split := strings.SplitN(line, "=", 2)
|
|
if len(split) != 2 {
|
|
continue
|
|
}
|
|
m[split[0]] = split[1]
|
|
}
|
|
// override existing environment variables with the ones we set
|
|
for k, v := range getMap(ctx) {
|
|
m[k] = v
|
|
}
|
|
return m
|
|
}
|