Add -label flag to set additional custom labels

2 files changed, 87 insertions, 28 deletions

diff --git a/flags.go b/flags.go
new file mode 100644
index 0000000..8241179
--- /dev/null
+++ b/flags.go
@@ -0,0 +1,40 @@
+package main
+
+import (
+	"fmt"
+	"strings"
+)
+
+type MapVar map[string]string
+
+func (this *MapVar) Set(value string) error {
+	if *this == nil {
+		*this = make(map[string]string)
+	}
+
+	splits := strings.SplitN(value, "=", 2)
+
+	(*this)[splits[0]] = splits[1]
+
+	return nil
+}
+
+func (this *MapVar) String() string {
+	s := ""
+
+	for key, value := range *this {
+		s += fmt.Sprintf("%s: %v\n", key, value)
+	}
+
+	return fmt.Sprintf("{%s}", s)
+}
+
+type SliceVar []string
+
+func (this *SliceVar) Set(value string) error {
+	return nil
+}
+
+func (this *SliceVar) String() string {
+	return ""
+}
diff --git a/main.go b/main.go
index 6408877..109b5cc 100644
--- a/main.go
+++ b/main.go
@@ -5,9 +5,11 @@ import (
 	"fmt"
 	"io/fs"
 	"log/slog"
+	"maps"
 	"net/http"
 	"os"
 	"path/filepath"
+	"slices"
 	"strconv"
 	"time"
 
@@ -16,12 +18,13 @@ import (
 )
 
 var (
-	httpHost  string
-	httpPort  int
-	interval  int
-	namespace string
-	path      string
-	subsystem string
+	httpHost     string
+	httpPort     int
+	interval     int
+	namespace    string
+	path         string
+	subsystem    string
+	customLabels MapVar
 
 	extensionKeys map[string]bool
 	gauges        map[string]*prometheus.GaugeVec
@@ -35,6 +38,7 @@ func main() {
 	flag.StringVar(&namespace, "namespace", "fsdir_exporter", "Namespace to use for the Prometheus metrics")
 	flag.StringVar(&path, "path", "", "The path to export metrics of")
 	flag.StringVar(&subsystem, "subsystem", "", "Subsystem to use for the Prometheus metrics")
+	flag.Var(&customLabels, "label", "Additional labels to apply to the metrics")
 	flag.Parse()
 
 	// Check required options
@@ -43,26 +47,32 @@ func main() {
 		os.Exit(1)
 	}
 
+	// Create a slice of all label keys to include user-defined additions
+	labelKeys := slices.Collect(maps.Keys(customLabels))
+	labelKeys = append(labelKeys, "extension")
+
 	// Instantiate Prometheus metrics
 	gauges = make(map[string]*prometheus.GaugeVec)
 
-	gauges["count"] = prometheus.NewGaugeVec(prometheus.GaugeOpts{
-		Namespace: namespace,
-		Subsystem: subsystem,
-		Name:      "count",
-		Help:      "Number of files found in the given directory",
-	}, []string{
-		"extension",
-	})
+	gauges["count"] = prometheus.NewGaugeVec(
+		prometheus.GaugeOpts{
+			Namespace: namespace,
+			Subsystem: subsystem,
+			Name:      "count",
+			Help:      "Number of files found in the given directory",
+		},
+		labelKeys,
+	)
 
-	gauges["size"] = prometheus.NewGaugeVec(prometheus.GaugeOpts{
-		Namespace: namespace,
-		Subsystem: subsystem,
-		Name:      "size",
-		Help:      "Size (in bytes) of files found in the given directory",
-	}, []string{
-		"extension",
-	})
+	gauges["size"] = prometheus.NewGaugeVec(
+		prometheus.GaugeOpts{
+			Namespace: namespace,
+			Subsystem: subsystem,
+			Name:      "size",
+			Help:      "Size (in bytes) of files found in the given directory",
+		},
+		labelKeys,
+	)
 
 	// Register Prometheus metrics
 	slog.Debug("Registering Prometheus metrics")
@@ -91,6 +101,15 @@ func main() {
 	slog.Error(err.Error())
 }
 
+func mergeLabels(extra map[string]string) prometheus.Labels {
+	labels := make(map[string]string)
+
+	maps.Copy(labels, customLabels)
+	maps.Copy(labels, extra)
+
+	return labels
+}
+
 func update() {
 	sizes := make(map[string]float64)
 	counts := make(map[string]float64)
@@ -129,19 +148,19 @@ func update() {
 	slog.Debug("Writing new values to metrics")
 	for key, _ := range extensionKeys {
 		if _, ok := counts[key]; ok {
-			gauges["count"].With(prometheus.Labels{
+			gauges["count"].With(mergeLabels(map[string]string{
 				"extension": key,
-			}).Set(counts[key])
-			gauges["size"].With(prometheus.Labels{
+			})).Set(float64(counts[key]))
+			gauges["size"].With(mergeLabels(map[string]string{
 				"extension": key,
-			}).Set(sizes[key])
+			})).Set(float64(sizes[key]))
 
 			continue
 		}
 
 		// Clean up extension which no longer exists in the directory
-		gauges["count"].Delete(prometheus.Labels{"extension": key})
-		gauges["size"].Delete(prometheus.Labels{"extension": key})
+		gauges["count"].Delete(mergeLabels(map[string]string{"extension": key}))
+		gauges["size"].Delete(mergeLabels(map[string]string{"extension": key}))
 		delete(extensionKeys, key)
 	}
 }