Refactor code into logical files

cmd/postgres_exporter/datasource.go

@@ -0,0 +1,168 @@
+// Copyright 2021 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 (
+	"fmt"
+	"io/ioutil"
+	"net/url"
+	"os"
+	"regexp"
+	"strings"
+
+	"github.com/go-kit/kit/log/level"
+	"github.com/prometheus/client_golang/prometheus"
+)
+
+func (e *Exporter) discoverDatabaseDSNs() []string {
+	// connstring syntax is complex (and not sure if even regular).
+	// we don't need to parse it, so just superficially validate that it starts
+	// with a valid-ish keyword pair
+	connstringRe := regexp.MustCompile(`^ *[a-zA-Z0-9]+ *= *[^= ]+`)
+
+	dsns := make(map[string]struct{})
+	for _, dsn := range e.dsn {
+		var dsnURI *url.URL
+		var dsnConnstring string
+
+		if strings.HasPrefix(dsn, "postgresql://") {
+			var err error
+			dsnURI, err = url.Parse(dsn)
+			if err != nil {
+				level.Error(logger).Log("msg", "Unable to parse DSN as URI", "dsn", loggableDSN(dsn), "err", err)
+				continue
+			}
+		} else if connstringRe.MatchString(dsn) {
+			dsnConnstring = dsn
+		} else {
+			level.Error(logger).Log("msg", "Unable to parse DSN as either URI or connstring", "dsn", loggableDSN(dsn))
+			continue
+		}
+
+		server, err := e.servers.GetServer(dsn)
+		if err != nil {
+			level.Error(logger).Log("msg", "Error opening connection to database", "dsn", loggableDSN(dsn), "err", err)
+			continue
+		}
+		dsns[dsn] = struct{}{}
+
+		// If autoDiscoverDatabases is true, set first dsn as master database (Default: false)
+		server.master = true
+
+		databaseNames, err := queryDatabases(server)
+		if err != nil {
+			level.Error(logger).Log("msg", "Error querying databases", "dsn", loggableDSN(dsn), "err", err)
+			continue
+		}
+		for _, databaseName := range databaseNames {
+			if contains(e.excludeDatabases, databaseName) {
+				continue
+			}
+
+			if len(e.includeDatabases) != 0 && !contains(e.includeDatabases, databaseName) {
+				continue
+			}
+
+			if dsnURI != nil {
+				dsnURI.Path = databaseName
+				dsn = dsnURI.String()
+			} else {
+				// replacing one dbname with another is complicated.
+				// just append new dbname to override.
+				dsn = fmt.Sprintf("%s dbname=%s", dsnConnstring, databaseName)
+			}
+			dsns[dsn] = struct{}{}
+		}
+	}
+
+	result := make([]string, len(dsns))
+	index := 0
+	for dsn := range dsns {
+		result[index] = dsn
+		index++
+	}
+
+	return result
+}
+
+func (e *Exporter) scrapeDSN(ch chan<- prometheus.Metric, dsn string) error {
+	server, err := e.servers.GetServer(dsn)
+
+	if err != nil {
+		return &ErrorConnectToServer{fmt.Sprintf("Error opening connection to database (%s): %s", loggableDSN(dsn), err.Error())}
+	}
+
+	// Check if autoDiscoverDatabases is false, set dsn as master database (Default: false)
+	if !e.autoDiscoverDatabases {
+		server.master = true
+	}
+
+	// Check if map versions need to be updated
+	if err := e.checkMapVersions(ch, server); err != nil {
+		level.Warn(logger).Log("msg", "Proceeding with outdated query maps, as the Postgres version could not be determined", "err", err)
+	}
+
+	return server.Scrape(ch, e.disableSettingsMetrics)
+}
+
+// try to get the DataSource
+// DATA_SOURCE_NAME always wins so we do not break older versions
+// reading secrets from files wins over secrets in environment variables
+// DATA_SOURCE_NAME > DATA_SOURCE_{USER|PASS}_FILE > DATA_SOURCE_{USER|PASS}
+func getDataSources() ([]string, error) {
+	var dsn = os.Getenv("DATA_SOURCE_NAME")
+	if len(dsn) != 0 {
+		return strings.Split(dsn, ","), nil
+	}
+
+	var user, pass, uri string
+
+	dataSourceUserFile := os.Getenv("DATA_SOURCE_USER_FILE")
+	if len(dataSourceUserFile) != 0 {
+		fileContents, err := ioutil.ReadFile(dataSourceUserFile)
+		if err != nil {
+			return nil, fmt.Errorf("failed loading data source user file %s: %s", dataSourceUserFile, err.Error())
+		}
+		user = strings.TrimSpace(string(fileContents))
+	} else {
+		user = os.Getenv("DATA_SOURCE_USER")
+	}
+
+	dataSourcePassFile := os.Getenv("DATA_SOURCE_PASS_FILE")
+	if len(dataSourcePassFile) != 0 {
+		fileContents, err := ioutil.ReadFile(dataSourcePassFile)
+		if err != nil {
+			return nil, fmt.Errorf("failed loading data source pass file %s: %s", dataSourcePassFile, err.Error())
+		}
+		pass = strings.TrimSpace(string(fileContents))
+	} else {
+		pass = os.Getenv("DATA_SOURCE_PASS")
+	}
+
+	ui := url.UserPassword(user, pass).String()
+	dataSrouceURIFile := os.Getenv("DATA_SOURCE_URI_FILE")
+	if len(dataSrouceURIFile) != 0 {
+		fileContents, err := ioutil.ReadFile(dataSrouceURIFile)
+		if err != nil {
+			return nil, fmt.Errorf("failed loading data source URI file %s: %s", dataSrouceURIFile, err.Error())
+		}
+		uri = strings.TrimSpace(string(fileContents))
+	} else {
+		uri = os.Getenv("DATA_SOURCE_URI")
+	}
+
+	dsn = "postgresql://" + ui + "@" + uri
+
+	return []string{dsn}, nil
+}

cmd/postgres_exporter/main.go

@@ -0,0 +1,129 @@
+// Copyright 2021 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"
+	"os"
+
+	"github.com/go-kit/kit/log"
+	"github.com/go-kit/kit/log/level"
+	"github.com/prometheus/client_golang/prometheus"
+	"github.com/prometheus/client_golang/prometheus/promhttp"
+	"github.com/prometheus/common/promlog"
+	"github.com/prometheus/common/promlog/flag"
+	"github.com/prometheus/common/version"
+	"github.com/prometheus/exporter-toolkit/web"
+	webflag "github.com/prometheus/exporter-toolkit/web/kingpinflag"
+	"gopkg.in/alecthomas/kingpin.v2"
+)
+
+var (
+	listenAddress          = kingpin.Flag("web.listen-address", "Address to listen on for web interface and telemetry.").Default(":9187").Envar("PG_EXPORTER_WEB_LISTEN_ADDRESS").String()
+	webConfig              = webflag.AddFlags(kingpin.CommandLine)
+	metricPath             = kingpin.Flag("web.telemetry-path", "Path under which to expose metrics.").Default("/metrics").Envar("PG_EXPORTER_WEB_TELEMETRY_PATH").String()
+	disableDefaultMetrics  = kingpin.Flag("disable-default-metrics", "Do not include default metrics.").Default("false").Envar("PG_EXPORTER_DISABLE_DEFAULT_METRICS").Bool()
+	disableSettingsMetrics = kingpin.Flag("disable-settings-metrics", "Do not include pg_settings metrics.").Default("false").Envar("PG_EXPORTER_DISABLE_SETTINGS_METRICS").Bool()
+	autoDiscoverDatabases  = kingpin.Flag("auto-discover-databases", "Whether to discover the databases on a server dynamically.").Default("false").Envar("PG_EXPORTER_AUTO_DISCOVER_DATABASES").Bool()
+	queriesPath            = kingpin.Flag("extend.query-path", "Path to custom queries to run.").Default("").Envar("PG_EXPORTER_EXTEND_QUERY_PATH").String()
+	onlyDumpMaps           = kingpin.Flag("dumpmaps", "Do not run, simply dump the maps.").Bool()
+	constantLabelsList     = kingpin.Flag("constantLabels", "A list of label=value separated by comma(,).").Default("").Envar("PG_EXPORTER_CONSTANT_LABELS").String()
+	excludeDatabases       = kingpin.Flag("exclude-databases", "A list of databases to remove when autoDiscoverDatabases is enabled").Default("").Envar("PG_EXPORTER_EXCLUDE_DATABASES").String()
+	includeDatabases       = kingpin.Flag("include-databases", "A list of databases to include when autoDiscoverDatabases is enabled").Default("").Envar("PG_EXPORTER_INCLUDE_DATABASES").String()
+	metricPrefix           = kingpin.Flag("metric-prefix", "A metric prefix can be used to have non-default (not \"pg\") prefixes for each of the metrics").Default("pg").Envar("PG_EXPORTER_METRIC_PREFIX").String()
+	logger                 = log.NewNopLogger()
+)
+
+// Metric name parts.
+const (
+	// Namespace for all metrics.
+	namespace = "pg"
+	// Subsystems.
+	exporter = "exporter"
+	// The name of the exporter.
+	exporterName = "postgres_exporter"
+	// Metric label used for static string data thats handy to send to Prometheus
+	// e.g. version
+	staticLabelName = "static"
+	// Metric label used for server identification.
+	serverLabelName = "server"
+)
+
+func main() {
+	kingpin.Version(version.Print(exporterName))
+	promlogConfig := &promlog.Config{}
+	flag.AddFlags(kingpin.CommandLine, promlogConfig)
+	kingpin.HelpFlag.Short('h')
+	kingpin.Parse()
+	logger = promlog.New(promlogConfig)
+
+	// landingPage contains the HTML served at '/'.
+	// TODO: Make this nicer and more informative.
+	var landingPage = []byte(`<html>
+	<head><title>Postgres exporter</title></head>
+	<body>
+	<h1>Postgres exporter</h1>
+	<p><a href='` + *metricPath + `'>Metrics</a></p>
+	</body>
+	</html>
+	`)
+
+	if *onlyDumpMaps {
+		dumpMaps()
+		return
+	}
+
+	dsn, 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)
+	}
+
+	opts := []ExporterOpt{
+		DisableDefaultMetrics(*disableDefaultMetrics),
+		DisableSettingsMetrics(*disableSettingsMetrics),
+		AutoDiscoverDatabases(*autoDiscoverDatabases),
+		WithUserQueriesPath(*queriesPath),
+		WithConstantLabels(*constantLabelsList),
+		ExcludeDatabases(*excludeDatabases),
+		IncludeDatabases(*includeDatabases),
+	}
+
+	exporter := NewExporter(dsn, opts...)
+	defer func() {
+		exporter.servers.Close()
+	}()
+
+	prometheus.MustRegister(version.NewCollector(exporterName))
+
+	prometheus.MustRegister(exporter)
+
+	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
+	})
+
+	level.Info(logger).Log("msg", "Listening on address", "address", *listenAddress)
+	srv := &http.Server{Addr: *listenAddress}
+	if err := web.ListenAndServe(srv, *webConfig, logger); err != nil {
+		level.Error(logger).Log("msg", "Error running HTTP server", "err", err)
+		os.Exit(1)
+	}
+}