From 04ae010378509ef6f9899354b3b1550aac942b29 Mon Sep 17 00:00:00 2001 From: Ben Johnson Date: Fri, 16 Apr 2021 09:28:01 -0600 Subject: [PATCH] Configuration file environment variable expansion This commit adds simple variable expansion using either `$FOO` or `${FOO}` when evaluating the config file. This can be disabled by any command by using the `-no-expand-env` flag. --- cmd/litestream/databases.go | 7 +++-- cmd/litestream/generations.go | 7 +++-- cmd/litestream/main.go | 22 +++++++++++---- cmd/litestream/main_test.go | 53 ++++++++++++++++++++++++++++++++++- cmd/litestream/replicate.go | 7 +++-- cmd/litestream/restore.go | 11 +++++--- cmd/litestream/snapshots.go | 7 +++-- cmd/litestream/wal.go | 7 +++-- 8 files changed, 100 insertions(+), 21 deletions(-) diff --git a/cmd/litestream/databases.go b/cmd/litestream/databases.go index dc32595..236c01e 100644 --- a/cmd/litestream/databases.go +++ b/cmd/litestream/databases.go @@ -15,7 +15,7 @@ type DatabasesCommand struct{} // Run executes the command. func (c *DatabasesCommand) Run(ctx context.Context, args []string) (err error) { fs := flag.NewFlagSet("litestream-databases", flag.ContinueOnError) - configPath := registerConfigFlag(fs) + configPath, noExpandEnv := registerConfigFlag(fs) fs.Usage = c.Usage if err := fs.Parse(args); err != nil { return err @@ -27,7 +27,7 @@ func (c *DatabasesCommand) Run(ctx context.Context, args []string) (err error) { if *configPath == "" { *configPath = DefaultConfigPath() } - config, err := ReadConfigFile(*configPath) + config, err := ReadConfigFile(*configPath, !*noExpandEnv) if err != nil { return err } @@ -72,6 +72,9 @@ Arguments: Specifies the configuration file. Defaults to %s + -no-expand-env + Disables environment variable expansion in configuration file. + `[1:], DefaultConfigPath(), ) diff --git a/cmd/litestream/generations.go b/cmd/litestream/generations.go index a6a5e2d..57a778e 100644 --- a/cmd/litestream/generations.go +++ b/cmd/litestream/generations.go @@ -18,7 +18,7 @@ type GenerationsCommand struct{} // Run executes the command. func (c *GenerationsCommand) Run(ctx context.Context, args []string) (err error) { fs := flag.NewFlagSet("litestream-generations", flag.ContinueOnError) - configPath := registerConfigFlag(fs) + configPath, noExpandEnv := registerConfigFlag(fs) replicaName := fs.String("replica", "", "replica name") fs.Usage = c.Usage if err := fs.Parse(args); err != nil { @@ -45,7 +45,7 @@ func (c *GenerationsCommand) Run(ctx context.Context, args []string) (err error) } // Load configuration. - config, err := ReadConfigFile(*configPath) + config, err := ReadConfigFile(*configPath, !*noExpandEnv) if err != nil { return err } @@ -131,6 +131,9 @@ Arguments: Specifies the configuration file. Defaults to %s + -no-expand-env + Disables environment variable expansion in configuration file. + -replica NAME Optional, filters by replica. diff --git a/cmd/litestream/main.go b/cmd/litestream/main.go index 5d0374b..ecafa76 100644 --- a/cmd/litestream/main.go +++ b/cmd/litestream/main.go @@ -179,7 +179,8 @@ func (c *Config) DBConfig(path string) *DBConfig { } // ReadConfigFile unmarshals config from filename. Expands path if needed. -func ReadConfigFile(filename string) (_ Config, err error) { +// If expandEnv is true then environment variables are expanded in the config. +func ReadConfigFile(filename string, expandEnv bool) (_ Config, err error) { config := DefaultConfig() // Expand filename, if necessary. @@ -188,12 +189,20 @@ func ReadConfigFile(filename string) (_ Config, err error) { return config, err } - // Read & deserialize configuration. - if buf, err := ioutil.ReadFile(filename); os.IsNotExist(err) { + // Read configuration. + buf, err := ioutil.ReadFile(filename) + if os.IsNotExist(err) { return config, fmt.Errorf("config file not found: %s", filename) } else if err != nil { return config, err - } else if err := yaml.Unmarshal(buf, &config); err != nil { + } + + // Expand environment variables, if enabled. + if expandEnv { + buf = []byte(os.ExpandEnv(string(buf))) + } + + if err := yaml.Unmarshal(buf, &config); err != nil { return config, err } @@ -461,8 +470,9 @@ func DefaultConfigPath() string { return defaultConfigPath } -func registerConfigFlag(fs *flag.FlagSet) *string { - return fs.String("config", "", "config path") +func registerConfigFlag(fs *flag.FlagSet) (configPath *string, noExpandEnv *bool) { + return fs.String("config", "", "config path"), + fs.Bool("no-expand-env", false, "do not expand env vars in config") } // expand returns an absolute path for s. diff --git a/cmd/litestream/main_test.go b/cmd/litestream/main_test.go index a1c89f0..bd30c28 100644 --- a/cmd/litestream/main_test.go +++ b/cmd/litestream/main_test.go @@ -2,6 +2,7 @@ package main_test import ( "io/ioutil" + "os" "path/filepath" "testing" @@ -26,7 +27,7 @@ dbs: t.Fatal(err) } - config, err := main.ReadConfigFile(filename) + config, err := main.ReadConfigFile(filename, true) if err != nil { t.Fatal(err) } else if got, want := config.AccessKeyID, `XXX`; got != want { @@ -39,6 +40,56 @@ dbs: t.Fatalf("Replica.SecretAccessKey=%v, want %v", got, want) } }) + + // Ensure environment variables are expanded. + t.Run("ExpandEnv", func(t *testing.T) { + os.Setenv("LITESTREAM_TEST_0129380", "/path/to/db") + os.Setenv("LITESTREAM_TEST_1872363", "s3://foo/bar") + + filename := filepath.Join(t.TempDir(), "litestream.yml") + if err := ioutil.WriteFile(filename, []byte(` +dbs: + - path: $LITESTREAM_TEST_0129380 + replicas: + - url: ${LITESTREAM_TEST_1872363} + - url: ${LITESTREAM_TEST_NO_SUCH_ENV} +`[1:]), 0666); err != nil { + t.Fatal(err) + } + + config, err := main.ReadConfigFile(filename, true) + if err != nil { + t.Fatal(err) + } else if got, want := config.DBs[0].Path, `/path/to/db`; got != want { + t.Fatalf("DB.Path=%v, want %v", got, want) + } else if got, want := config.DBs[0].Replicas[0].URL, `s3://foo/bar`; got != want { + t.Fatalf("Replica[0].URL=%v, want %v", got, want) + } else if got, want := config.DBs[0].Replicas[1].URL, ``; got != want { + t.Fatalf("Replica[1].URL=%v, want %v", got, want) + } + }) + + // Ensure environment variables are not expanded. + t.Run("NoExpandEnv", func(t *testing.T) { + os.Setenv("LITESTREAM_TEST_9847533", "s3://foo/bar") + + filename := filepath.Join(t.TempDir(), "litestream.yml") + if err := ioutil.WriteFile(filename, []byte(` +dbs: + - path: /path/to/db + replicas: + - url: ${LITESTREAM_TEST_9847533} +`[1:]), 0666); err != nil { + t.Fatal(err) + } + + config, err := main.ReadConfigFile(filename, false) + if err != nil { + t.Fatal(err) + } else if got, want := config.DBs[0].Replicas[0].URL, `${LITESTREAM_TEST_9847533}`; got != want { + t.Fatalf("Replica.URL=%v, want %v", got, want) + } + }) } func TestNewFileReplicaFromConfig(t *testing.T) { diff --git a/cmd/litestream/replicate.go b/cmd/litestream/replicate.go index f748ae1..2f6584e 100644 --- a/cmd/litestream/replicate.go +++ b/cmd/litestream/replicate.go @@ -32,7 +32,7 @@ func NewReplicateCommand() *ReplicateCommand { func (c *ReplicateCommand) ParseFlags(ctx context.Context, args []string) (err error) { fs := flag.NewFlagSet("litestream-replicate", flag.ContinueOnError) tracePath := fs.String("trace", "", "trace path") - configPath := registerConfigFlag(fs) + configPath, noExpandEnv := registerConfigFlag(fs) fs.Usage = c.Usage if err := fs.Parse(args); err != nil { return err @@ -58,7 +58,7 @@ func (c *ReplicateCommand) ParseFlags(ctx context.Context, args []string) (err e if *configPath == "" { *configPath = DefaultConfigPath() } - if c.Config, err = ReadConfigFile(*configPath); err != nil { + if c.Config, err = ReadConfigFile(*configPath, !*noExpandEnv); err != nil { return err } } @@ -168,6 +168,9 @@ Arguments: Specifies the configuration file. Defaults to %s + -no-expand-env + Disables environment variable expansion in configuration file. + -trace PATH Write verbose trace logging to PATH. diff --git a/cmd/litestream/restore.go b/cmd/litestream/restore.go index 66c5d00..e35733d 100644 --- a/cmd/litestream/restore.go +++ b/cmd/litestream/restore.go @@ -21,7 +21,7 @@ func (c *RestoreCommand) Run(ctx context.Context, args []string) (err error) { opt.Verbose = true fs := flag.NewFlagSet("litestream-restore", flag.ContinueOnError) - configPath := registerConfigFlag(fs) + configPath, noExpandEnv := registerConfigFlag(fs) fs.StringVar(&opt.OutputPath, "o", "", "output path") fs.StringVar(&opt.ReplicaName, "replica", "", "replica name") fs.StringVar(&opt.Generation, "generation", "", "generation name") @@ -69,7 +69,7 @@ func (c *RestoreCommand) Run(ctx context.Context, args []string) (err error) { if *configPath == "" { *configPath = DefaultConfigPath() } - if r, err = c.loadFromConfig(ctx, fs.Arg(0), *configPath, &opt); err != nil { + if r, err = c.loadFromConfig(ctx, fs.Arg(0), *configPath, !*noExpandEnv, &opt); err != nil { return err } } @@ -98,9 +98,9 @@ func (c *RestoreCommand) loadFromURL(ctx context.Context, replicaURL string, opt } // loadFromConfig returns a replica & updates the restore options from a DB reference. -func (c *RestoreCommand) loadFromConfig(ctx context.Context, dbPath, configPath string, opt *litestream.RestoreOptions) (litestream.Replica, error) { +func (c *RestoreCommand) loadFromConfig(ctx context.Context, dbPath, configPath string, expandEnv bool, opt *litestream.RestoreOptions) (litestream.Replica, error) { // Load configuration. - config, err := ReadConfigFile(configPath) + config, err := ReadConfigFile(configPath, expandEnv) if err != nil { return nil, err } @@ -150,6 +150,9 @@ Arguments: Specifies the configuration file. Defaults to %s + -no-expand-env + Disables environment variable expansion in configuration file. + -replica NAME Restore from a specific replica. Defaults to replica with latest data. diff --git a/cmd/litestream/snapshots.go b/cmd/litestream/snapshots.go index 742e7cb..920ddb5 100644 --- a/cmd/litestream/snapshots.go +++ b/cmd/litestream/snapshots.go @@ -17,7 +17,7 @@ type SnapshotsCommand struct{} // Run executes the command. func (c *SnapshotsCommand) Run(ctx context.Context, args []string) (err error) { fs := flag.NewFlagSet("litestream-snapshots", flag.ContinueOnError) - configPath := registerConfigFlag(fs) + configPath, noExpandEnv := registerConfigFlag(fs) replicaName := fs.String("replica", "", "replica name") fs.Usage = c.Usage if err := fs.Parse(args); err != nil { @@ -43,7 +43,7 @@ func (c *SnapshotsCommand) Run(ctx context.Context, args []string) (err error) { } // Load configuration. - config, err := ReadConfigFile(*configPath) + config, err := ReadConfigFile(*configPath, !*noExpandEnv) if err != nil { return err } @@ -112,6 +112,9 @@ Arguments: Specifies the configuration file. Defaults to %s + -no-expand-env + Disables environment variable expansion in configuration file. + -replica NAME Optional, filter by a specific replica. diff --git a/cmd/litestream/wal.go b/cmd/litestream/wal.go index 467c740..bd60a28 100644 --- a/cmd/litestream/wal.go +++ b/cmd/litestream/wal.go @@ -17,7 +17,7 @@ type WALCommand struct{} // Run executes the command. func (c *WALCommand) Run(ctx context.Context, args []string) (err error) { fs := flag.NewFlagSet("litestream-wal", flag.ContinueOnError) - configPath := registerConfigFlag(fs) + configPath, noExpandEnv := registerConfigFlag(fs) replicaName := fs.String("replica", "", "replica name") generation := fs.String("generation", "", "generation name") fs.Usage = c.Usage @@ -44,7 +44,7 @@ func (c *WALCommand) Run(ctx context.Context, args []string) (err error) { } // Load configuration. - config, err := ReadConfigFile(*configPath) + config, err := ReadConfigFile(*configPath, !*noExpandEnv) if err != nil { return err } @@ -118,6 +118,9 @@ Arguments: Specifies the configuration file. Defaults to %s + -no-expand-env + Disables environment variable expansion in configuration file. + -replica NAME Optional, filter by a specific replica.