Merge pull request #21 from ncabatoff/feature-queries_from_file

Feature queries from file

postgres_exporter.go

@@ -4,12 +4,14 @@ import (
 	"database/sql"
 	"flag"
 	"fmt"
+	"io/ioutil"
 	"math"
 	"net/http"
 	"os"
+	"strconv"
 	"time"
 
-	"strconv"
+	"gopkg.in/yaml.v2"
 
 	_ "github.com/lib/pq"
 	"github.com/prometheus/client_golang/prometheus"
@@ -27,6 +29,14 @@ var (
 		"web.telemetry-path", "/metrics",
 		"Path under which to expose metrics.",
 	)
+	queriesPath = flag.String(
+		"extend.query-path", "",
+		"Path to custom queries to run.",
+	)
+	onlyDumpMaps = flag.Bool(
+		"dumpmaps", false,
+		"Do not run, simply dump the maps.",
+	)
 )
 
 // Metric name parts.
@@ -102,6 +112,21 @@ var variableMaps = map[string]map[string]ColumnMapping{
 	},
 }
 
+func dumpMaps() {
+	for name, cmap := range metricMaps {
+		query, ok := queryOverrides[name]
+		if ok {
+			fmt.Printf("%s: %s\n", name, query)
+		} else {
+			fmt.Println(name)
+		}
+		for column, details := range cmap {
+			fmt.Printf("  %-40s %v\n", column, details)
+		}
+		fmt.Println()
+	}
+}
+
 var metricMaps = map[string]map[string]ColumnMapping{
 	"pg_stat_bgwriter": map[string]ColumnMapping{
 		"checkpoints_timed":     {COUNTER, "Number of scheduled checkpoints that have been performed", nil},
@@ -223,6 +248,68 @@ var queryOverrides = map[string]string{
       ON tmp.state = tmp2.state AND pg_database.datname = tmp2.datname`,
 }
 
+// Add queries to the metricMaps and queryOverrides maps
+func addQueries(queriesPath string) (err error) {
+	var extra map[string]interface{}
+
+	content, err := ioutil.ReadFile(queriesPath)
+	if err != nil {
+		return err
+	}
+
+	err = yaml.Unmarshal(content, &extra)
+	if err != nil {
+		return err
+	}
+
+	for metric, specs := range extra {
+		for key, value := range specs.(map[interface{}]interface{}) {
+			switch key.(string) {
+			case "query":
+				query := value.(string)
+				queryOverrides[metric] = query
+
+			case "metrics":
+				for _, c := range value.([]interface{}) {
+					column := c.(map[interface{}]interface{})
+
+					for n, a := range column {
+						var cmap ColumnMapping
+
+						metric_map, ok := metricMaps[metric]
+						if !ok {
+							metric_map = make(map[string]ColumnMapping)
+						}
+
+						name := n.(string)
+
+						for attr_key, attr_val := range a.(map[interface{}]interface{}) {
+							switch attr_key.(string) {
+							case "usage":
+								usage, err := stringToColumnUsage(attr_val.(string))
+								if err != nil {
+									return err
+								}
+								cmap.usage = usage
+							case "description":
+								cmap.description = attr_val.(string)
+							}
+						}
+
+						cmap.mapping = nil
+
+						metric_map[name] = cmap
+
+						metricMaps[metric] = metric_map
+					}
+				}
+			}
+		}
+	}
+
+	return
+}
+
 // Turn the MetricMap column mapping into a prometheus descriptor mapping.
 func makeDescMap(metricMaps map[string]map[string]ColumnMapping) map[string]MetricMapNamespace {
 	var metricMap = make(map[string]MetricMapNamespace)
@@ -317,6 +404,34 @@ func makeDescMap(metricMaps map[string]map[string]ColumnMapping) map[string]Metr
 	return metricMap
 }
 
+// convert a string to the corresponding ColumnUsage
+func stringToColumnUsage(s string) (u ColumnUsage, err error) {
+	switch s {
+	case "DISCARD":
+		u = DISCARD
+
+	case "LABEL":
+		u = LABEL
+
+	case "COUNTER":
+		u = COUNTER
+
+	case "GAUGE":
+		u = GAUGE
+
+	case "MAPPEDMETRIC":
+		u = MAPPEDMETRIC
+
+	case "DURATION":
+		u = DURATION
+
+	default:
+		err = fmt.Errorf("wrong ColumnUsage given : %s", s)
+	}
+
+	return
+}
+
 // Convert database.sql types to float64s for Prometheus consumption. Null types are mapped to NaN. string and []byte
 // types are mapped as NaN and !ok
 func dbToFloat64(t interface{}) (float64, bool) {
@@ -590,6 +705,18 @@ func (e *Exporter) scrape(ch chan<- prometheus.Metric) {
 func main() {
 	flag.Parse()
 
+	if *queriesPath != "" {
+		err := addQueries(*queriesPath)
+		if err != nil {
+			log.Warnln("Unparseable queries file - discarding merge: ", *queriesPath, err)
+		}
+	}
+
+	if *onlyDumpMaps {
+		dumpMaps()
+		return
+	}
+
 	dsn := os.Getenv("DATA_SOURCE_NAME")
 	if len(dsn) == 0 {
 		log.Fatal("couldn't find environment variable DATA_SOURCE_NAME")