mirror of https://github.com/databricks/cli.git
64 lines
1.3 KiB
Go
64 lines
1.3 KiB
Go
|
package env
|
||
|
|
|
||
|
|
import (
|
||
|
|
"context"
|
||
|
|
"os"
|
||
|
|
)
|
||
|
|
|
||
|
|
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)
|
||
|
|
}
|