From 285502a8ae29a8402074443b9cb749f94795bbc0 Mon Sep 17 00:00:00 2001
From: Johnny Steenbergen <jonathansteenbergen@gmail.com>
Date: Wed, 17 Jun 2020 09:28:41 -0700
Subject: [PATCH] feat(cli): extend Program with config file

closes: #18565
---
 CHANGELOG.md          |  1 +
 kit/cli/viper.go      | 29 ++++++++++++-
 kit/cli/viper_test.go | 96 +++++++++++++++++++++++++++++++++++++++++++
 3 files changed, 125 insertions(+), 1 deletion(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 10b10e41dad..c9a9aefc201 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -8,6 +8,7 @@
 1. [18541](https://github.com/influxdata/influxdb/pull/18541): Pkger allow raw github.com host URLs for yaml|json|jsonnet URLs
 1. [18546](https://github.com/influxdata/influxdb/pull/18546): Influx allow for files to be remotes for all template commands
 1. [18560](https://github.com/influxdata/influxdb/pull/18560): Extend stacks API with update capability
+1. [18568](https://github.com/influxdata/influxdb/pull/18568): Add support for config files to influxd and any cli.NewCommand use case
 
 ## v2.0.0-beta.12 [2020-06-12]
 
diff --git a/kit/cli/viper.go b/kit/cli/viper.go
index c22a9e635b7..7041b262928 100644
--- a/kit/cli/viper.go
+++ b/kit/cli/viper.go
@@ -2,6 +2,7 @@ package cli
 
 import (
 	"fmt"
+	"os"
 	"strings"
 	"time"
 
@@ -52,7 +53,7 @@ type Program struct {
 //
 // This is to simplify the viper/cobra boilerplate.
 func NewCommand(p *Program) *cobra.Command {
-	var cmd = &cobra.Command{
+	cmd := &cobra.Command{
 		Use:  p.Name,
 		Args: cobra.NoArgs,
 		RunE: func(_ *cobra.Command, _ []string) error {
@@ -65,11 +66,37 @@ func NewCommand(p *Program) *cobra.Command {
 	// This normalizes "-" to an underscore in env names.
 	viper.SetEnvKeyReplacer(strings.NewReplacer("-", "_"))
 
+	configFile := viper.GetString("CONFIG_FILE")
+	if configFile == "" {
+		// defaults to looking in same directory as program running for
+		// a file `config.yaml`
+		configFile = "config.yaml"
+	}
+	viper.SetConfigFile(configFile)
+
+	// done before we bind flags to viper keys.
+	// order of precedence (1 highest -> 3 lowest):
+	//	1. flags
+	//  2. env vars
+	//	3. config file
+	if err := initializeConfig(); err != nil {
+		panic(fmt.Sprintf("invalid config file[%s] caused panic: %s", configFile, err))
+	}
 	BindOptions(cmd, p.Opts)
 
 	return cmd
 }
 
+func initializeConfig() error {
+	err := viper.ReadInConfig()
+	if err != nil && !os.IsNotExist(err) {
+		if _, ok := err.(viper.ConfigFileNotFoundError); !ok {
+			return err
+		}
+	}
+	return nil
+}
+
 // BindOptions adds opts to the specified command and automatically
 // registers those options with viper.
 func BindOptions(cmd *cobra.Command, opts []Opt) {
diff --git a/kit/cli/viper_test.go b/kit/cli/viper_test.go
index 3fbbd0738ec..de49af94c93 100644
--- a/kit/cli/viper_test.go
+++ b/kit/cli/viper_test.go
@@ -1,9 +1,15 @@
 package cli
 
 import (
+	"encoding/json"
 	"fmt"
+	"io/ioutil"
 	"os"
+	"path"
+	"testing"
 	"time"
+
+	"github.com/stretchr/testify/require"
 )
 
 type customFlag bool
@@ -101,3 +107,93 @@ func ExampleNewCommand() {
 	// [foo bar]
 	// on
 }
+
+func Test_NewProgram(t *testing.T) {
+	testFilePath, cleanup := newConfigFile(t, map[string]string{
+		"FOO": "bar",
+	})
+	defer cleanup()
+	defer setEnvVar("TEST_CONFIG_FILE", testFilePath)()
+
+	tests := []struct {
+		name      string
+		envVarVal string
+		args      []string
+		expected  string
+	}{
+		{
+			name:     "no vals reads from config",
+			expected: "bar",
+		},
+		{
+			name:      "reads from env var",
+			envVarVal: "foobar",
+			expected:  "foobar",
+		},
+		{
+			name:     "reads from flag",
+			args:     []string{"--foo=baz"},
+			expected: "baz",
+		},
+		{
+			name:      "flag has highest precedence",
+			envVarVal: "foobar",
+			args:      []string{"--foo=baz"},
+			expected:  "baz",
+		},
+	}
+
+	for _, tt := range tests {
+		fn := func(t *testing.T) {
+			if tt.envVarVal != "" {
+				defer setEnvVar("TEST_FOO", tt.envVarVal)()
+			}
+
+			var testVar string
+			program := &Program{
+				Name: "test",
+				Opts: []Opt{
+					{
+						DestP:    &testVar,
+						Flag:     "foo",
+						Required: true,
+					},
+				},
+				Run: func() error { return nil },
+			}
+
+			cmd := NewCommand(program)
+			cmd.SetArgs(append([]string{}, tt.args...))
+			require.NoError(t, cmd.Execute())
+
+			require.Equal(t, tt.expected, testVar)
+		}
+
+		t.Run(tt.name, fn)
+	}
+}
+
+func setEnvVar(key, val string) func() {
+	old := os.Getenv(key)
+	os.Setenv(key, val)
+	return func() {
+		os.Setenv(key, old)
+	}
+}
+
+func newConfigFile(t *testing.T, config interface{}) (string, func()) {
+	t.Helper()
+
+	testDir, err := ioutil.TempDir("", "")
+	require.NoError(t, err)
+
+	b, err := json.Marshal(config)
+	require.NoError(t, err)
+
+	testFile := path.Join(testDir, "config.json")
+	require.NoError(t, ioutil.WriteFile(testFile, b, os.ModePerm))
+
+	return testFile, func() {
+		os.RemoveAll(testDir)
+	}
+}