github.com/yankunsam/loki/v2@v2.6.3-0.20220817130409-389df5235c27/pkg/util/cfg/files.go (about)

     1  package cfg
     2  
     3  import (
     4  	"encoding/json"
     5  	"flag"
     6  	"fmt"
     7  	"io/ioutil"
     8  	"os"
     9  	"strconv"
    10  	"strings"
    11  
    12  	"github.com/drone/envsubst"
    13  	"github.com/pkg/errors"
    14  	"gopkg.in/yaml.v2"
    15  )
    16  
    17  // JSON returns a Source that opens the supplied `.json` file and loads it.
    18  func JSON(f *string) Source {
    19  	return func(dst Cloneable) error {
    20  		if f == nil {
    21  			return nil
    22  		}
    23  
    24  		j, err := ioutil.ReadFile(*f)
    25  		if err != nil {
    26  			return err
    27  		}
    28  
    29  		err = dJSON(j)(dst)
    30  		return errors.Wrap(err, *f)
    31  	}
    32  }
    33  
    34  // dJSON returns a JSON source and allows dependency injection
    35  func dJSON(y []byte) Source {
    36  	return func(dst Cloneable) error {
    37  		return json.Unmarshal(y, dst)
    38  	}
    39  }
    40  
    41  // YAML returns a Source that opens the supplied `.yaml` file and loads it.
    42  // When expandEnvVars is true, variables in the supplied '.yaml\ file are expanded
    43  // using https://pkg.go.dev/github.com/drone/envsubst?tab=overview
    44  func YAML(f string, expandEnvVars bool) Source {
    45  	return func(dst Cloneable) error {
    46  		y, err := ioutil.ReadFile(f)
    47  		if err != nil {
    48  			return err
    49  		}
    50  		if expandEnvVars {
    51  			s, err := envsubst.EvalEnv(string(y))
    52  			if err != nil {
    53  				return err
    54  			}
    55  			y = []byte(s)
    56  		}
    57  		err = dYAML(y)(dst)
    58  		return errors.Wrap(err, f)
    59  	}
    60  }
    61  
    62  // dYAML returns a YAML source and allows dependency injection
    63  func dYAML(y []byte) Source {
    64  	return func(dst Cloneable) error {
    65  		return yaml.UnmarshalStrict(y, dst)
    66  	}
    67  }
    68  
    69  func ConfigFileLoader(args []string, name string) Source {
    70  
    71  	return func(dst Cloneable) error {
    72  		freshFlags := flag.NewFlagSet("config-file-loader", flag.ContinueOnError)
    73  
    74  		// Ensure we register flags on a copy of the config so as to not mutate it while
    75  		// parsing out the config file location.
    76  		dst.Clone().RegisterFlags(freshFlags)
    77  
    78  		usage := freshFlags.Usage
    79  		freshFlags.Usage = func() { /* don't do anything by default, we will print usage ourselves, but only when requested. */ }
    80  
    81  		err := freshFlags.Parse(args)
    82  		if err == flag.ErrHelp {
    83  			// print available parameters to stdout, so that users can grep/less it easily
    84  			freshFlags.SetOutput(os.Stdout)
    85  			usage()
    86  			os.Exit(2)
    87  		} else if err != nil {
    88  			fmt.Fprintln(freshFlags.Output(), "Run with -help to get list of available parameters")
    89  			os.Exit(2)
    90  		}
    91  
    92  		f := freshFlags.Lookup(name)
    93  		if f == nil || f.Value.String() == "" {
    94  			return nil
    95  		}
    96  
    97  		for _, val := range strings.Split(f.Value.String(), ",") {
    98  			val := strings.TrimSpace(val)
    99  			expandEnv := false
   100  			expandEnvFlag := freshFlags.Lookup("config.expand-env")
   101  			if expandEnvFlag != nil {
   102  				expandEnv, _ = strconv.ParseBool(expandEnvFlag.Value.String()) // Can ignore error as false returned
   103  			}
   104  			if _, err := os.Stat(val); err == nil {
   105  				return YAML(val, expandEnv)(dst)
   106  			}
   107  		}
   108  		return fmt.Errorf("%s does not exist, set %s for custom config path", f.Value.String(), name)
   109  	}
   110  }