storj/pkg/process/exec_conf.go

// Copyright (C) 2018 Storj Labs, Inc.
// See LICENSE for copying information.

package process

import (
	"context"
	"flag"
	"log"
	"os"
	"strings"
	"sync"

	"github.com/spf13/cobra"
	"github.com/spf13/pflag"
	"github.com/spf13/viper"
	"go.uber.org/zap"
	monkit "gopkg.in/spacemonkeygo/monkit.v2"

	"storj.io/storj/pkg/telemetry"
	"storj.io/storj/pkg/utils"
)

// ExecuteWithConfig runs a Cobra command with the provided default config
func ExecuteWithConfig(cmd *cobra.Command, defaultConfig string) {
	cfgFile := flag.String("config", os.ExpandEnv(defaultConfig), "config file")
	pflag.CommandLine.AddGoFlagSet(flag.CommandLine)
	cleanup(cmd, cfgFile)
	_ = cmd.Execute()
}

var (
	mon = monkit.Package()

	contextMtx sync.Mutex
	contexts   = map[*cobra.Command]context.Context{}
)

// Ctx returns the appropriate context.Context for ExecuteWithConfig commands
func Ctx(cmd *cobra.Command) context.Context {
	contextMtx.Lock()
	defer contextMtx.Unlock()
	ctx := contexts[cmd]
	if ctx == nil {
		return context.Background()
	}
	return ctx
}

type ctxKey int

const (
	ctxKeyVip ctxKey = iota
	ctxKeyCfg
)

// SaveConfig outputs the configuration to the configured (or default) config
// file given to ExecuteWithConfig
func SaveConfig(cmd *cobra.Command) error {
	ctx := Ctx(cmd)
	return getViper(ctx).WriteConfigAs(CfgPath(ctx))
}

// SaveConfigAs outputs the configuration to the provided path assuming the
// command was executed with ExecuteWithConfig
func SaveConfigAs(cmd *cobra.Command, path string) error {
	return getViper(Ctx(cmd)).WriteConfigAs(path)
}

func getViper(ctx context.Context) *viper.Viper {
	if v, ok := ctx.Value(ctxKeyVip).(*viper.Viper); ok {
		return v
	}
	return nil
}

// CfgPath returns the configuration path used with ExecuteWithConfig
func CfgPath(ctx context.Context) string {
	if v, ok := ctx.Value(ctxKeyCfg).(string); ok {
		return v
	}
	return ""
}

func cleanup(cmd *cobra.Command, cfgFile *string) {
	for _, ccmd := range cmd.Commands() {
		cleanup(ccmd, cfgFile)
	}
	if cmd.Run != nil {
		panic("Please use cobra's RunE instead of Run")
	}
	internalRun := cmd.RunE
	if internalRun == nil {
		return
	}
	cmd.RunE = func(cmd *cobra.Command, args []string) (err error) {
		ctx := context.Background()
		defer mon.TaskNamed("root")(&ctx)(&err)

		vip := viper.New()
		err = vip.BindPFlags(cmd.Flags())
		if err != nil {
			return err
		}
		vip.SetEnvPrefix("storj")
		vip.AutomaticEnv()
		if *cfgFile != "" && fileExists(*cfgFile) {
			vip.SetConfigFile(*cfgFile)
			err = vip.ReadInConfig()
			if err != nil {
				return err
			}
		}

		// go back and propagate changed config values to appropriate flags
		var brokenKeys []string
		for _, key := range vip.AllKeys() {
			if cmd.Flags().Lookup(key) == nil {
				// flag couldn't be found
				brokenKeys = append(brokenKeys, key)
			} else {
				err := cmd.Flags().Set(key, vip.GetString(key))
				if err != nil {
					// flag couldn't be set
					brokenKeys = append(brokenKeys, key)
				}
			}
		}

		ctx = context.WithValue(ctx, ctxKeyVip, vip)
		ctx = context.WithValue(ctx, ctxKeyCfg, *cfgFile)

		logger, err := utils.NewLogger(*logDisposition)
		if err != nil {
			return err
		}
		defer func() { _ = logger.Sync() }()
		defer zap.ReplaceGlobals(logger)()
		defer zap.RedirectStdLog(logger)()

		// okay now that logging is working, inform about the broken keys
		// these keys are almost certainly broken because they have capital
		// letters
		if len(brokenKeys) > 0 {
			logger.Sugar().Infof("TODO: these flags are not configurable via "+
				"config file, probably due to having uppercase letters: %s",
				strings.Join(brokenKeys, ", "))
		}

		err = initMetrics(ctx, monkit.Default,
			telemetry.DefaultInstanceID())
		if err != nil {
			logger.Error("failed to configure telemetry", zap.Error(err))
		}

		err = initDebug(logger, monkit.Default)
		if err != nil {
			logger.Error("failed to start debug endpoints", zap.Error(err))
		}

		contextMtx.Lock()
		contexts[cmd] = ctx
		contextMtx.Unlock()
		defer func() {
			contextMtx.Lock()
			delete(contexts, cmd)
			contextMtx.Unlock()
		}()

		err = internalRun(cmd, args)
		if err != nil {
			log.Fatalf("%+v", err)
		}
		return err
	}
}