Use reflection and tags for env variables.

This commit is contained in:
David Newhall II 2019-08-25 23:19:29 -07:00
parent 09e2886554
commit 96db7429d8
3 changed files with 45 additions and 60 deletions

View File

@ -22,27 +22,6 @@ const (
defaultUnifURL = "https://127.0.0.1:8443" defaultUnifURL = "https://127.0.0.1:8443"
) )
// These are environment variables that can be used to override configuration.
// Useful for Docker users.
const (
ENVConfigMode = "UP_POLLING_MODE"
ENVConfigInfluxDB = "UP_INFLUX_DB"
ENVConfigInfluxUser = "UP_INFLUX_USER"
ENVConfigInfluxPass = "UP_INFLUX_PASS"
ENVConfigInfluxURL = "UP_INFLUX_URL"
ENVConfigUnifiUser = "UP_UNIFI_USER"
ENVConfigUnifiPass = "UP_UNIFI_PASS"
ENVConfigUnifiBase = "UP_UNIFI_URL"
ENVConfigReAuth = "UP_REAUTHENTICATE"
ENVConfigVerifySSL = "UP_VERIFY_SSL"
ENVConfigCollectIDS = "UP_COLLECT_IDS"
ENVConfigQuiet = "UP_QUIET_MODE"
ENVConfigDebug = "UP_DEBUG_MODE"
ENVConfigInterval = "UP_POLLING_INTERVAL"
ENVConfigMaxErrors = "UP_MAX_ERRORS"
ENVConfigSites = "UP_POLL_SITES"
)
// UnifiPoller contains the application startup data, and auth info for UniFi & Influx. // UnifiPoller contains the application startup data, and auth info for UniFi & Influx.
type UnifiPoller struct { type UnifiPoller struct {
Influx influx.Client Influx influx.Client
@ -74,22 +53,22 @@ type Metrics struct {
// Config represents the data needed to poll a controller and report to influxdb. // Config represents the data needed to poll a controller and report to influxdb.
// This is all of the data stored in the config file. // This is all of the data stored in the config file.
type Config struct { type Config struct {
MaxErrors int `json:"max_errors,_omitempty" toml:"max_errors,_omitempty" xml:"max_errors" yaml:"max_errors"` MaxErrors int `json:"max_errors,_omitempty" toml:"max_errors,_omitempty" xml:"max_errors" yaml:"max_errors" env:"UP_MAX_ERRORS"`
Interval Duration `json:"interval,_omitempty" toml:"interval,_omitempty" xml:"interval" yaml:"interval"` Interval Duration `json:"interval,_omitempty" toml:"interval,_omitempty" xml:"interval" yaml:"interval" env:"UP_POLLING_INTERVAL"`
Debug bool `json:"debug" toml:"debug" xml:"debug" yaml:"debug"` Debug bool `json:"debug" toml:"debug" xml:"debug" yaml:"debug" env:"UP_DEBUG_MODE"`
Quiet bool `json:"quiet,_omitempty" toml:"quiet,_omitempty" xml:"quiet" yaml:"quiet"` Quiet bool `json:"quiet,_omitempty" toml:"quiet,_omitempty" xml:"quiet" yaml:"quiet" env:"UP_QUIET_MODE"`
VerifySSL bool `json:"verify_ssl" toml:"verify_ssl" xml:"verify_ssl" yaml:"verify_ssl"` VerifySSL bool `json:"verify_ssl" toml:"verify_ssl" xml:"verify_ssl" yaml:"verify_ssl" env:"UP_VERIFY_SSL"`
CollectIDS bool `json:"collect_ids" toml:"collect_ids" xml:"collect_ids" yaml:"collect_ids"` CollectIDS bool `json:"collect_ids" toml:"collect_ids" xml:"collect_ids" yaml:"collect_ids" env:"UP_COLLECT_IDS"`
ReAuth bool `json:"reauthenticate" toml:"reauthenticate" xml:"reauthenticate" yaml:"reauthenticate"` ReAuth bool `json:"reauthenticate" toml:"reauthenticate" xml:"reauthenticate" yaml:"reauthenticate" env:"UP_REAUTHENTICATE"`
Mode string `json:"mode" toml:"mode" xml:"mode" yaml:"mode"` Mode string `json:"mode" toml:"mode" xml:"mode" yaml:"mode" env:"UP_POLLING_MODE"`
InfluxURL string `json:"influx_url,_omitempty" toml:"influx_url,_omitempty" xml:"influx_url" yaml:"influx_url"` InfluxURL string `json:"influx_url,_omitempty" toml:"influx_url,_omitempty" xml:"influx_url" yaml:"influx_url" env:"UP_INFLUX_URL"`
InfluxUser string `json:"influx_user,_omitempty" toml:"influx_user,_omitempty" xml:"influx_user" yaml:"influx_user"` InfluxUser string `json:"influx_user,_omitempty" toml:"influx_user,_omitempty" xml:"influx_user" yaml:"influx_user" env:"UP_INFLUX_USER"`
InfluxPass string `json:"influx_pass,_omitempty" toml:"influx_pass,_omitempty" xml:"influx_pass" yaml:"influx_pass"` InfluxPass string `json:"influx_pass,_omitempty" toml:"influx_pass,_omitempty" xml:"influx_pass" yaml:"influx_pass" env:"UP_INFLUX_PASS"`
InfluxDB string `json:"influx_db,_omitempty" toml:"influx_db,_omitempty" xml:"influx_db" yaml:"influx_db"` InfluxDB string `json:"influx_db,_omitempty" toml:"influx_db,_omitempty" xml:"influx_db" yaml:"influx_db" env:"UP_INFLUX_DB"`
UnifiUser string `json:"unifi_user,_omitempty" toml:"unifi_user,_omitempty" xml:"unifi_user" yaml:"unifi_user"` UnifiUser string `json:"unifi_user,_omitempty" toml:"unifi_user,_omitempty" xml:"unifi_user" yaml:"unifi_user" env:"UP_UNIFI_USER"`
UnifiPass string `json:"unifi_pass,_omitempty" toml:"unifi_pass,_omitempty" xml:"unifi_pass" yaml:"unifi_pass"` UnifiPass string `json:"unifi_pass,_omitempty" toml:"unifi_pass,_omitempty" xml:"unifi_pass" yaml:"unifi_pass" env:"UP_UNIFI_PASS"`
UnifiBase string `json:"unifi_url,_omitempty" toml:"unifi_url,_omitempty" xml:"unifi_url" yaml:"unifi_url"` UnifiBase string `json:"unifi_url,_omitempty" toml:"unifi_url,_omitempty" xml:"unifi_url" yaml:"unifi_url" env:"UP_UNIFI_URL"`
Sites []string `json:"sites,_omitempty" toml:"sites,_omitempty" xml:"sites" yaml:"sites"` Sites []string `json:"sites,_omitempty" toml:"sites,_omitempty" xml:"sites" yaml:"sites" env:"UP_POLL_SITES"`
} }
// Duration is used to UnmarshalTOML into a time.Duration value. // Duration is used to UnmarshalTOML into a time.Duration value.

View File

@ -7,6 +7,7 @@ import (
"io/ioutil" "io/ioutil"
"log" "log"
"os" "os"
"reflect"
"strconv" "strconv"
"strings" "strings"
"time" "time"
@ -51,29 +52,33 @@ func (f *Flag) Parse(args []string) {
// setEnvVarOptions copies environment variables into configuration values. // setEnvVarOptions copies environment variables into configuration values.
// This is useful for Docker users that find it easier to pass ENV variables // This is useful for Docker users that find it easier to pass ENV variables
// that a specific configuration file. // than a specific configuration file. Uses reflection to find struct tags.
func (u *UnifiPoller) setEnvVarOptions() { func (u *UnifiPoller) setEnvVarOptions() {
u.Config.Mode = pick(os.Getenv(ENVConfigMode), u.Config.Mode) t := reflect.TypeOf(u.Config) // whole struct
u.Config.InfluxDB = pick(os.Getenv(ENVConfigInfluxDB), u.Config.InfluxDB) // Loop each Config struct member and check for a reflect tag / env var setting.
u.Config.InfluxUser = pick(os.Getenv(ENVConfigInfluxUser), u.Config.InfluxUser) for i := 0; i < t.NumField(); i++ {
u.Config.InfluxPass = pick(os.Getenv(ENVConfigInfluxPass), u.Config.InfluxPass) tag := t.Field(i).Tag.Get("env") // struct member tag
u.Config.InfluxURL = pick(os.Getenv(ENVConfigInfluxURL), u.Config.InfluxURL) env := os.Getenv(tag) // value of "tag" env variable
u.Config.UnifiUser = pick(os.Getenv(ENVConfigUnifiUser), u.Config.UnifiUser) if tag == "" || env == "" {
u.Config.UnifiPass = pick(os.Getenv(ENVConfigUnifiPass), u.Config.UnifiPass) continue
u.Config.UnifiBase = pick(os.Getenv(ENVConfigUnifiBase), u.Config.UnifiBase)
u.Config.ReAuth = parseBool(os.Getenv(ENVConfigReAuth), u.Config.ReAuth)
u.Config.VerifySSL = parseBool(os.Getenv(ENVConfigVerifySSL), u.Config.VerifySSL)
u.Config.CollectIDS = parseBool(os.Getenv(ENVConfigCollectIDS), u.Config.CollectIDS)
u.Config.Quiet = parseBool(os.Getenv(ENVConfigQuiet), u.Config.Quiet)
u.Config.Debug = parseBool(os.Getenv(ENVConfigDebug), u.Config.Debug)
if e := os.Getenv(ENVConfigInterval); e != "" {
_ = u.Config.Interval.UnmarshalText([]byte(e))
} }
if e := os.Getenv(ENVConfigMaxErrors); e != "" { // Reflect and update the u.Config struct member at position i.
u.Config.MaxErrors, _ = strconv.Atoi(e) switch c := reflect.ValueOf(u.Config).Elem().Field(i); c.Type().String() {
// Handle each member type appropriately (differently).
case "string":
c.SetString(env)
case "int":
val, _ := strconv.Atoi(env)
c.Set(reflect.ValueOf(val))
case "[]string":
c.Set(reflect.ValueOf(strings.Split(env, ",")))
case "Duration":
val, _ := time.ParseDuration(env)
c.Set(reflect.ValueOf(Duration{val}))
case "bool":
val, _ := strconv.ParseBool(env)
c.SetBool(val)
} }
if e := os.Getenv(ENVConfigSites); e != "" {
u.Config.Sites = strings.Split(e, ",")
} }
} }

View File

@ -62,7 +62,8 @@ func (u *UnifiPoller) PollController() error {
_ = u.CollectAndReport() _ = u.CollectAndReport()
} }
if u.Config.MaxErrors >= 0 && u.errorCount > u.Config.MaxErrors { if u.Config.MaxErrors >= 0 && u.errorCount > u.Config.MaxErrors {
return fmt.Errorf("reached maximum error count, stopping poller (%d > %d)", u.errorCount, u.Config.MaxErrors) return fmt.Errorf("reached maximum error count, stopping poller (%d > %d)",
u.errorCount, u.Config.MaxErrors)
} }
} }
return nil return nil