cosolidate code

This commit is contained in:
davidnewhall2 2019-11-28 04:14:05 -08:00
parent 632936d772
commit 0a15850dc8
4 changed files with 79 additions and 82 deletions

View File

@ -1,13 +1,29 @@
package poller
import (
"crypto/tls"
"fmt"
"github.com/davidnewhall/unifi-poller/influxunifi"
"github.com/davidnewhall/unifi-poller/metrics"
client "github.com/influxdata/influxdb1-client/v2"
influx "github.com/influxdata/influxdb1-client/v2"
)
// GetInfluxDB returns an InfluxDB interface.
func (u *UnifiPoller) GetInfluxDB() (err error) {
u.Influx, err = influx.NewHTTPClient(influx.HTTPConfig{
Addr: u.Config.InfluxURL,
Username: u.Config.InfluxUser,
Password: u.Config.InfluxPass,
TLSConfig: &tls.Config{InsecureSkipVerify: u.Config.InfxBadSSL},
})
if err != nil {
return fmt.Errorf("influxdb: %v", err)
}
u.Logf("Logging Measurements to InfluxDB at %s as user %s", u.Config.InfluxURL, u.Config.InfluxUser)
return nil
}
// CollectAndProcess collects measurements and then reports them to InfluxDB
// Can be called once or in a ticker loop. This function and all the ones below
// handle their own logging. An error is returned so the calling function may
@ -31,7 +47,7 @@ func (u *UnifiPoller) ReportMetrics(metrics *metrics.Metrics) error {
m := &influxunifi.Metrics{Metrics: metrics}
// Make a new Influx Points Batcher.
var err error
m.BatchPoints, err = client.NewBatchPoints(client.BatchPointsConfig{Database: u.Config.InfluxDB})
m.BatchPoints, err = influx.NewBatchPoints(influx.BatchPointsConfig{Database: u.Config.InfluxDB})
if err != nil {
return fmt.Errorf("influx.NewBatchPoints: %v", err)
}

View File

@ -2,12 +2,29 @@ package poller
import (
"fmt"
"net/http"
"strings"
"time"
"github.com/davidnewhall/unifi-poller/metrics"
"github.com/davidnewhall/unifi-poller/promunifi"
"github.com/prometheus/client_golang/prometheus"
"github.com/prometheus/client_golang/prometheus/promhttp"
)
// RunPrometheus starts the web server and registers the collector.
func (u *UnifiPoller) RunPrometheus() error {
u.Logf("Exporting Measurements at https://%s/metrics for Prometheus", u.Config.HTTPListen)
http.Handle("/metrics", promhttp.Handler())
prometheus.MustRegister(promunifi.NewUnifiCollector(promunifi.UnifiCollectorCnfg{
Namespace: strings.Replace(u.Config.Namespace, "-", "", -1),
CollectFn: u.ExportMetrics,
LoggingFn: u.LogExportReport,
ReportErrors: true, // XXX: Does this need to be configurable?
}))
return http.ListenAndServe(u.Config.HTTPListen, nil)
}
// ExportMetrics updates the internal metrics provided via
// HTTP at /metrics for prometheus collection.
// This is run by Prometheus as CollectFn.

View File

@ -2,20 +2,13 @@
package poller
import (
"crypto/tls"
"fmt"
"log"
"net/http"
"os"
"strings"
"time"
"github.com/davidnewhall/unifi-poller/promunifi"
influx "github.com/influxdata/influxdb1-client/v2"
"github.com/prometheus/client_golang/prometheus"
"github.com/prometheus/client_golang/prometheus/promhttp"
"github.com/spf13/pflag"
"golift.io/unifi"
)
// New returns a new poller struct preloaded with default values.
@ -40,7 +33,7 @@ func New() *UnifiPoller {
// Start begins the application from a CLI.
// Parses cli flags, parses config file, parses env vars, sets up logging, then:
// - dumps a json payload OR - authenticates unifi controller and executes Run().
// - dumps a json payload OR - executes Run().
func (u *UnifiPoller) Start() error {
log.SetFlags(log.LstdFlags)
u.Flag.Parse(os.Args[1:])
@ -75,10 +68,6 @@ func (u *UnifiPoller) Start() error {
}
log.Printf("[INFO] UniFi Poller v%v Starting Up! PID: %d", Version, os.Getpid())
if err := u.GetUnifi(); err != nil {
return err
}
return u.Run()
}
@ -102,6 +91,9 @@ func (f *Flag) Parse(args []string) {
// 2. Run the collector one time and report the metrics to influxdb. (lambda)
// 3. Start a web server and wait for Prometheus to poll the application for metrics.
func (u *UnifiPoller) Run() error {
if err := u.GetUnifi(); err != nil {
return err
}
u.Logf("Polling UniFi Controller at %s v%s as user %s. Sites: %v",
u.Config.UnifiBase, u.Unifi.ServerVersion, u.Config.UnifiUser, u.Config.Sites)
@ -110,62 +102,43 @@ func (u *UnifiPoller) Run() error {
if err := u.GetInfluxDB(); err != nil {
return err
}
u.Logf("Logging Measurements to InfluxDB at %s as user %s", u.Config.InfluxURL, u.Config.InfluxUser)
u.Config.Mode = "influx poller"
return u.PollController()
case "influxlambda", "lambdainflux", "lambda_influx", "influx_lambda":
if err := u.GetInfluxDB(); err != nil {
return err
}
u.Logf("Logging Measurements to InfluxDB at %s as user %s one time (lambda mode)",
u.Config.InfluxURL, u.Config.InfluxUser)
u.LastCheck = time.Now()
return u.CollectAndProcess()
case "prometheus", "exporter":
u.Logf("Exporting Measurements at https://%s/metrics for Prometheus", u.Config.HTTPListen)
http.Handle("/metrics", promhttp.Handler())
prometheus.MustRegister(promunifi.NewUnifiCollector(promunifi.UnifiCollectorCnfg{
Namespace: strings.Replace(u.Config.Namespace, "-", "", -1),
CollectFn: u.ExportMetrics,
LoggingFn: u.LogExportReport,
ReportErrors: true, // XXX: Does this need to be configurable?
}))
return http.ListenAndServe(u.Config.HTTPListen, nil)
return u.RunPrometheus()
}
}
// GetInfluxDB returns an InfluxDB interface.
func (u *UnifiPoller) GetInfluxDB() (err error) {
u.Influx, err = influx.NewHTTPClient(influx.HTTPConfig{
Addr: u.Config.InfluxURL,
Username: u.Config.InfluxUser,
Password: u.Config.InfluxPass,
TLSConfig: &tls.Config{InsecureSkipVerify: u.Config.InfxBadSSL},
})
if err != nil {
return fmt.Errorf("influxdb: %v", err)
// PollController runs forever, polling UniFi and pushing to InfluxDB
// This is started by Run() after everything checks out.
func (u *UnifiPoller) PollController() error {
interval := u.Config.Interval.Round(time.Second)
log.Printf("[INFO] Everything checks out! Poller started, interval: %v", interval)
ticker := time.NewTicker(interval)
defer ticker.Stop()
for u.LastCheck = range ticker.C {
var err error
if u.Config.ReAuth {
u.LogDebugf("Re-authenticating to UniFi Controller")
// Some users need to re-auth every interval because the cookie times out.
if err = u.Unifi.Login(); err != nil {
u.LogError(err, "re-authenticating")
}
}
if err == nil {
// Only run this if the authentication procedure didn't return error.
_ = u.CollectAndProcess()
}
if u.errorCount > 0 {
return fmt.Errorf("too many errors, stopping poller")
}
}
return nil
}
// GetUnifi returns a UniFi controller interface.
func (u *UnifiPoller) GetUnifi() (err error) {
// Create an authenticated session to the Unifi Controller.
u.Unifi, err = unifi.NewUnifi(&unifi.Config{
User: u.Config.UnifiUser,
Pass: u.Config.UnifiPass,
URL: u.Config.UnifiBase,
VerifySSL: u.Config.VerifySSL,
ErrorLog: u.LogErrorf, // Log all errors.
DebugLog: u.LogDebugf, // Log debug messages.
})
if err != nil {
return fmt.Errorf("unifi controller: %v", err)
}
u.LogDebugf("Authenticated with controller successfully")
return u.CheckSites()
}

View File

@ -2,7 +2,6 @@ package poller
import (
"fmt"
"log"
"strings"
"time"
@ -10,31 +9,23 @@ import (
"golift.io/unifi"
)
// PollController runs forever, polling UniFi
// and pushing to influx OR exporting for prometheus.
// This is started by Run() after everything checks out.
func (u *UnifiPoller) PollController() error {
interval := u.Config.Interval.Round(time.Second)
log.Printf("[INFO] Everything checks out! Poller started in %v mode, interval: %v", u.Config.Mode, interval)
ticker := time.NewTicker(interval)
for u.LastCheck = range ticker.C {
var err error
if u.Config.ReAuth {
u.LogDebugf("Re-authenticating to UniFi Controller")
// Some users need to re-auth every interval because the cookie times out.
if err = u.Unifi.Login(); err != nil {
u.LogError(err, "re-authenticating")
}
}
if err == nil {
// Only run this if the authentication procedure didn't return error.
_ = u.CollectAndProcess()
}
if u.errorCount > 0 {
return fmt.Errorf("too many errors, stopping poller")
}
// GetUnifi returns a UniFi controller interface.
func (u *UnifiPoller) GetUnifi() (err error) {
// Create an authenticated session to the Unifi Controller.
u.Unifi, err = unifi.NewUnifi(&unifi.Config{
User: u.Config.UnifiUser,
Pass: u.Config.UnifiPass,
URL: u.Config.UnifiBase,
VerifySSL: u.Config.VerifySSL,
ErrorLog: u.LogErrorf, // Log all errors.
DebugLog: u.LogDebugf, // Log debug messages.
})
if err != nil {
return fmt.Errorf("unifi controller: %v", err)
}
return nil
u.LogDebugf("Authenticated with controller successfully")
return u.CheckSites()
}
// CheckSites makes sure the list of provided sites exists on the controller.