Add collectors settings: make possible to specify collecting metrics …

…from multiple databases.

internal/collector/collector_common.go

@@ -25,6 +25,7 @@ func (d *typedDesc) mustNewConstMetric(value float64, labels ...string) promethe
 
 // typedDescSet unions metrics in a set, which could be collected using query.
 type typedDescSet struct {
+	databases      []string             // list of databases from which metrics should be collected
 	query          string               // query used for requesting stats
 	variableLabels []string             // ordered list of labels names
 	metricNames    []string             // ordered list of metrics short names (with no namespace/subsystem)
@@ -34,6 +35,13 @@ type typedDescSet struct {
 // newDescSet creates new typedDescSet based on passed metrics attributes.
 func newDescSet(constLabels prometheus.Labels, namespace, subsystem string, settings model.MetricsSubsystem) typedDescSet {
 	var variableLabels []string
+
+	// Add extra "database" label to metrics collected from different databases.
+	if len(settings.Databases) > 0 {
+		variableLabels = append(variableLabels, "database")
+	}
+
+	// Construct the rest of labels slice.
 	for _, m := range settings.Metrics {
 		if m.Usage == "LABEL" {
 			variableLabels = append(variableLabels, m.ShortName)
@@ -42,11 +50,13 @@ func newDescSet(constLabels prometheus.Labels, namespace, subsystem string, sett
 
 	descs := make(map[string]typedDesc)
 
+	// typeMap is auxiliary dictionary for selecting proper Prometheus data type depending on 'usage' property.
 	typeMap := map[string]prometheus.ValueType{
 		"COUNTER": prometheus.CounterValue,
 		"GAUGE":   prometheus.GaugeValue,
 	}
 
+	// Construct metrics names and descriptors slices.
 	var metricNames []string
 	for _, m := range settings.Metrics {
 		if m.Usage == "LABEL" {
@@ -68,6 +78,7 @@ func newDescSet(constLabels prometheus.Labels, namespace, subsystem string, sett
 	}
 
 	return typedDescSet{
+		databases:      settings.Databases,
 		query:          settings.Query,
 		metricNames:    metricNames,
 		variableLabels: variableLabels,

internal/collector/collector_common_test.go

@@ -8,26 +8,54 @@ import (
 )
 
 func Test_newDescSet(t *testing.T) {
-	subsystem := model.MetricsSubsystem{
-		Query: "SELECT 'l1' as label1, 'l2' as label2, 0.123 as metric1, 0.456 as metric2, 0.789 as metric3",
-		Metrics: model.Metrics{
-			{ShortName: "label1", Usage: "LABEL", Description: "label1 description"},
-			{ShortName: "label2", Usage: "LABEL", Description: "label2 description"},
-			{ShortName: "metric1", Usage: "GAUGE", Description: "metric1 description"},
-			{ShortName: "metric2", Usage: "GAUGE", Description: "metric2 description"},
-			{ShortName: "metric3", Usage: "GAUGE", Description: "metric3 description"},
+	testcases := []struct {
+		subsystem       model.MetricsSubsystem
+		wantQuery       string
+		wantVarLabels   []string
+		wantMetricNames []string
+	}{
+		{
+			// With no databases specified
+			subsystem: model.MetricsSubsystem{
+				Query: "SELECT 'l1' as label1, 'l2' as label2, 0.123 as metric1, 0.456 as metric2, 0.789 as metric3",
+				Metrics: model.Metrics{
+					{ShortName: "label1", Usage: "LABEL", Description: "label1 description"},
+					{ShortName: "label2", Usage: "LABEL", Description: "label2 description"},
+					{ShortName: "metric1", Usage: "GAUGE", Description: "metric1 description"},
+					{ShortName: "metric2", Usage: "GAUGE", Description: "metric2 description"},
+					{ShortName: "metric3", Usage: "GAUGE", Description: "metric3 description"},
+				},
+			},
+			wantQuery:       "SELECT 'l1' as label1, 'l2' as label2, 0.123 as metric1, 0.456 as metric2, 0.789 as metric3",
+			wantVarLabels:   []string{"label1", "label2"},
+			wantMetricNames: []string{"metric1", "metric2", "metric3"},
+		},
+		{
+			// With databases specified
+			subsystem: model.MetricsSubsystem{
+				Databases: []string{"pgscv_fixtures"},
+				Query:     "SELECT 'l1' as label1, 'l2' as label2, 0.123 as metric1, 0.456 as metric2, 0.789 as metric3",
+				Metrics: model.Metrics{
+					{ShortName: "label1", Usage: "LABEL", Description: "label1 description"},
+					{ShortName: "label2", Usage: "LABEL", Description: "label2 description"},
+					{ShortName: "metric1", Usage: "GAUGE", Description: "metric1 description"},
+					{ShortName: "metric2", Usage: "GAUGE", Description: "metric2 description"},
+					{ShortName: "metric3", Usage: "GAUGE", Description: "metric3 description"},
+				},
+			},
+			wantQuery:       "SELECT 'l1' as label1, 'l2' as label2, 0.123 as metric1, 0.456 as metric2, 0.789 as metric3",
+			wantVarLabels:   []string{"database", "label1", "label2"},
+			wantMetricNames: []string{"metric1", "metric2", "metric3"},
 		},
 	}
 
 	constLabels := prometheus.Labels{"constlabel": "example"}
 
-	wantQuery := "SELECT 'l1' as label1, 'l2' as label2, 0.123 as metric1, 0.456 as metric2, 0.789 as metric3"
-	wantVariableLabels := []string{"label1", "label2"}
-	wantMetricNames := []string{"metric1", "metric2", "metric3"}
-
-	descSet := newDescSet(constLabels, "postgres", "class", subsystem)
-	assert.Equal(t, wantQuery, descSet.query)
-	assert.Equal(t, wantVariableLabels, descSet.variableLabels)
-	assert.Equal(t, wantMetricNames, descSet.metricNames)
-	assert.NotNil(t, descSet.descs)
+	for _, tc := range testcases {
+		descSet := newDescSet(constLabels, "postgres", "class", tc.subsystem)
+		assert.Equal(t, tc.wantQuery, descSet.query)
+		assert.Equal(t, tc.wantVarLabels, descSet.variableLabels)
+		assert.Equal(t, tc.wantMetricNames, descSet.metricNames)
+		assert.NotNil(t, descSet.descs)
+	}
 }

internal/collector/postgres_custom.go

@@ -1,6 +1,7 @@
 package collector
 
 import (
+	"github.com/jackc/pgx/v4"
 	"github.com/prometheus/client_golang/prometheus"
 	"github.com/weaponry/pgscv/internal/log"
 	"github.com/weaponry/pgscv/internal/model"
@@ -30,6 +31,39 @@ func NewPostgresCustomCollector(constLabels prometheus.Labels, settings model.Co
 
 // Update method collects statistics, parse it and produces metrics that are sent to Prometheus.
 func (c *postgresCustomCollector) Update(config Config, ch chan<- prometheus.Metric) error {
+
+	// Make list of databases should be visited for collecting metrics.
+	databases := []string{}
+	for _, set := range c.descSets {
+		m := map[string]bool{}
+		for _, dbname := range set.databases {
+			m[dbname] = true
+		}
+
+		for dbname := range m {
+			databases = append(databases, dbname)
+		}
+	}
+
+	// Collect multiple-databases metrics.
+	if len(databases) > 0 {
+		err := c.updateFromMultipleDatabases(config, databases, ch)
+		if err != nil {
+			log.Errorf("collect failed: %s; skip", err)
+		}
+	}
+
+	// Collect once-database metrics.
+	err := c.updateFromSingleDatabase(config, ch)
+	if err != nil {
+		log.Errorf("collect failed: %s; skip", err)
+	}
+
+	return nil
+}
+
+// updateFromSingleDatabase method visit only one database and collect necessary metrics.
+func (c *postgresCustomCollector) updateFromSingleDatabase(config Config, ch chan<- prometheus.Metric) error {
 	conn, err := store.New(config.ConnString)
 	if err != nil {
 		return err
@@ -41,22 +75,107 @@ func (c *postgresCustomCollector) Update(config Config, ch chan<- prometheus.Met
 	// and translate stats into metrics.
 
 	for _, s := range c.descSets {
-		res, err := conn.Query(s.query)
+		// Skip sets with multiple databases.
+		if len(s.databases) > 0 {
+			continue
+		}
+
+		err = updateDescSet(conn, s, ch)
 		if err != nil {
-			log.Errorf("query failed: %s; skip", err)
+			log.Errorf("collect failed: %s; skip", err)
 			continue
 		}
+	}
 
-		stats := parsePostgresCustomStats(res, s.variableLabels)
+	return nil
+}
+
+// updateFromMultipleDatabases method visits all requested databases and collects necessary metrics.
+func (c *postgresCustomCollector) updateFromMultipleDatabases(config Config, userDatabases []string, ch chan<- prometheus.Metric) error {
+	conn, err := store.New(config.ConnString)
+	if err != nil {
+		return err
+	}
+
+	realDatabases, err := listDatabases(conn)
+	if err != nil {
+		return err
+	}
 
-		// iterate over stats, extract labels and values, wrap to metric and send to receiver.
-		for key, stat := range stats {
-			labelValues := strings.Split(key, "/")
+	conn.Close()
 
-			for name, value := range stat {
-				d := s.descs[name]
-				ch <- d.mustNewConstMetric(value, labelValues...)
+	pgconfig, err := pgx.ParseConfig(config.ConnString)
+	if err != nil {
+		return err
+	}
+
+	// walk through all databases, connect to it and collect schema-specific stats
+	for _, dbname := range userDatabases {
+		// Skip user-specified databases which are not really exist.
+		if !stringsContains(realDatabases, dbname) {
+			continue
+		}
+
+		// Create
+		pgconfig.Database = dbname
+		conn, err := store.NewWithConfig(pgconfig)
+		if err != nil {
+			return err
+		}
+
+		for _, s := range c.descSets {
+			// Skip sets with single databases, and databases which are not listed in set's databases.
+			if len(s.databases) == 0 || !stringsContains(s.databases, dbname) {
+				continue
+			}
+
+			// Swap descriptors labels, add database as first label
+			if len(s.variableLabels) > 0 && s.variableLabels[0] != "database" {
+				s.variableLabels = append([]string{"database"}, s.variableLabels...)
 			}
+
+			err = updateDescSet(conn, s, ch)
+			if err != nil {
+				log.Errorf("collect failed: %s; skip", err)
+				continue
+			}
+		}
+
+		// Close connection.
+		conn.Close()
+	}
+
+	return nil
+}
+
+// updateDescSet using passed connection collects metrics for specified descSet and receives it to metric channel.
+func updateDescSet(conn *store.DB, set typedDescSet, ch chan<- prometheus.Metric) error {
+	res, err := conn.Query(set.query)
+	if err != nil {
+		return err
+	}
+
+	stats := parsePostgresCustomStats(res, set.variableLabels)
+
+	// Get database name from config.
+	// Database name used as value for 'database' label in case of
+	// user-defined metrics collected from multiple databases.
+	dbname := conn.Conn().Config().Database
+
+	// iterate over stats, extract labels and values, wrap to metric and send to receiver.
+	for key, stat := range stats {
+
+		// If database label present in variable labels, prepend label values with database name.
+		var labelValues []string
+		if len(set.variableLabels) > 0 && set.variableLabels[0] == "database" {
+			labelValues = append([]string{dbname}, strings.Split(key, "/")...)
+		} else {
+			labelValues = strings.Split(key, "/")
+		}
+
+		for name, value := range stat {
+			d := set.descs[name]
+			ch <- d.mustNewConstMetric(value, labelValues...)
 		}
 	}