Use new libraries.

cmd/unifi-poller/main.go

@@ -8,11 +8,16 @@ import (
 	"time"
 
 	"github.com/golift/unifi"
-	influx "github.com/influxdata/influxdb/client/v2"
+	influx "github.com/influxdata/influxdb1-client/v2"
 	"github.com/naoina/toml"
 	flag "github.com/ogier/pflag"
 )
 
+// Asset is used to give all devices and clients a common interface.
+type Asset interface {
+	Points() ([]*influx.Point, error)
+}
+
 func main() {
 	configFile := parseFlags()
 	log.Println("Unifi-Poller Starting Up! PID:", os.Getpid())
@@ -24,14 +29,14 @@ func main() {
 		log.SetFlags(log.Lshortfile | log.Lmicroseconds | log.Ldate)
 		log.Println("Debug Logging Enabled")
 	}
-	log.Println("Loaded Configuration:", configFile)
 	// Create an authenticated session to the Unifi Controller.
-	device, err := unifi.AuthController(config.UnifiUser, config.UnifiPass, config.UnifiBase, config.VerifySSL)
+	device, err := unifi.GetController(config.UnifiUser, config.UnifiPass, config.UnifiBase, config.VerifySSL)
 	if err != nil {
 		log.Fatalln("Unifi Controller Error:", err)
 	} else if !config.Quiet {
 		log.Println("Authenticated to Unifi Controller @", config.UnifiBase, "as user", config.UnifiUser)
 	}
+	device.ErrorLog = log.Printf
 	infdb, err := influx.NewHTTPClient(influx.HTTPConfig{
 		Addr:     config.InfluxURL,
 		Username: config.InfluxUser,
@@ -41,13 +46,12 @@ func main() {
 		log.Fatalln("InfluxDB Error:", err)
 	} else if config.Quiet {
 		// Doing it this way allows debug error logs (line numbers, etc)
-		unifi.Debug = false
+		device.DebugLog = log.Printf
 	} else {
 		log.Println("Logging Unifi Metrics to InfluXDB @", config.InfluxURL, "as user", config.InfluxUser)
 		log.Println("Polling Unifi Controller, interval:", config.Interval.value)
 	}
-	log.Println("Everyting checks out! Beginning Poller Routine.")
-	config.PollUnifiController(infdb, device, config.Quiet)
+	config.PollUnifiController(infdb, device)
 }
 
 func parseFlags() string {
@@ -86,39 +90,65 @@ func GetConfig(configFile string) (Config, error) {
 	} else if err := toml.Unmarshal(buf, &config); err != nil {
 		return config, err
 	}
+	log.Println("Loaded Configuration:", configFile)
 	return config, nil
 }
 
 // PollUnifiController runs forever, polling and pushing.
-func (c *Config) PollUnifiController(infdb influx.Client, device *unifi.AuthedReq, quiet bool) {
+func (c *Config) PollUnifiController(infdb influx.Client, uni *unifi.Unifi) {
+	log.Println("Everyting checks out! Beginning Poller Routine.")
 	ticker := time.NewTicker(c.Interval.value)
 	for range ticker.C {
-		var clients, devices []unifi.Asset
-		var bp influx.BatchPoints
-		var err error
-		if clients, err = device.GetUnifiClientAssets(); err != nil {
-			log.Println("ERROR unifi.GetUnifiClientsAssets():", err)
-		} else if devices, err = device.GetUnifiDeviceAssets(); err != nil {
-			log.Println("ERROR unifi.GetUnifiDeviceAssets():", err)
-		} else if bp, err = influx.NewBatchPoints(influx.BatchPointsConfig{Database: c.InfluxDB}); err != nil {
-			log.Println("ERROR influx.NewBatchPoints:", err)
-		}
-		if err != nil {
-			continue
-		}
-		for _, asset := range append(clients, devices...) {
-			if pt, errr := asset.Points(); errr != nil {
-				log.Println("ERROR asset.Points():", errr)
-			} else {
-				bp.AddPoints(pt)
-			}
-		}
-		if err = infdb.Write(bp); err != nil {
-			log.Println("ERROR infdb.Write(bp):", err)
-			continue
-		}
-		if !quiet {
-			log.Println("Logged Unifi States. Clients:", len(clients), "- Devices:", len(devices))
+		if clients, err := uni.GetClients(); err != nil {
+			logErrors([]error{err}, "uni.GetClients()")
+		} else if devices, err := uni.GetDevices(); err != nil {
+			logErrors([]error{err}, "uni.GetDevices()")
+		} else if bp, err := influx.NewBatchPoints(influx.BatchPointsConfig{Database: c.InfluxDB}); err != nil {
+			logErrors([]error{err}, "influx.NewBatchPoints")
+		} else if errs := batchPoints(devices, clients, bp); errs != nil {
+			logErrors(errs, "asset.Points()")
+		} else if err := infdb.Write(bp); err != nil {
+			logErrors([]error{err}, "infdb.Write(bp)")
+		} else if !c.Quiet {
+			log.Println("Logged Unifi States. Clients:", len(clients.UCLs), "- Wireless APs:",
+				len(devices.UAPs), "Gateways:", len(devices.USGs), "Switches:", len(devices.USWs))
+		}
+	}
+}
+
+// batchPoints combines all device and client data into influxdb data points.
+func batchPoints(devices *unifi.Devices, clients *unifi.Clients, batchPoints influx.BatchPoints) (errs []error) {
+	process := func(asset Asset) error {
+		influxPoints, err := asset.Points()
+		if err != nil {
+			return err
+		}
+		batchPoints.AddPoints(influxPoints)
+		return nil
+	}
+	for _, asset := range devices.UAPs {
+		errs = append(errs, process(asset))
+	}
+	for _, asset := range devices.USGs {
+		errs = append(errs, process(asset))
+	}
+	for _, asset := range devices.USWs {
+		errs = append(errs, process(asset))
+	}
+	for _, asset := range clients.UCLs {
+		errs = append(errs, process(asset))
+	}
+	return
+}
+
+// logErrors writes a slice of errors, with a prefix, to log-out.
+func logErrors(errs []error, prefix string) {
+	if errs == nil {
+		return
+	}
+	for _, err := range errs {
+		if err != nil {
+			log.Println("[ERROR]", prefix+":", err.Error())
 		}
 	}
 }