From 3e58509026211ac22001104ab18b53e161c3f5c8 Mon Sep 17 00:00:00 2001
From: Guillaume Lecerf <guillaume.lecerf@iziwork.com>
Date: Wed, 3 Nov 2021 16:29:17 +0100
Subject: [PATCH] Add an option to override value for empty result set

---
 collector/collector.go      | 52 +++++++++++++++++++++++++------------
 collector/collector_test.go | 16 ++++++++++++
 example/configv2.yaml       |  2 ++
 3 files changed, 53 insertions(+), 17 deletions(-)

diff --git a/collector/collector.go b/collector/collector.go
index 39664538..e34f51f9 100644
--- a/collector/collector.go
+++ b/collector/collector.go
@@ -51,21 +51,23 @@ type Config struct {
 // A metric defines what metric should be generated from what MongoDB aggregation.
 // The pipeline configures (as JSON) the aggreation query
 type Metric struct {
-	Name        string
-	Type        string
-	Servers     []string
-	Help        string
-	Value       string
-	Cache       int64
-	ConstLabels prometheus.Labels
-	Mode        string
-	Labels      []string
-	Database    string
-	Collection  string
-	Pipeline    string
-	desc        *prometheus.Desc
-	pipeline    bson.A
-	validUntil  time.Time
+	Name          string
+	Type          string
+	Servers       []string
+	Help          string
+	Value         string
+	OverrideEmpty bool
+	EmptyValue    int64
+	Cache         int64
+	ConstLabels   prometheus.Labels
+	Mode          string
+	Labels        []string
+	Database      string
+	Collection    string
+	Pipeline      string
+	desc          *prometheus.Desc
+	pipeline      bson.A
+	validUntil    time.Time
 }
 
 var (
@@ -358,10 +360,10 @@ func (c *Collector) generateMetrics(metric *Metric, srv *server, ch chan<- prome
 
 	var multierr *multierror.Error
 	var i int
+	var result = make(AggregationResult)
 
 	for cursor.Next(ctx) {
 		i++
-		var result AggregationResult
 
 		err := cursor.Decode(&result)
 		c.logger.Debugf("found record %s from metric %s", result, metric.Name)
@@ -383,7 +385,23 @@ func (c *Collector) generateMetrics(metric *Metric, srv *server, ch chan<- prome
 	}
 
 	if i == 0 {
-		return fmt.Errorf("metric %s aggregation returned an emtpy result set", metric.Name)
+		if ! metric.OverrideEmpty {
+			return fmt.Errorf("metric %s aggregation returned an empty result set", metric.Name)
+		}
+
+		c.logger.Debugf("OverrideEmpty option is set for metric %s, overriding value with %d", metric.Name, metric.EmptyValue)
+
+		result[metric.Value] = int64(metric.EmptyValue)
+		for _, label := range metric.Labels {
+			result[label] = ""
+		}
+		m, err := createMetric(metric, result)
+		if err != nil {
+			return err
+		}
+
+		c.updateCache(metric, srv, m)
+		ch <- m
 	}
 
 	return multierr.ErrorOrNil()
diff --git a/collector/collector_test.go b/collector/collector_test.go
index 9fd40be8..ffcc7e8c 100644
--- a/collector/collector_test.go
+++ b/collector/collector_test.go
@@ -123,6 +123,22 @@ func TestInitializeMetrics(t *testing.T) {
 			//error: "1 error occurred:\n\t* value not found in result set\n\n",
 			expected: ``,
 		},
+		metricTest{
+			name: "Unlabeled gauge no value found in result but OverrideEmpty is set with EmptyValue 0",
+			metric: &Metric{
+				Name:          "simple_gauge_value_not_found_overriden",
+				Type:          "gauge",
+				Help:          "overridden",
+				OverrideEmpty: true,
+				EmptyValue:    12,
+				Pipeline:      "[{\"$match\":{\"foo\":\"bar\"}}]",
+			},
+			expected: `
+				# HELP simple_gauge_value_not_found_overriden overridden
+				# TYPE simple_gauge_value_not_found_overriden gauge
+				simple_gauge_value_not_found_overriden 12
+			`,
+		},
 		metricTest{
 			name: "Unlabeled gauge value not of type float",
 			metric: &Metric{
diff --git a/example/configv2.yaml b/example/configv2.yaml
index 8d1cd482..9d7f6729 100644
--- a/example/configv2.yaml
+++ b/example/configv2.yaml
@@ -19,6 +19,8 @@ metrics:
   servers: [main] #Can also be empty, if empty the metric will be used for every server defined
   help: 'Simple gauge metric'
   value: total
+  overrideEmpty: true # if an empty result set is returned..
+  emptyValue: 0       # create a metric with value 0
   labels: []
   mode: pull
   cache: 0