Add pg_database collector

CHANGELOG.md

@@ -1,3 +1,7 @@
+## master / unreleased
+
+* [ENHANCEMENT] Add pg_database_size_bytes metric #613
+
 ## 0.10.1 / 2022-01-14
 
 * [BUGFIX] Fix broken log-level for values other than debug. #560

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

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

@@ -0,0 +1,73 @@
+// 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"
+
+	"github.com/go-kit/log"
+	"github.com/prometheus/client_golang/prometheus"
+)
+
+func init() {
+	registerCollector("database", defaultEnabled, NewPGDatabaseCollector)
+}
+
+type PGDatabaseCollector struct {
+	log log.Logger
+}
+
+func NewPGDatabaseCollector(logger log.Logger) (Collector, error) {
+	return &PGDatabaseCollector{log: logger}, nil
+}
+
+var pgDatabase = map[string]*prometheus.Desc{
+	"size_bytes": prometheus.NewDesc(
+		"pg_database_size_bytes",
+		"Disk space used by the database",
+		[]string{"datname", "server"}, nil,
+	),
+}
+
+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 err
+	}
+	defer rows.Close()
+
+	for rows.Next() {
+		var datname string
+		var size int64
+		if err := rows.Scan(&datname, &size); err != nil {
+			return err
+		}
+
+		ch <- prometheus.MustNewConstMetric(
+			pgDatabase["size_bytes"],
+			prometheus.GaugeValue, float64(size), datname, server.GetName(),
+		)
+	}
+	if err := rows.Err(); err != nil {
+		return err
+	}
+	return nil
+}