From 5e8a4143dcefa68c57d35b4ea73f0c9f8ab0e06d Mon Sep 17 00:00:00 2001
From: Raghu Udiyar <raghusiddarth@gmail.com>
Date: Mon, 28 Dec 2020 20:15:50 +0530
Subject: [PATCH] Add exclude and include to filter by regex

Regex filtering is applied against the complete metric name and tags
based on the Prometheus exposition format, allowing filtering both
on metric name as well as tags, for e.g.

-include-regex 'go_gc_duration_seconds{quantile="0\..?.?"}'
---
 main.go | 52 ++++++++++++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 52 insertions(+)

diff --git a/main.go b/main.go
index 1c6e076..aee6b14 100644
--- a/main.go
+++ b/main.go
@@ -10,6 +10,7 @@ import (
 	"log"
 	"net/http"
 	"os"
+	"regexp"
 	"strconv"
 	"strings"
 	"time"
@@ -117,6 +118,47 @@ func CreateInfluxMetrics(samples model.Vector, metricPrefix string) string {
 	return metrics
 }
 
+func FilterSamples(samples model.Vector, includeRegex string, excludeRegex string) (model.Vector, error) {
+	var reInclude, reExclude *regexp.Regexp
+	var err error
+
+	if includeRegex != "" {
+		reInclude, err = regexp.Compile(includeRegex)
+		if err != nil {
+			return nil, err
+		}
+	}
+
+	if excludeRegex != "" {
+		reExclude, err = regexp.Compile(excludeRegex)
+		if err != nil {
+			return nil, err
+		}
+	}
+
+	var filteredSamples model.Vector
+	for _, sample := range samples {
+		metricString := sample.Metric.String()
+
+		var matchInclude, matchExclude bool
+		if reInclude == nil {
+			// include all metrics
+			matchInclude = true
+		} else {
+			matchInclude = reInclude.MatchString(metricString)
+		}
+
+		if reExclude != nil {
+			matchExclude = reExclude.MatchString(metricString)
+		}
+
+		if matchInclude == true && matchExclude == false {
+			filteredSamples = append(filteredSamples, sample)
+		}
+	}
+	return filteredSamples, nil
+}
+
 func OutputMetrics(samples model.Vector, outputFormat string, metricPrefix string) error {
 	output := ""
 
@@ -239,6 +281,8 @@ func main() {
 	promURL := flag.String("prom-url", "http://localhost:9090", "Prometheus API URL.")
 	queryString := flag.String("prom-query", "up", "Prometheus API query string.")
 	outputFormat := flag.String("output-format", "influx", "The check output format to use for metrics {influx|graphite|json}.")
+	includeRegex := flag.String("include-regex", "", "Regex to include metrics applied agasint the metric in Prometheus exposition format")
+	excludeRegex := flag.String("exclude-regex", "", "Regex to exclude metrics, applied after -include-regex")
 	metricPrefix := flag.String("metric-prefix", "", "Metric name prefix, only supported by line protocol output formats.")
 	insecureSkipVerify := flag.Bool("insecure-skip-verify", false, "Skip TLS peer verification.")
 	flag.Parse()
@@ -270,6 +314,14 @@ func main() {
 		}
 	}
 
+	if *includeRegex != "" || *excludeRegex != "" {
+		samples, err = FilterSamples(samples, *includeRegex, *excludeRegex)
+		if err != nil {
+			log.Fatal(err)
+			os.Exit(2)
+		}
+	}
+
 	err = OutputMetrics(samples, *outputFormat, *metricPrefix)
 
 	if err != nil {