Merge pull request #9154 from ajm188/vtadmin-viper

[vtadmin] viper config support

go.mod

@@ -111,7 +111,6 @@ require (
 	gopkg.in/gcfg.v1 v1.2.3
 	gopkg.in/ldap.v2 v2.5.0
 	gopkg.in/warnings.v0 v0.1.2 // indirect
-	gopkg.in/yaml.v2 v2.4.0
 	gotest.tools v2.2.0+incompatible
 	honnef.co/go/tools v0.0.1-2020.1.4
 	k8s.io/apiextensions-apiserver v0.18.19
@@ -198,6 +197,7 @@ require (
 	google.golang.org/appengine v1.6.7 // indirect
 	gopkg.in/inf.v0 v0.9.1 // indirect
 	gopkg.in/ini.v1 v1.62.0 // indirect
+	gopkg.in/yaml.v2 v2.4.0 // indirect
 	gopkg.in/yaml.v3 v3.0.0-20210107192922-496545a6307b // indirect
 	k8s.io/api v0.20.6 // indirect
 	k8s.io/gengo v0.0.0-20200413195148-3a45101e95ac // indirect

go/vt/vtadmin/cluster/config.go

@@ -87,14 +87,7 @@ func (cfg *Config) Set(value string) error {
 	return parseFlag(cfg, value)
 }
 
-// UnmarshalYAML implements the yaml.Unmarshaler interface.
-func (cfg *Config) UnmarshalYAML(unmarshal func(interface{}) error) error {
-	attributes := map[string]string{}
-
-	if err := unmarshal(attributes); err != nil {
-		return err
-	}
-
+func (cfg *Config) unmarshalMap(attributes map[string]string) error {
 	for k, v := range attributes {
 		if err := parseOne(cfg, k, v); err != nil {
 			return err

go/vt/vtadmin/cluster/config_test.go

@@ -20,8 +20,6 @@ import (
 	"testing"
 
 	"github.com/stretchr/testify/assert"
-	"github.com/stretchr/testify/require"
-	"gopkg.in/yaml.v2"
 )
 
 func TestMergeConfig(t *testing.T) {
@@ -143,75 +141,3 @@ func TestMergeConfig(t *testing.T) {
 		})
 	}
 }
-
-func TestConfigUnmarshalYAML(t *testing.T) {
-	t.Parallel()
-
-	tests := []struct {
-		name   string
-		yaml   string
-		config Config
-		err    error
-	}{
-		{
-			name: "simple",
-			yaml: `name: cluster1
-id: c1`,
-			config: Config{
-				ID:                   "c1",
-				Name:                 "cluster1",
-				DiscoveryFlagsByImpl: map[string]map[string]string{},
-			},
-			err: nil,
-		},
-		{
-			name: "discovery flags",
-			yaml: `name: cluster1
-id: c1
-discovery: consul
-discovery-consul-vtgate-datacenter-tmpl: "dev-{{ .Cluster }}"
-discovery-zk-whatever: 5
-`,
-			config: Config{
-				ID:            "c1",
-				Name:          "cluster1",
-				DiscoveryImpl: "consul",
-				DiscoveryFlagsByImpl: map[string]map[string]string{
-					"consul": {
-						"vtgate-datacenter-tmpl": "dev-{{ .Cluster }}",
-					},
-					"zk": {
-						"whatever": "5",
-					},
-				},
-			},
-		},
-		{
-			name:   "errors",
-			yaml:   `name: "cluster1`,
-			config: Config{},
-			err:    assert.AnError,
-		},
-	}
-
-	for _, tt := range tests {
-		tt := tt
-
-		t.Run(tt.name, func(t *testing.T) {
-			t.Parallel()
-
-			cfg := Config{
-				DiscoveryFlagsByImpl: map[string]map[string]string{},
-			}
-
-			err := yaml.Unmarshal([]byte(tt.yaml), &cfg)
-			if tt.err != nil {
-				assert.Error(t, err)
-				return
-			}
-
-			require.NoError(t, err)
-			assert.Equal(t, tt.config, cfg)
-		})
-	}
-}

go/vt/vtadmin/cluster/file_config.go

@@ -17,17 +17,17 @@ limitations under the License.
 package cluster
 
 import (
-	"os"
 	"strings"
 
-	"gopkg.in/yaml.v2"
+	"github.com/spf13/viper"
 )
 
 // FileConfig represents the structure of a set of cluster configs on disk. It
-// contains both a default config, and cluster-specific overrides. Currently
-// only YAML config files are supported.
+// contains both a default config, and cluster-specific overrides. Viper is used
+// internally to load the config file, so any file format supported by viper is
+// permitted.
 //
-// A valid config looks like:
+// A valid YAML config looks like:
 //		defaults:
 //			discovery: k8s
 //		clusters:
@@ -42,30 +42,6 @@ type FileConfig struct {
 	Clusters map[string]Config
 }
 
-// UnmarshalYAML is part of the yaml.Unmarshaler interface.
-func (fc *FileConfig) UnmarshalYAML(unmarshal func(interface{}) error) error {
-	tmp := struct {
-		Defaults Config
-		Clusters map[string]Config
-	}{
-		Defaults: fc.Defaults,
-		Clusters: fc.Clusters,
-	}
-
-	if err := unmarshal(&tmp); err != nil {
-		return err
-	}
-
-	fc.Defaults = tmp.Defaults
-	fc.Clusters = make(map[string]Config, len(tmp.Clusters))
-
-	for id, cfg := range tmp.Clusters {
-		fc.Clusters[id] = cfg.Merge(Config{ID: id})
-	}
-
-	return nil
-}
-
 // String is part of the flag.Value interface.
 func (fc *FileConfig) String() string {
 	buf := strings.Builder{}
@@ -88,14 +64,57 @@ func (fc *FileConfig) Type() string {
 }
 
 // Set is part of the flag.Value interface. It loads the file configuration
-// found at the path passed to the flag.
+// found at the path passed to the flag. Any config file format supported by
+// viper is supported by FileConfig.
 func (fc *FileConfig) Set(value string) error {
-	data, err := os.ReadFile(value)
-	if err != nil {
+	v := viper.New()
+	v.SetConfigFile(value)
+	if err := v.ReadInConfig(); err != nil {
+		return err
+	}
+
+	return fc.unmarshalViper(v)
+}
+
+func (fc *FileConfig) unmarshalViper(v *viper.Viper) error {
+	tmp := struct { // work around mapstructure's interface; see https://github.com/spf13/viper/issues/338#issuecomment-382376136
+		Defaults map[string]string
+		Clusters map[string]map[string]string
+	}{
+		Defaults: map[string]string{},
+		Clusters: map[string]map[string]string{},
+	}
+
+	if err := v.Unmarshal(&tmp); err != nil {
 		return err
 	}
 
-	return yaml.Unmarshal(data, fc)
+	if err := fc.Defaults.unmarshalMap(tmp.Defaults); err != nil {
+		return err
+	}
+
+	if fc.Clusters == nil {
+		fc.Clusters = map[string]Config{}
+	}
+	for id, clusterMap := range tmp.Clusters {
+		c, ok := fc.Clusters[id]
+		if !ok {
+			c = Config{
+				ID:                   id,
+				DiscoveryFlagsByImpl: map[string]map[string]string{},
+				VtSQLFlags:           map[string]string{},
+				VtctldFlags:          map[string]string{},
+			}
+		}
+
+		if err := c.unmarshalMap(clusterMap); err != nil {
+			return err
+		}
+
+		fc.Clusters[id] = c
+	}
+
+	return nil
 }
 
 // Combine combines a FileConfig with a default Config and a ClustersFlag (each