feat: implements support for environment variable substitutions in co…

…nfig file (#1283)

* feat: implements support for env substitutions in config

* improvement

* improvements

* fix for $1 strings

pkg/cli/command.go

@@ -3,7 +3,9 @@ package cli
 import (
 	"errors"
 	"fmt"
+	"io/ioutil"
 	"os"
+	"regexp"
 	"strings"
 
 	"github.com/spf13/cobra"
@@ -159,10 +161,30 @@ func loadConfigFile(cmd *cobra.Command, vpr *viper.Viper) error {
 		return nil
 	}
 	vpr.SetConfigFile(configPath)
-	err := vpr.ReadInConfig()
-	if err == nil || (errors.Is(err, os.ErrNotExist) && !userDefined) {
-		// The default config file can be missing.
+	data, err := ioutil.ReadFile(configPath)
+	if err != nil && errors.Is(err, os.ErrNotExist) && !userDefined {
+		// If user did not specify the config file, and the file does not exist,
+		//   this is okay
 		return nil
 	}
+
+	if err == nil {
+		return vpr.ReadConfig(strings.NewReader(performSubstitutions(data)))
+	}
+
 	return fmt.Errorf("loading configuration file: %w", err)
 }
+
+var digitCheck = regexp.MustCompile(`^[0-9]`)
+
+func performSubstitutions(data []byte) string {
+	// return string(data)
+	return os.Expand(string(data), func(name string) string {
+		// this is here so that $1, $2, etc. work in the config file
+		if digitCheck.MatchString(name) {
+			return "$" + name
+		}
+		s := os.Getenv(name)
+		return s
+	})
+}

pkg/cli/flags_test.go

@@ -76,38 +76,60 @@ var _ = Describe("flags", func() {
 				cfg := FlagsStruct{}
 				vpr := viper.New()
 				exampleCommand := &cobra.Command{
-					RunE: func(cmd *cobra.Command, args []string) error {
-						if cfg.Config != "" {
-							// Use config file from the flag.
-							vpr.SetConfigFile(cfg.Config)
+					RunE: CreateCmdRunFn(&cfg, vpr, func(cmd *cobra.Command, args []string) error {
+						return nil
+					}),
+				}
 
-							// If a config file is found, read it in.
-							if err := vpr.ReadInConfig(); err == nil {
-								fmt.Fprintln(os.Stderr, "Using config file:", vpr.ConfigFileUsed())
-							}
+				PopulateFlagSet(&cfg, exampleCommand.Flags(), vpr)
+				vpr.BindPFlags(exampleCommand.Flags())
 
-							if err := Unmarshal(vpr, &cfg); err != nil {
-								fmt.Fprintln(os.Stderr, "Unable to unmarshal:", err)
-							}
+				b := bytes.NewBufferString("")
+				exampleCommand.SetOut(b)
+				exampleCommand.SetArgs([]string{fmt.Sprintf("--config=%s", "testdata/example.yml")})
 
-							fmt.Printf("configuration is %+v \n", cfg)
-						}
+				err := exampleCommand.Execute()
+				Expect(err).ToNot(HaveOccurred())
+				Expect(cfg.Foo).To(Equal("test-val-1"))
+				Expect(cfg.Foos).To(Equal([]string{"test-val-2", "test-val-3"}))
+				Expect(cfg.Bar).To(Equal(123))
+				Expect(cfg.Baz).To(Equal(10 * time.Hour))
+				Expect(cfg.FooBar).To(Equal("test-val-4"))
+				Expect(cfg.FooFoo).To(Equal(10.23))
+				Expect(cfg.FooBytes).To(Equal(100 * bytesize.MB))
+				Expect(cfg.FooDur).To(Equal(5*time.Minute + 23*time.Second))
+			})
 
+			It("correctly works with substitutions", func() {
+				os.Setenv("VALUE1", "test-val-1")
+				os.Setenv("VALUE2", "test-val-2")
+				// os.Setenv("VALUE3", "test-val-3")
+				os.Setenv("VALUE4", "123")
+				os.Setenv("VALUE5", "10h")
+				os.Setenv("VALUE6", "test-val-4")
+				os.Setenv("VALUE7", "10.23")
+				os.Setenv("VALUE8", "100mb")
+				os.Setenv("VALUE9", "5m23s")
+				cfg := FlagsStruct{}
+				vpr := viper.New()
+
+				exampleCommand := &cobra.Command{
+					RunE: CreateCmdRunFn(&cfg, vpr, func(cmd *cobra.Command, args []string) error {
 						return nil
-					},
+					}),
 				}
 
 				PopulateFlagSet(&cfg, exampleCommand.Flags(), vpr)
 				vpr.BindPFlags(exampleCommand.Flags())
 
 				b := bytes.NewBufferString("")
 				exampleCommand.SetOut(b)
-				exampleCommand.SetArgs([]string{fmt.Sprintf("--config=%s", "testdata/example.yml")})
+				exampleCommand.SetArgs([]string{fmt.Sprintf("--config=%s", "testdata/substitutions.yml")})
 
 				err := exampleCommand.Execute()
 				Expect(err).ToNot(HaveOccurred())
 				Expect(cfg.Foo).To(Equal("test-val-1"))
-				Expect(cfg.Foos).To(Equal([]string{"test-val-2", "test-val-3"}))
+				Expect(cfg.Foos).To(Equal([]string{"test-val-2", ""}))
 				Expect(cfg.Bar).To(Equal(123))
 				Expect(cfg.Baz).To(Equal(10 * time.Hour))
 				Expect(cfg.FooBar).To(Equal("test-val-4"))
@@ -161,26 +183,11 @@ var _ = Describe("flags", func() {
 				var cfg config.Server
 				vpr := viper.New()
 				exampleCommand := &cobra.Command{
-					RunE: func(cmd *cobra.Command, args []string) error {
-						if cfg.Config != "" {
-							// Use config file from the flag.
-							vpr.SetConfigFile(cfg.Config)
-
-							// If a config file is found, read it in.
-							if err := vpr.ReadInConfig(); err == nil {
-								fmt.Fprintln(os.Stderr, "Using config file:", vpr.ConfigFileUsed())
-							}
-
-							if err := Unmarshal(vpr, &cfg); err != nil {
-								fmt.Fprintln(os.Stderr, "Unable to unmarshal:", err)
-							}
-
-							Expect(loadScrapeConfigsFromFile(&cfg)).ToNot(HaveOccurred())
-							fmt.Printf("configuration is %+v \n", cfg)
-						}
-
+					RunE: CreateCmdRunFn(&cfg, vpr, func(cmd *cobra.Command, args []string) error {
+						Expect(loadScrapeConfigsFromFile(&cfg)).ToNot(HaveOccurred())
+						fmt.Printf("configuration is %+v \n", cfg)
 						return nil
-					},
+					}),
 				}
 
 				PopulateFlagSet(&cfg, exampleCommand.Flags(), vpr, WithSkip("scrape-configs"))

pkg/cli/server.go

@@ -6,6 +6,7 @@ import (
 	"os"
 	"time"
 
+	"github.com/davecgh/go-spew/spew"
 	"github.com/prometheus/client_golang/prometheus"
 	"github.com/pyroscope-io/client/pyroscope"
 	"github.com/sirupsen/logrus"
@@ -227,6 +228,11 @@ func newServerService(c *config.Server) (*serverService, error) {
 		svc.analyticsService = analytics.NewService(c, svc.storage, svc.controller)
 	}
 
+	if os.Getenv("PYROSCOPE_CONFIG_DEBUG") != "" {
+		fmt.Println("parsed config:")
+		spew.Dump(svc.config)
+	}
+
 	return &svc, nil
 }
 
@@ -383,7 +389,7 @@ func loadScrapeConfigsFromFile(c *config.Server) error {
 		ScrapeConfigs []*sc.Config `yaml:"scrape-configs" mapstructure:"-"`
 	}
 	var s scrapeConfig
-	if err = yaml.Unmarshal(b, &s); err != nil {
+	if err = yaml.Unmarshal([]byte(performSubstitutions(b)), &s); err != nil {
 		return err
 	}
 	// Populate scrape configs.
@@ -408,7 +414,7 @@ func loadRemoteWriteTargetConfigsFromFile(c *config.Server) error {
 	}
 
 	var s cfg
-	if err = yaml.Unmarshal(b, &s); err != nil {
+	if err = yaml.Unmarshal([]byte(performSubstitutions(b)), &s); err != nil {
 		return err
 	}
 

pkg/cli/testdata/substitutions.yml

@@ -0,0 +1,11 @@
+---
+foo: '${VALUE1}'
+foos:
+  - '${VALUE2}'
+  - '${VALUE3}'
+bar: ${VALUE4}
+baz: '${VALUE5}'
+foo-bar: '${VALUE6}'
+foo-foo: ${VALUE7}
+foo-bytes: '${VALUE8}'
+foo-dur: '${VALUE9}'