WIP: Add prelim multi-target support

- Remove multi server support from new collector package - Add http handler for multi-target support Signed-off-by: Joe Adams <github@joeadams.io>

Author: sysadmind

cmd/postgres_exporter/datasource.go

@@ -162,6 +162,12 @@ func getDataSources() ([]string, error) {
 uri = os.Getenv("DATA_SOURCE_URI")
 }
 
+// No datasources found. This allows us to support the multi-target pattern
+// withouth an explicit datasource.
+if uri == "" {
+return []string{}, nil
+}
+
 dsn = "postgresql://" + ui + "@" + uri
 
 return []string{dsn}, nil

cmd/postgres_exporter/main.go

@@ -85,16 +85,17 @@ func main() {
 return
 }
 
-dsn, err := getDataSources()
+dsns, err := getDataSources()
 if err != nil {
 level.Error(logger).Log("msg", "Failed reading data sources", "err", err.Error())
 os.Exit(1)
 }
 
-if len(dsn) == 0 {
-level.Error(logger).Log("msg", "Couldn't find environment variables describing the datasource to use")
-os.Exit(1)
-}
+// TODO(@sysadmind): Remove this with multi-target support
+// if len(dsn) == 0 {
+//	level.Error(logger).Log("msg", "Couldn't find environment variables describing the datasource to use")
+//	os.Exit(1)
+// }
 
 opts := []ExporterOpt{
 DisableDefaultMetrics(*disableDefaultMetrics),
@@ -106,7 +107,7 @@ func main() {
 IncludeDatabases(*includeDatabases),
 }
 
-exporter := NewExporter(dsn, opts...)
+exporter := NewExporter(dsns, opts...)
 defer func() {
 exporter.servers.Close()
 }()
@@ -115,23 +116,31 @@ func main() {
 
 prometheus.MustRegister(exporter)
 
+// TODO(@sysadmind): Remove this with multi-target support. We are removing multiple DSN support
+dsn := ""
+if len(dsns) > 0 {
+dsn = dsns[0]
+}
+
 pe, err := collector.NewPostgresCollector(
 logger,
 dsn,
 []string{},
 )
 if err != nil {
 level.Error(logger).Log("msg", "Failed to create PostgresCollector", "err", err.Error())
-os.Exit(1)
+} else {
+prometheus.MustRegister(pe)
 }
-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
 w.Write(landingPage) // nolint: errcheck
 })
 
+http.HandleFunc("/probe", handleProbe(logger))
+
 level.Info(logger).Log("msg", "Listening on address", "address", *listenAddress)
 srv := &http.Server{Addr: *listenAddress}
 if err := web.ListenAndServe(srv, *webConfig, logger); err != nil {

cmd/postgres_exporter/probe.go

@@ -0,0 +1,91 @@
+// 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 main
+
+import (
+"net/http"
+"time"
+
+"github.com/go-kit/log"
+"github.com/prometheus-community/postgres_exporter/collector"
+"github.com/prometheus/client_golang/prometheus"
+"github.com/prometheus/client_golang/prometheus/promhttp"
+)
+
+func handleProbe(logger log.Logger) http.HandlerFunc {
+return func(w http.ResponseWriter, r *http.Request) {
+ctx := r.Context()
+params := r.URL.Query()
+target := params.Get("target")
+if target == "" {
+http.Error(w, "target is required", http.StatusBadRequest)
+return
+}
+
+// TODO: Timeout
+// TODO: Auth Module
+
+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)
+_ = tl
+
+start := time.Now()
+registry := prometheus.NewRegistry()
+registry.MustRegister(probeSuccessGauge)
+registry.MustRegister(probeDurationGauge)
+
+// TODO(@sysadmind): this is a temp hack until we have a proper auth module
+target = "postgres://postgres:test@localhost:5432/circle_test?sslmode=disable"
+
+// Run the probe
+pc, err := collector.NewProbeCollector(tl, registry, target)
+if err != nil {
+probeSuccessGauge.Set(0)
+probeDurationGauge.Set(time.Since(start).Seconds())
+http.Error(w, err.Error(), http.StatusInternalServerError)
+return
+}
+_ = ctx
+
+// TODO: Which way should this be? Register or handle the collection manually?
+// Also, what about the context?
+
+// Option 1: Register the collector
+registry.MustRegister(pc)
+
+// Option 2: Handle the collection manually. This allows us to collect duration metrics.
+// The collectors themselves already support their own duration metrics.
+// err = pc.Update(ctx)
+// if err != nil {
+//	probeSuccessGauge.Set(0)
+// } else {
+//	probeSuccessGauge.Set(1)
+// }
+
+duration := time.Since(start).Seconds()
+probeDurationGauge.Set(duration)
+
+// TODO check success, etc
+h := promhttp.HandlerFor(registry, promhttp.HandlerOpts{})
+h.ServeHTTP(w, r)
+}
+}

collector/collector.go

@@ -15,6 +15,7 @@ package collector
 
 import (
 "context"
+"database/sql"
 "errors"
 "fmt"
 "sync"
@@ -58,7 +59,7 @@ var (
 )
 
 type Collector interface {
-Update(ctx context.Context, server *server, ch chan<- prometheus.Metric) error
+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)) {
@@ -86,13 +87,13 @@ type PostgresCollector struct {
 Collectors map[string]Collector
 logger log.Logger
 
-servers map[string]*server
+db *sql.DB
 }
 
 type Option func(*PostgresCollector) error
 
 // NewPostgresCollector creates a new PostgresCollector.
-func NewPostgresCollector(logger log.Logger, dsns []string, filters []string, options ...Option) (*PostgresCollector, error) {
+func NewPostgresCollector(logger log.Logger, dsn string, filters []string, options ...Option) (*PostgresCollector, error) {
 p := &PostgresCollector{
 logger: logger,
 }
@@ -136,17 +137,18 @@ func NewPostgresCollector(logger log.Logger, dsns []string, filters []string, op
 
 p.Collectors = collectors
 
-servers := make(map[string]*server)
-for _, dsn := range dsns {
-s, err := makeServer(dsn)
-if err != nil {
-return nil, err
-}
+if dsn == "" {
+return nil, errors.New("empty dsn")
+}
 
-servers[dsn] = s
+db, err := sql.Open("postgres", dsn)
+if err != nil {
+return nil, err
 }
+db.SetMaxOpenConns(1)
+db.SetMaxIdleConns(1)
 
-p.servers = servers
+p.db = db
 
 return p, nil
 }
@@ -160,32 +162,20 @@ func (p PostgresCollector) Describe(ch chan<- *prometheus.Desc) {
 // Collect implements the prometheus.Collector interface.
 func (p PostgresCollector) Collect(ch chan<- prometheus.Metric) {
 ctx := context.TODO()
-wg := sync.WaitGroup{}
-wg.Add(len(p.servers))
-for _, s := range p.servers {
-go func(s *server) {
-p.subCollect(ctx, s, ch)
-wg.Done()
-}(s)
-}
-wg.Wait()
-}
-
-func (p PostgresCollector) subCollect(ctx context.Context, server *server, ch chan<- prometheus.Metric) {
 wg := sync.WaitGroup{}
 wg.Add(len(p.Collectors))
 for name, c := range p.Collectors {
 go func(name string, c Collector) {
-execute(ctx, name, c, server, ch, p.logger)
+execute(ctx, name, c, p.db, ch, p.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) {
+func execute(ctx context.Context, name string, c Collector, db *sql.DB, ch chan<- prometheus.Metric, logger log.Logger) {
 begin := time.Now()
-err := c.Update(ctx, s, ch)
+err := c.Update(ctx, db, ch)
 duration := time.Since(begin)
 var success float64
 

collector/pg_database.go

@@ -15,6 +15,7 @@ package collector
 
 import (
 "context"
+"database/sql"
 
 "github.com/go-kit/log"
 "github.com/prometheus/client_golang/prometheus"
@@ -36,15 +37,11 @@ var pgDatabase = map[string]*prometheus.Desc{
 "size_bytes": prometheus.NewDesc(
 "pg_database_size_bytes",
 "Disk space used by the database",
-[]string{"datname", "server"}, nil,
+[]string{"datname"}, nil,
 ),
 }
 
-func (PGDatabaseCollector) Update(ctx context.Context, server *server, ch chan<- prometheus.Metric) error {
-db, err := server.GetDB()
-if err != nil {
-return err
-}
+func (PGDatabaseCollector) Update(ctx context.Context, db *sql.DB, ch chan<- prometheus.Metric) error {
 rows, err := db.QueryContext(ctx,
 `SELECT pg_database.datname
 ,pg_database_size(pg_database.datname)
@@ -63,7 +60,7 @@ func (PGDatabaseCollector) Update(ctx context.Context, server *server, ch chan<-
 
 ch <- prometheus.MustNewConstMetric(
 pgDatabase["size_bytes"],
-prometheus.GaugeValue, float64(size), datname, server.GetName(),
+prometheus.GaugeValue, float64(size), datname,
 )
 }
 if err := rows.Err(); err != nil {