viperconfig.go

/*
Copyright 2018 The Kubernetes Authors.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/

package viperconfig

import (
	"flag"
	"fmt"
	"github.com/pkg/errors"
	"path/filepath"

	"github.com/spf13/viper"
)

const (
	viperFileNotFound = "Unsupported Config Type \"\""
)

// ViperizeFlags checks whether a configuration file was specified, reads it, and updates
// the configuration variables accordingly. Must be called after framework.HandleFlags()
// and before framework.AfterReadingAllFlags().
//
// The logic is so that a required configuration file must be present. If empty,
// the optional configuration file is used instead, unless also empty.
//
// Files can be specified with just a base name ("e2e", matches "e2e.json/yaml/..." in
// the current directory) or with path and suffix.
func ViperizeFlags(requiredConfig, optionalConfig string) error {
	viperConfig := optionalConfig
	required := false
	if requiredConfig != "" {
		viperConfig = requiredConfig
		required = true
	}
	if viperConfig == "" {
		return nil
	}
	viper.SetConfigName(filepath.Base(viperConfig))
	viper.AddConfigPath(filepath.Dir(viperConfig))
	wrapError := func(err error) error {
		if err == nil {
			return nil
		}
		errorPrefix := fmt.Sprintf("viper config %q", viperConfig)
		actualFile := viper.ConfigFileUsed()
		if actualFile != "" && actualFile != viperConfig {
			errorPrefix = fmt.Sprintf("%s = %q", errorPrefix, actualFile)
		}
		return errors.Wrap(err, errorPrefix)
	}

	if err := viper.ReadInConfig(); err != nil {
		// If the user specified a file suffix, the Viper won't
		// find the file because it always appends its known set
		// of file suffices. Therefore try once more without
		// suffix.
		ext := filepath.Ext(viperConfig)
		if ext != "" && err.Error() == viperFileNotFound {
			viper.SetConfigName(filepath.Base(viperConfig[0 : len(viperConfig)-len(ext)]))
			err = viper.ReadInConfig()
		}
		if err != nil {
			// If a config was required, then parsing must
			// succeed. This catches syntax errors and
			// "file not found". Unfortunately error
			// messages are sometimes hard to understand,
			// so try to help the user a bit.
			switch err.Error() {
			case viperFileNotFound:
				if required {
					return wrapError(errors.New("not found or not using a supported file format"))
				}
				// Proceed without config.
				return nil
			default:
				// Something isn't right in the file.
				return wrapError(err)
			}
		}
	}

	// Update all flag values not already set with values found
	// via Viper. We do this ourselves instead of calling
	// something like viper.Unmarshal(&TestContext) because we
	// want to support all values, regardless where they are
	// stored.
	return wrapError(viperUnmarshal())
}

// viperUnmarshall updates all command line flags with the corresponding values found
// via Viper, regardless whether the flag value is stored in TestContext, some other
// context or a local variable.
func viperUnmarshal() error {
	var result error
	set := make(map[string]bool)

	// Determine which values were already set explicitly via
	// flags. Those we don't overwrite because command line
	// flags have a higher priority.
	flag.Visit(func(f *flag.Flag) {
		set[f.Name] = true
	})

	flag.VisitAll(func(f *flag.Flag) {
		if result != nil ||
			set[f.Name] ||
			!viper.IsSet(f.Name) {
			return
		}

		// In contrast to viper.Unmarshal(), values
		// that have the wrong type (for example, a
		// list instead of a plain string) will not
		// trigger an error here. This could be fixed
		// by checking the type ourselves, but
		// probably isn't worth the effort.
		//
		// "%v" correctly turns bool, int, strings into
		// the representation expected by flag, so those
		// can be used in config files. Plain strings
		// always work there, just as on the command line.
		str := fmt.Sprintf("%v", viper.Get(f.Name))
		if err := f.Value.Set(str); err != nil {
			result = fmt.Errorf("setting option %q from config file value: %s", f.Name, err)
		}
	})

	return result
}