databricks-cli/libs/env/context.go

107 lines
2.3 KiB
Go
Raw Normal View History

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
}