Support multiple -topic-pattern flags

radhus · Nov 14, 2020 · 424354c · 424354c
1 parent 08b965b
commit 424354c
Show file tree

Hide file tree

Showing 2 changed files with 24 additions and 8 deletions.
diff --git a/go.sum b/go.sum
@@ -83,6 +83,7 @@ github.com/golang/protobuf v1.4.0-rc.2/go.mod h1:LlEzMj4AhA7rCAGe4KMBDvJI+AwstrU
 github.com/golang/protobuf v1.4.0-rc.4.0.20200313231945-b860323f09d0/go.mod h1:WU3c8KckQ9AFe+yFwt9sWVRKCVIyN9cPHBJSNnbL67w=
 github.com/golang/protobuf v1.4.0/go.mod h1:jodUvKwWbYaEsadDk5Fwe5c77LiNKVO9IDvqG2KuDX0=
 github.com/golang/protobuf v1.4.2/go.mod h1:oDoupMAO8OvCJWAcko0GGGIgR6R6ocIYbsSw735rRwI=
+github.com/golang/protobuf v1.4.3 h1:JjCZWpVbqXDqFVmTfYWEVTMIYrL/NPdPSCHPJ0T/raM=
 github.com/golang/protobuf v1.4.3/go.mod h1:oDoupMAO8OvCJWAcko0GGGIgR6R6ocIYbsSw735rRwI=
 github.com/golang/snappy v0.0.0-20180518054509-2e65f85255db/go.mod h1:/XxbfmMg8lxefKM7IXC3fBNl/7bRcc72aCRzEWrmP2Q=
 github.com/google/btree v0.0.0-20180813153112-4030bb1f1f0c/go.mod h1:lNA+9X1NB3Zf8V7Ke586lFgjr2dZNuvo3lPJSGZ5JPQ=
@@ -386,6 +387,7 @@ google.golang.org/protobuf v0.0.0-20200221191635-4d8936d0db64/go.mod h1:kwYJMbMJ
 google.golang.org/protobuf v0.0.0-20200228230310-ab0ca4ff8a60/go.mod h1:cfTl7dwQJ+fmap5saPgwCLgHXTUD7jkjRqWcaiX5VyM=
 google.golang.org/protobuf v1.20.1-0.20200309200217-e05f789c0967/go.mod h1:A+miEFZTKqfCUM6K7xSMQL9OKL/b6hQv+e19PK+JZNE=
 google.golang.org/protobuf v1.21.0/go.mod h1:47Nbq4nVaFHyn7ilMalzfO3qCViNmqZ2kzikPIcrTAo=
+google.golang.org/protobuf v1.23.0 h1:4MY060fB1DLGMB/7MBTLnwQUY6+F09GEiz6SsrNqyzM=
 google.golang.org/protobuf v1.23.0/go.mod h1:EGpADcykh3NcUnDUJcl1+ZksZNG86OlYog2l/sGQquU=
 gopkg.in/alecthomas/kingpin.v2 v2.2.6/go.mod h1:FMv+mEhP44yOT+4EoQTLFTRgOQ1FBLkstjWtayDeSgw=
 gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0=

diff --git a/main.go b/main.go
@@ -15,11 +15,22 @@ import (
 
 const topicLabel = "topic"
 
+type stringList []string
+
+func (sl *stringList) String() string {
+	return ""
+}
+
+func (sl *stringList) Set(s string) error {
+	*sl = append(*sl, s)
+	return nil
+}
+
 var (
-	url          = flag.String("url", "", "MQTT URL")
-	clientID     = flag.String("client-id", "mqtt-exporter", "MQTT client ID")
-	topicPattern = flag.String("topic-pattern", "", "Subscribe to topic pattern")
-	listenAddr   = flag.String("addr", ":8080", "HTTP server listening address")
+	url           = flag.String("url", "", "MQTT URL")
+	clientID      = flag.String("client-id", "mqtt-exporter", "MQTT client ID")
+	topicPatterns = stringList{}
+	listenAddr    = flag.String("addr", ":8080", "HTTP server listening address")
 
 	topicGauge = prometheus.NewGaugeVec(
 		prometheus.GaugeOpts{
@@ -31,8 +42,9 @@ var (
 )
 
 func main() {
+	flag.Var(&topicPatterns, "topic-pattern", "Subscribe to topic pattern (can be repeated)")
 	flag.Parse()
-	if *url == "" || *topicPattern == "" {
+	if *url == "" || len(topicPatterns) == 0 || (len(topicPatterns) == 1 && topicPatterns[0] == "") {
 		fmt.Println("-url and -topic-pattern are both required")
 		flag.Usage()
 		os.Exit(1)
@@ -54,9 +66,11 @@ func main() {
 		log.Fatalln("Failed to connect to MQTT:", token.Error())
 	}
 
-	if token := client.Subscribe(*topicPattern, 0, messageHandler); token.Wait() && token.Error() != nil {
-		client.Disconnect(0)
-		log.Fatalln("Failed to subscribe to MQTT:", token.Error())
+	for _, topicPattern := range topicPatterns {
+		if token := client.Subscribe(topicPattern, 0, messageHandler); token.Wait() && token.Error() != nil {
+			client.Disconnect(0)
+			log.Fatalln("Failed to subscribe to MQTT:", token.Error())
+		}
 	}
 
 	log.Println("Listening on:", *listenAddr)