Yaml parsing refacto

cmd/postgres_exporter/postgres_exporter.go

@@ -85,6 +85,26 @@ func (cu *ColumnUsage) UnmarshalYAML(unmarshal func(interface{}) error) error {
 	return nil
 }
 
+// MappingOptions is a copy of ColumnMapping used only for parsing
+type MappingOptions struct {
+	Usage             string             `yaml:"usage"`
+	Description       string             `yaml:"description"`
+	Mapping           map[string]float64 `yaml:"metric_mapping"` // Optional column mapping for MAPPEDMETRIC
+	SupportedVersions semver.Range       `yaml:"pg_version"`     // Semantic version ranges which are supported. Unsupported columns are not queried (internally converted to DISCARD).
+}
+
+// nolint: golint
+type Mapping map[string]MappingOptions
+
+// nolint: golint
+type UserQuery struct {
+	Query   string    `yaml:"query"`
+	Metrics []Mapping `yaml:"metrics"`
+}
+
+// nolint: golint
+type UserQueries map[string]UserQuery
+
 // Regex used to get the "short-version" from the postgres version field.
 var versionRegex = regexp.MustCompile(`^\w+ ((\d+)(\.\d+)?(\.\d+)?)`)
 var lowestSupportedVersion = semver.MustParse("9.1.0")
@@ -392,78 +412,58 @@ func makeQueryOverrideMap(pgVersion semver.Version, queryOverrides map[string][]
 	return resultMap
 }
 
-// Add queries to the builtinMetricMaps and queryOverrides maps. Added queries do not
-// respect version requirements, because it is assumed that the user knows
-// what they are doing with their version of postgres.
-//
-// This function modifies metricMap and queryOverrideMap to contain the new
-// queries.
-// TODO: test code for all cu.
-// TODO: use proper struct type system
-// TODO: the YAML this supports is "non-standard" - we should move away from it.
-func addQueries(content []byte, pgVersion semver.Version, server *Server) error {
-	var extra map[string]interface{}
+func parseUserQueries(content []byte) (map[string]map[string]ColumnMapping, map[string]string, error) {
+	var userQueries UserQueries
 
-	err := yaml.Unmarshal(content, &extra)
+	err := yaml.Unmarshal(content, &userQueries)
 	if err != nil {
-		return err
+		return nil, nil, err
 	}
 
 	// Stores the loaded map representation
 	metricMaps := make(map[string]map[string]ColumnMapping)
 	newQueryOverrides := make(map[string]string)
 
-	for metric, specs := range extra {
+	for metric, specs := range userQueries {
 		log.Debugln("New user metric namespace from YAML:", metric)
-		for key, value := range specs.(map[interface{}]interface{}) {
-			switch key.(string) {
-			case "query":
-				query := value.(string)
-				newQueryOverrides[metric] = query
-
-			case "metrics":
-				for _, c := range value.([]interface{}) {
-					column := c.(map[interface{}]interface{})
-
-					for n, a := range column {
-						var columnMapping ColumnMapping
-
-						// Fetch the metric map we want to work on.
-						metricMap, ok := metricMaps[metric]
-						if !ok {
-							// Namespace for metric not found - add it.
-							metricMap = make(map[string]ColumnMapping)
-							metricMaps[metric] = metricMap
-						}
-
-						// Get name.
-						name := n.(string)
-
-						for attrKey, attrVal := range a.(map[interface{}]interface{}) {
-							switch attrKey.(string) {
-							case "usage":
-								usage, err := stringToColumnUsage(attrVal.(string))
-								if err != nil {
-									return err
-								}
-								columnMapping.usage = usage
-							case "description":
-								columnMapping.description = attrVal.(string)
-							}
-						}
-
-						// TODO: we should support cu
-						columnMapping.mapping = nil
-						// Should we support this for users?
-						columnMapping.supportedVersions = nil
-
-						metricMap[name] = columnMapping
-					}
-				}
+		newQueryOverrides[metric] = specs.Query
+		metricMap, ok := metricMaps[metric]
+		if !ok {
+			// Namespace for metric not found - add it.
+			metricMap = make(map[string]ColumnMapping)
+			metricMaps[metric] = metricMap
+		}
+		for _, metric := range specs.Metrics {
+			for name, mappingOption := range metric {
+				var columnMapping ColumnMapping
+				tmpUsage, _ := stringToColumnUsage(mappingOption.Usage)
+				columnMapping.usage = tmpUsage
+				columnMapping.description = mappingOption.Description
+
+				// TODO: we should support cu
+				columnMapping.mapping = nil
+				// Should we support this for users?
+				columnMapping.supportedVersions = nil
+				metricMap[name] = columnMapping
 			}
 		}
 	}
+	return metricMaps, newQueryOverrides, nil
+}
 
+// Add queries to the builtinMetricMaps and queryOverrides maps. Added queries do not
+// respect version requirements, because it is assumed that the user knows
+// what they are doing with their version of postgres.
+//
+// This function modifies metricMap and queryOverrideMap to contain the new
+// queries.
+// TODO: test code for all cu.
+// TODO: the YAML this supports is "non-standard" - we should move away from it.
+func addQueries(content []byte, pgVersion semver.Version, server *Server) error {
+	metricMaps, newQueryOverrides, err := parseUserQueries(content)
+	if err != nil {
+		return nil
+	}
 	// Convert the loaded metric map into exporter representation
 	partialExporterMap := makeDescMap(pgVersion, server.labels, metricMaps)
 
@@ -488,7 +488,6 @@ func addQueries(content []byte, pgVersion semver.Version, server *Server) error
 		}
 		server.queryOverrides[k] = v
 	}
-
 	return nil
 }
 

cmd/postgres_exporter/postgres_exporter_test.go

@@ -3,6 +3,7 @@
 package main
 
 import (
+	"io/ioutil"
 	"os"
 	"reflect"
 	"testing"
@@ -305,3 +306,17 @@ func (s *FunctionalSuite) TestBooleanConversionToValueAndString(c *C) {
 		c.Assert(ok, Equals, cs.expectedOK)
 	}
 }
+
+func (s *FunctionalSuite) TestParseUserQueries(c *C) {
+	userQueriesData, err := ioutil.ReadFile("./tests/user_queries_ok.yaml")
+	if err == nil {
+		metricMaps, newQueryOverrides, err := parseUserQueries(userQueriesData)
+		c.Assert(err, Equals, nil)
+		c.Assert(metricMaps, NotNil)
+		c.Assert(newQueryOverrides, NotNil)
+
+		if len(metricMaps) != 2 {
+			c.Errorf("Expected 2 metrics from user file, got %d", len(metricMaps))
+		}
+	}
+}

cmd/postgres_exporter/tests/user_queries_ok.yaml

@@ -0,0 +1,23 @@
+pg_locks_mode:
+  query: "WITH q_locks AS (select * from pg_locks where pid != pg_backend_pid() and database = (select oid from pg_database where datname = current_database())) SELECT (select current_database()) as datname,
+  lockmodes AS tag_lockmode, coalesce((select count(*) FROM q_locks WHERE mode = lockmodes), 0) AS count FROM
+  unnest('{AccessShareLock, ExclusiveLock, RowShareLock, RowExclusiveLock, ShareLock, ShareRowExclusiveLock,  AccessExclusiveLock, ShareUpdateExclusiveLock}'::text[]) lockmodes;"
+  metrics:
+    - datname:
+        usage: "LABEL"
+        description: "Database name"
+    - tag_lockmode:
+        usage: "LABEL"
+        description: "Lock type"
+    - count:
+        usage: "GAUGE"
+        description: "Number of lock"
+pg_wal:
+  query: "select current_database() as datname, case when pg_is_in_recovery() = false then pg_xlog_location_diff(pg_current_xlog_location(), '0/0')::int8 else pg_xlog_location_diff(pg_last_xlog_replay_location(), '0/0')::int8 end as xlog_location_b;"
+  metrics:
+    - datname:
+        usage: "LABEL"
+        description: "Database name"
+    - xlog_location_b:
+        usage: "COUNTER"
+        description: "current transaction log write location"