Fix exclude-databases for collector package

The pg_database collector was not respecting the --exclude-databases flag and causing problems where databases were not accessible. This now respects the list of databases to exclude.

- Adjusts the Collector create func to take a config struct instead of a logger. This allows more changes like this in the future. I figured we would need to do this at some point but I wasn't sure if we could hold off.
- Split the database size collection to a separate query when database is not excluded.
- Comment some probe code that was not useful/accurate

Signed-off-by: Joe Adams <github@joeadams.io>

cmd/postgres_exporter/main.go

@@ -14,8 +14,10 @@
 package main
 
 import (
+	"fmt"
 	"net/http"
 	"os"
+	"strings"
 
 	"github.com/go-kit/log"
 	"github.com/go-kit/log/level"
@@ -101,13 +103,16 @@ func main() {
 		os.Exit(1)
 	}
 
+	excludedDatabases := strings.Split(*excludeDatabases, ",")
+	logger.Log("msg", "Excluded databases", "databases", fmt.Sprintf("%v", excludedDatabases))
+
 	opts := []ExporterOpt{
 		DisableDefaultMetrics(*disableDefaultMetrics),
 		DisableSettingsMetrics(*disableSettingsMetrics),
 		AutoDiscoverDatabases(*autoDiscoverDatabases),
 		WithUserQueriesPath(*queriesPath),
 		WithConstantLabels(*constantLabelsList),
-		ExcludeDatabases(*excludeDatabases),
+		ExcludeDatabases(excludedDatabases),
 		IncludeDatabases(*includeDatabases),
 	}
 
@@ -128,6 +133,7 @@ func main() {
 
 	pe, err := collector.NewPostgresCollector(
 		logger,
+		excludedDatabases,
 		dsn,
 		[]string{},
 	)
@@ -143,7 +149,7 @@ func main() {
 		w.Write(landingPage)                                       // nolint: errcheck
 	})
 
-	http.HandleFunc("/probe", handleProbe(logger))
+	http.HandleFunc("/probe", handleProbe(logger, excludedDatabases))
 
 	srv := &http.Server{}
 	if err := web.ListenAndServe(srv, webConfig, logger); err != nil {

cmd/postgres_exporter/postgres_exporter.go

@@ -484,9 +484,9 @@ func AutoDiscoverDatabases(b bool) ExporterOpt {
 }
 
 // ExcludeDatabases allows to filter out result from AutoDiscoverDatabases
-func ExcludeDatabases(s string) ExporterOpt {
+func ExcludeDatabases(s []string) ExporterOpt {
 	return func(e *Exporter) {
-		e.excludeDatabases = strings.Split(s, ",")
+		e.excludeDatabases = s
 	}
 }
 

cmd/postgres_exporter/probe.go

@@ -16,7 +16,6 @@ package main
 import (
 	"fmt"
 	"net/http"
-	"time"
 
 	"github.com/go-kit/log"
 	"github.com/go-kit/log/level"
@@ -26,7 +25,7 @@ import (
 	"github.com/prometheus/client_golang/prometheus/promhttp"
 )
 
-func handleProbe(logger log.Logger) http.HandlerFunc {
+func handleProbe(logger log.Logger, excludeDatabases []string) http.HandlerFunc {
 	return func(w http.ResponseWriter, r *http.Request) {
 		ctx := r.Context()
 		conf := c.GetConfig()
@@ -62,21 +61,21 @@ func handleProbe(logger log.Logger) http.HandlerFunc {
 
 		// TODO(@sysadmind): Timeout
 
-		probeSuccessGauge := prometheus.NewGauge(prometheus.GaugeOpts{
-			Name: "probe_success",
-			Help: "Displays whether or not the probe was a success",
-		})
-		probeDurationGauge := prometheus.NewGauge(prometheus.GaugeOpts{
-			Name: "probe_duration_seconds",
-			Help: "Returns how long the probe took to complete in seconds",
-		})
+		// probeSuccessGauge := prometheus.NewGauge(prometheus.GaugeOpts{
+		// 	Name: "probe_success",
+		// 	Help: "Displays whether or not the probe was a success",
+		// })
+		// probeDurationGauge := prometheus.NewGauge(prometheus.GaugeOpts{
+		// 	Name: "probe_duration_seconds",
+		// 	Help: "Returns how long the probe took to complete in seconds",
+		// })
 
 		tl := log.With(logger, "target", target)
 
-		start := time.Now()
+		// start := time.Now()
 		registry := prometheus.NewRegistry()
-		registry.MustRegister(probeSuccessGauge)
-		registry.MustRegister(probeDurationGauge)
+		// registry.MustRegister(probeSuccessGauge)
+		// registry.MustRegister(probeDurationGauge)
 
 		opts := []ExporterOpt{
 			DisableDefaultMetrics(*disableDefaultMetrics),
@@ -96,10 +95,10 @@ func handleProbe(logger log.Logger) http.HandlerFunc {
 		registry.MustRegister(exporter)
 
 		// Run the probe
-		pc, err := collector.NewProbeCollector(tl, registry, dsn)
+		pc, err := collector.NewProbeCollector(tl, excludeDatabases, registry, dsn)
 		if err != nil {
-			probeSuccessGauge.Set(0)
-			probeDurationGauge.Set(time.Since(start).Seconds())
+			// probeSuccessGauge.Set(0)
+			// probeDurationGauge.Set(time.Since(start).Seconds())
 			http.Error(w, err.Error(), http.StatusInternalServerError)
 			return
 		}
@@ -115,10 +114,6 @@ func handleProbe(logger log.Logger) http.HandlerFunc {
 
 		registry.MustRegister(pc)
 
-		duration := time.Since(start).Seconds()
-		probeDurationGauge.Set(duration)
-		probeSuccessGauge.Set(1)
-
 		// TODO check success, etc
 		h := promhttp.HandlerFor(registry, promhttp.HandlerOpts{})
 		h.ServeHTTP(w, r)

collector/collector.go

@@ -28,7 +28,7 @@ import (
 )
 
 var (
-	factories              = make(map[string]func(logger log.Logger) (Collector, error))
+	factories              = make(map[string]func(collectorConfig) (Collector, error))
 	initiatedCollectorsMtx = sync.Mutex{}
 	initiatedCollectors    = make(map[string]Collector)
 	collectorState         = make(map[string]*bool)
@@ -62,7 +62,12 @@ type Collector interface {
 	Update(ctx context.Context, db *sql.DB, ch chan<- prometheus.Metric) error
 }
 
-func registerCollector(name string, isDefaultEnabled bool, createFunc func(logger log.Logger) (Collector, error)) {
+type collectorConfig struct {
+	logger           log.Logger
+	excludeDatabases []string
+}
+
+func registerCollector(name string, isDefaultEnabled bool, createFunc func(collectorConfig) (Collector, error)) {
 	var helpDefaultState string
 	if isDefaultEnabled {
 		helpDefaultState = "enabled"
@@ -93,7 +98,7 @@ type PostgresCollector struct {
 type Option func(*PostgresCollector) error
 
 // NewPostgresCollector creates a new PostgresCollector.
-func NewPostgresCollector(logger log.Logger, dsn string, filters []string, options ...Option) (*PostgresCollector, error) {
+func NewPostgresCollector(logger log.Logger, excludeDatabases []string, dsn string, filters []string, options ...Option) (*PostgresCollector, error) {
 	p := &PostgresCollector{
 		logger: logger,
 	}
@@ -126,7 +131,10 @@ func NewPostgresCollector(logger log.Logger, dsn string, filters []string, optio
 		if collector, ok := initiatedCollectors[key]; ok {
 			collectors[key] = collector
 		} else {
-			collector, err := factories[key](log.With(logger, "collector", key))
+			collector, err := factories[key](collectorConfig{
+				logger:           log.With(logger, "collector", key),
+				excludeDatabases: excludeDatabases,
+			})
 			if err != nil {
 				return nil, err
 			}

collector/pg_database.go

@@ -26,11 +26,19 @@ func init() {
 }
 
 type PGDatabaseCollector struct {
-	log log.Logger
+	log               log.Logger
+	excludedDatabases []string
 }
 
-func NewPGDatabaseCollector(logger log.Logger) (Collector, error) {
-	return &PGDatabaseCollector{log: logger}, nil
+func NewPGDatabaseCollector(config collectorConfig) (Collector, error) {
+	exclude := config.excludeDatabases
+	if exclude == nil {
+		exclude = []string{}
+	}
+	return &PGDatabaseCollector{
+		log:               config.logger,
+		excludedDatabases: exclude,
+	}, nil
 }
 
 var pgDatabase = map[string]*prometheus.Desc{
@@ -41,20 +49,49 @@ var pgDatabase = map[string]*prometheus.Desc{
 	),
 }
 
-func (PGDatabaseCollector) Update(ctx context.Context, db *sql.DB, ch chan<- prometheus.Metric) error {
+// Update implements Collector and exposes database size.
+// It is called by the Prometheus registry when collecting metrics.
+// The list of databases is retrieved from pg_database and filtered
+// by the excludeDatabase config parameter. The tradeoff here is that
+// we have to query the list of databases and then query the size of
+// each database individually. This is because we can't filter the
+// list of databases in the query because the list of excluded
+// databases is dynamic.
+func (c PGDatabaseCollector) Update(ctx context.Context, db *sql.DB, ch chan<- prometheus.Metric) error {
+	// Query the list of databases
 	rows, err := db.QueryContext(ctx,
 		`SELECT pg_database.datname
-		,pg_database_size(pg_database.datname)
-		FROM pg_database;`)
+		FROM pg_database;
+		`,
+	)
 	if err != nil {
 		return err
 	}
 	defer rows.Close()
 
+	var databases []string
+
 	for rows.Next() {
 		var datname string
+		if err := rows.Scan(&datname); err != nil {
+			return err
+		}
+
+		// Ignore excluded databases
+		// Filtering is done here instead of in the query to avoid
+		// a complicated NOT IN query with a variable number of parameters
+		if sliceContains(c.excludedDatabases, datname) {
+			continue
+		}
+
+		databases = append(databases, datname)
+	}
+
+	// Query the size of the databases
+	for _, datname := range databases {
 		var size int64
-		if err := rows.Scan(&datname, &size); err != nil {
+		err = db.QueryRowContext(ctx, "SELECT pg_database_size($1)", datname).Scan(&size)
+		if err != nil {
 			return err
 		}
 
@@ -68,3 +105,12 @@ func (PGDatabaseCollector) Update(ctx context.Context, db *sql.DB, ch chan<- pro
 	}
 	return nil
 }
+
+func sliceContains(slice []string, s string) bool {
+	for _, item := range slice {
+		if item == s {
+			return true
+		}
+	}
+	return false
+}

collector/pg_stat_bgwriter.go

@@ -18,7 +18,6 @@ import (
 	"database/sql"
 	"time"
 
-	"github.com/go-kit/log"
 	"github.com/prometheus/client_golang/prometheus"
 )
 
@@ -29,7 +28,7 @@ func init() {
 type PGStatBGWriterCollector struct {
 }
 
-func NewPGStatBGWriterCollector(logger log.Logger) (Collector, error) {
+func NewPGStatBGWriterCollector(collectorConfig) (Collector, error) {
 	return &PGStatBGWriterCollector{}, nil
 }
 

collector/probe.go

@@ -30,7 +30,7 @@ type ProbeCollector struct {
 	db         *sql.DB
 }
 
-func NewProbeCollector(logger log.Logger, registry *prometheus.Registry, dsn config.DSN) (*ProbeCollector, error) {
+func NewProbeCollector(logger log.Logger, excludeDatabases []string, registry *prometheus.Registry, dsn config.DSN) (*ProbeCollector, error) {
 	collectors := make(map[string]Collector)
 	initiatedCollectorsMtx.Lock()
 	defer initiatedCollectorsMtx.Unlock()
@@ -45,7 +45,11 @@ func NewProbeCollector(logger log.Logger, registry *prometheus.Registry, dsn con
 		if collector, ok := initiatedCollectors[key]; ok {
 			collectors[key] = collector
 		} else {
-			collector, err := factories[key](log.With(logger, "collector", key))
+			collector, err := factories[key](
+				collectorConfig{
+					logger:           log.With(logger, "collector", key),
+					excludeDatabases: excludeDatabases,
+				})
 			if err != nil {
 				return nil, err
 			}