github.com/grafana/pyroscope@v1.18.0/pkg/cfg/files.go (about)

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