2020-02-10 12:26:18 +01:00
|
|
|
package envflag
|
|
|
|
|
|
|
|
import (
|
|
|
|
"flag"
|
2020-02-10 15:08:04 +01:00
|
|
|
"log"
|
2020-02-10 12:26:18 +01:00
|
|
|
"os"
|
2020-02-24 20:14:22 +01:00
|
|
|
"strings"
|
2022-10-26 00:52:05 +02:00
|
|
|
|
|
|
|
"github.com/VictoriaMetrics/VictoriaMetrics/lib/envtemplate"
|
2020-02-10 12:26:18 +01:00
|
|
|
)
|
|
|
|
|
2020-03-30 14:51:19 +02:00
|
|
|
var (
|
|
|
|
enable = flag.Bool("envflag.enable", false, "Whether to enable reading flags from environment variables additionally to command line. "+
|
|
|
|
"Command line flag values have priority over values from environment vars. "+
|
2021-08-11 09:29:33 +02:00
|
|
|
"Flags are read only from command line if this flag isn't set. See https://docs.victoriametrics.com/#environment-variables for more details")
|
2020-03-30 14:51:19 +02:00
|
|
|
prefix = flag.String("envflag.prefix", "", "Prefix for environment variables if -envflag.enable is set")
|
|
|
|
)
|
2020-02-10 14:58:30 +01:00
|
|
|
|
2020-02-10 12:26:18 +01:00
|
|
|
// Parse parses environment vars and command-line flags.
|
|
|
|
//
|
|
|
|
// Flags set via command-line override flags set via environment vars.
|
|
|
|
//
|
|
|
|
// This function must be called instead of flag.Parse() before using any flags in the program.
|
|
|
|
func Parse() {
|
2022-10-26 00:52:05 +02:00
|
|
|
// Substitute %{ENV_VAR} inside args with the corresponding environment variable values
|
|
|
|
args := os.Args[1:]
|
|
|
|
dstArgs := args[:0]
|
|
|
|
for _, arg := range args {
|
2022-10-26 13:49:20 +02:00
|
|
|
s, err := envtemplate.ReplaceString(arg)
|
2022-10-26 00:52:05 +02:00
|
|
|
if err != nil {
|
2022-10-26 13:49:20 +02:00
|
|
|
// Do not use lib/logger here, since it is uninitialized yet.
|
2022-10-26 00:52:05 +02:00
|
|
|
log.Fatalf("cannot process arg %q: %s", arg, err)
|
|
|
|
}
|
2022-10-26 13:49:20 +02:00
|
|
|
if len(s) > 0 {
|
|
|
|
dstArgs = append(dstArgs, s)
|
2022-10-26 00:52:05 +02:00
|
|
|
}
|
|
|
|
}
|
|
|
|
os.Args = os.Args[:1+len(dstArgs)]
|
|
|
|
|
|
|
|
// Parse flags
|
2020-02-10 12:26:18 +01:00
|
|
|
flag.Parse()
|
2020-02-10 14:58:30 +01:00
|
|
|
if !*enable {
|
|
|
|
return
|
|
|
|
}
|
2020-02-10 12:26:18 +01:00
|
|
|
|
|
|
|
// Remember explicitly set command-line flags.
|
|
|
|
flagsSet := make(map[string]bool)
|
|
|
|
flag.Visit(func(f *flag.Flag) {
|
|
|
|
flagsSet[f.Name] = true
|
|
|
|
})
|
|
|
|
|
|
|
|
// Obtain the remaining flag values from environment vars.
|
|
|
|
flag.VisitAll(func(f *flag.Flag) {
|
|
|
|
if flagsSet[f.Name] {
|
|
|
|
// The flag is explicitly set via command-line.
|
|
|
|
return
|
|
|
|
}
|
2020-02-10 14:58:30 +01:00
|
|
|
// Get flag value from environment var.
|
2020-02-24 20:14:22 +01:00
|
|
|
fname := getEnvFlagName(f.Name)
|
2022-10-26 13:49:20 +02:00
|
|
|
if v, ok := envtemplate.LookupEnv(fname); ok {
|
2021-10-19 23:41:02 +02:00
|
|
|
if err := flag.Set(f.Name, v); err != nil {
|
2020-02-10 15:08:04 +01:00
|
|
|
// Do not use lib/logger here, since it is uninitialized yet.
|
2020-02-24 20:14:22 +01:00
|
|
|
log.Fatalf("cannot set flag %s to %q, which is read from environment variable %q: %s", f.Name, v, fname, err)
|
2020-02-10 15:08:04 +01:00
|
|
|
}
|
2020-02-10 12:26:18 +01:00
|
|
|
}
|
|
|
|
})
|
|
|
|
}
|
2020-02-24 20:14:22 +01:00
|
|
|
|
|
|
|
func getEnvFlagName(s string) string {
|
|
|
|
// Substitute dots with underscores, since env var names cannot contain dots.
|
|
|
|
// See https://github.com/VictoriaMetrics/VictoriaMetrics/issues/311#issuecomment-586354129 for details.
|
2020-03-30 14:51:19 +02:00
|
|
|
s = strings.ReplaceAll(s, ".", "_")
|
|
|
|
return *prefix + s
|
2020-02-24 20:14:22 +01:00
|
|
|
}
|