Add collector interface

Uses node_exporter style collector registration Signed-off-by: Joe Adams <github@joeadams.io>

Author: sysadmind

cmd/postgres_exporter/main.go

@@ -19,6 +19,7 @@ import (
 
 "github.com/go-kit/log"
 "github.com/go-kit/log/level"
+"github.com/prometheus-community/postgres_exporter/collector"
 "github.com/prometheus/client_golang/prometheus"
 "github.com/prometheus/client_golang/prometheus/promhttp"
 "github.com/prometheus/common/promlog"
@@ -114,6 +115,13 @@ func main() {
 
 prometheus.MustRegister(exporter)
 
+pe, err := collector.NewPostgresCollector(logger, dsn)
+if err != nil {
+level.Error(logger).Log("msg", "Failed to create PostgresCollector", "err", err.Error())
+os.Exit(1)
+}
+prometheus.MustRegister(pe)
+
 http.Handle(*metricPath, promhttp.Handler())
 http.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) {
 w.Header().Set("Content-Type", "text/html; charset=UTF-8") // nolint: errcheck

cmd/postgres_exporter/server.go

@@ -14,16 +14,13 @@
 package main
 
 import (
-"context"
 "database/sql"
 "fmt"
-"log"
 "sync"
 "time"
 
 "github.com/blang/semver"
 "github.com/go-kit/log/level"
-"github.com/prometheus-community/postgres_exporter/collector"
 "github.com/prometheus/client_golang/prometheus"
 )
 
@@ -131,17 +128,6 @@ func (s *Server) Scrape(ch chan<- prometheus.Metric, disableSettingsMetrics bool
 err = fmt.Errorf("queryNamespaceMappings returned %d errors", len(errMap))
 }
 
-{
-pgdb := collector.NewPGDatabaseCollector()
-metrics, err := pgdb.Update(context.TODO(), s.db, s.String())
-if err != nil {
-log.Printf("Failed to scrape pg_database metrics: %s", err)
-}
-for _, m := range metrics {
-ch <- m
-}
-}
-
 return err
 }
 

collector/collector.go

@@ -0,0 +1,210 @@
+// Copyright 2022 The Prometheus Authors
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package collector
+
+import (
+"context"
+"errors"
+"fmt"
+"sync"
+"time"
+
+"github.com/go-kit/log"
+"github.com/go-kit/log/level"
+"github.com/prometheus/client_golang/prometheus"
+"gopkg.in/alecthomas/kingpin.v2"
+)
+
+var (
+factories = make(map[string]func(logger log.Logger) (Collector, error))
+initiatedCollectorsMtx = sync.Mutex{}
+initiatedCollectors = make(map[string]Collector)
+collectorState = make(map[string]*bool)
+forcedCollectors = map[string]bool{} // collectors which have been explicitly enabled or disabled
+)
+
+const (
+// Namespace for all metrics.
+namespace = "pg"
+
+defaultEnabled = true
+defaultDisabled = false
+)
+
+var (
+scrapeDurationDesc = prometheus.NewDesc(
+prometheus.BuildFQName(namespace, "scrape", "collector_duration_seconds"),
+"postgres_exporter: Duration of a collector scrape.",
+[]string{"collector"},
+nil,
+)
+scrapeSuccessDesc = prometheus.NewDesc(
+prometheus.BuildFQName(namespace, "scrape", "collector_success"),
+"postgres_exporter: Whether a collector succeeded.",
+[]string{"collector"},
+nil,
+)
+)
+
+type Collector interface {
+Update(ctx context.Context, server *server, ch chan<- prometheus.Metric) error
+}
+
+func registerCollector(name string, isDefaultEnabled bool, createFunc func(logger log.Logger) (Collector, error)) {
+var helpDefaultState string
+if isDefaultEnabled {
+helpDefaultState = "enabled"
+} else {
+helpDefaultState = "disabled"
+}
+
+// Create flag for this collector
+flagName := fmt.Sprintf("collector.%s", name)
+flagHelp := fmt.Sprintf("Enable the %s collector (default: %s).", name, helpDefaultState)
+defaultValue := fmt.Sprintf("%v", isDefaultEnabled)
+
+flag := kingpin.Flag(flagName, flagHelp).Default(defaultValue).Action(collectorFlagAction(name)).Bool()
+collectorState[name] = flag
+
+// Register the create function for this collector
+factories[name] = createFunc
+}
+
+// PostgresCollector implements the prometheus.Collector interface.
+type PostgresCollector struct {
+Collectors map[string]Collector
+logger log.Logger
+
+servers map[string]*server
+}
+
+// NewPostgresCollector creates a new PostgresCollector.
+func NewPostgresCollector(logger log.Logger, dsns []string, filters ...string) (*PostgresCollector, error) {
+f := make(map[string]bool)
+for _, filter := range filters {
+enabled, exist := collectorState[filter]
+if !exist {
+return nil, fmt.Errorf("missing collector: %s", filter)
+}
+if !*enabled {
+return nil, fmt.Errorf("disabled collector: %s", filter)
+}
+f[filter] = true
+}
+collectors := make(map[string]Collector)
+initiatedCollectorsMtx.Lock()
+defer initiatedCollectorsMtx.Unlock()
+for key, enabled := range collectorState {
+if !*enabled || (len(f) > 0 && !f[key]) {
+continue
+}
+if collector, ok := initiatedCollectors[key]; ok {
+collectors[key] = collector
+} else {
+collector, err := factories[key](log.With(logger, "collector", key))
+if err != nil {
+return nil, err
+}
+collectors[key] = collector
+initiatedCollectors[key] = collector
+}
+}
+
+servers := make(map[string]*server)
+for _, dsn := range dsns {
+s, err := makeServer(dsn)
+if err != nil {
+return nil, err
+}
+servers[dsn] = s
+}
+
+return &PostgresCollector{
+Collectors: collectors,
+logger: logger,
+servers: servers,
+}, nil
+}
+
+// Describe implements the prometheus.Collector interface.
+func (n PostgresCollector) Describe(ch chan<- *prometheus.Desc) {
+ch <- scrapeDurationDesc
+ch <- scrapeSuccessDesc
+}
+
+// Collect implements the prometheus.Collector interface.
+func (n PostgresCollector) Collect(ch chan<- prometheus.Metric) {
+ctx := context.TODO()
+wg := sync.WaitGroup{}
+wg.Add(len(n.servers))
+for _, s := range n.servers {
+go func(s *server) {
+n.subCollect(ctx, s, ch)
+wg.Done()
+}(s)
+}
+wg.Wait()
+}
+
+func (n PostgresCollector) subCollect(ctx context.Context, server *server, ch chan<- prometheus.Metric) {
+wg := sync.WaitGroup{}
+wg.Add(len(n.Collectors))
+for name, c := range n.Collectors {
+go func(name string, c Collector) {
+execute(ctx, name, c, server, ch, n.logger)
+wg.Done()
+}(name, c)
+}
+wg.Wait()
+}
+
+func execute(ctx context.Context, name string, c Collector, s *server, ch chan<- prometheus.Metric, logger log.Logger) {
+begin := time.Now()
+err := c.Update(ctx, s, ch)
+duration := time.Since(begin)
+var success float64
+
+if err != nil {
+if IsNoDataError(err) {
+level.Debug(logger).Log("msg", "collector returned no data", "name", name, "duration_seconds", duration.Seconds(), "err", err)
+} else {
+level.Error(logger).Log("msg", "collector failed", "name", name, "duration_seconds", duration.Seconds(), "err", err)
+}
+success = 0
+} else {
+level.Debug(logger).Log("msg", "collector succeeded", "name", name, "duration_seconds", duration.Seconds())
+success = 1
+}
+ch <- prometheus.MustNewConstMetric(scrapeDurationDesc, prometheus.GaugeValue, duration.Seconds(), name)
+ch <- prometheus.MustNewConstMetric(scrapeSuccessDesc, prometheus.GaugeValue, success, name)
+}
+
+// collectorFlagAction generates a new action function for the given collector
+// to track whether it has been explicitly enabled or disabled from the command line.
+// A new action function is needed for each collector flag because the ParseContext
+// does not contain information about which flag called the action.
+// See: https://github.com/alecthomas/kingpin/issues/294
+func collectorFlagAction(collector string) func(ctx *kingpin.ParseContext) error {
+return func(ctx *kingpin.ParseContext) error {
+forcedCollectors[collector] = true
+return nil
+}
+}
+
+// ErrNoData indicates the collector found no data to collect, but had no other error.
+var ErrNoData = errors.New("collector returned no data")
+
+func IsNoDataError(err error) bool {
+return err == ErrNoData
+}

collector/pg_database.go

@@ -1,4 +1,4 @@
-// Copyright 2021 The Prometheus Authors
+// Copyright 2022 The Prometheus Authors
 // Licensed under the Apache License, Version 2.0 (the "License");
 // you may not use this file except in compliance with the License.
 // You may obtain a copy of the License at
@@ -15,15 +15,21 @@ package collector
 
 import (
 "context"
-"database/sql"
 
+"github.com/go-kit/log"
 "github.com/prometheus/client_golang/prometheus"
 )
 
-type PGDatabaseCollector struct{}
+func init() {
+registerCollector("database", defaultEnabled, NewPGDatabaseCollector)
+}
+
+type PGDatabaseCollector struct {
+log log.Logger
+}
 
-func NewPGDatabaseCollector() *PGDatabaseCollector {
-return &PGDatabaseCollector{}
+func NewPGDatabaseCollector(logger log.Logger) (Collector, error) {
+return &PGDatabaseCollector{log: logger}, nil
 }
 
 var pgDatabase = map[string]*prometheus.Desc{
@@ -34,30 +40,34 @@ var pgDatabase = map[string]*prometheus.Desc{
 ),
 }
 
-func (PGDatabaseCollector) Update(ctx context.Context, db *sql.DB, server string) ([]prometheus.Metric, error) {
-metrics := []prometheus.Metric{}
+func (PGDatabaseCollector) Update(ctx context.Context, server *server, ch chan<- prometheus.Metric) error {
+db, err := server.GetDB()
+if err != nil {
+return err
+}
 rows, err := db.QueryContext(ctx,
 `SELECT pg_database.datname
 ,pg_database_size(pg_database.datname)
 FROM pg_database;`)
 if err != nil {
-return metrics, err
+return err
 }
 defer rows.Close()
 
 for rows.Next() {
 var datname string
 var size int64
 if err := rows.Scan(&datname, &size); err != nil {
-return metrics, err
+return err
 }
-metrics = append(metrics, prometheus.MustNewConstMetric(
+
+ch <- prometheus.MustNewConstMetric(
 pgDatabase["size_bytes"],
-prometheus.GaugeValue, float64(size), datname, server,
-))
+prometheus.GaugeValue, float64(size), datname, server.GetName(),
+)
 }
 if err := rows.Err(); err != nil {
-return metrics, err
+return err
 }
-return metrics, nil
+return nil
 }